package com.dexorder.flink.ingestor;

import com.dexorder.flink.zmq.ZmqChannelManager;
import com.dexorder.proto.DataRequest;
import com.dexorder.proto.RealtimeParams;
import com.dexorder.proto.SlotType;
import com.dexorder.proto.SubmitHistoricalRequest;
import com.dexorder.proto.WorkComplete;
import com.dexorder.proto.WorkHeartbeat;
import com.dexorder.proto.WorkReject;
import com.dexorder.proto.WorkStop;
import com.dexorder.proto.WorkerReady;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.zeromq.ZMQ;

import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Deque;
import java.util.List;
import java.util.Map;
import java.util.Queue;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedQueue;

/**
 * Slot-based work broker for ingestors.
 *
 * Each WorkerReady message from an ingestor represents ONE available slot for a
 * specific exchange and job type (HISTORICAL or REALTIME). Flink consumes the slot
 * by dispatching a DataRequest to it. The ingestor re-offers the slot (sends another
 * WorkerReady) once the job completes, subject to any rate-limit backoff.
 *
 * Also receives SubmitHistoricalRequest messages forwarded by the relay on the PULL
 * socket (5566), and realtime job requests from RealtimeSubscriptionManager.
 *
 * Message type IDs (ZMQ framing):
 *   0x10  SubmitHistoricalRequest (relay → Flink via PULL)
 *   0x20  WorkerReady             (ingestor → Flink: one slot offer)
 *   0x21  WorkComplete            (ingestor → Flink)
 *   0x22  WorkHeartbeat           (ingestor → Flink)
 *   0x23  WorkReject              (ingestor → Flink)
 *   0x01  DataRequest/WorkAssign  (Flink → ingestor via ROUTER)
 *   0x25  WorkStop                (Flink → ingestor via ROUTER)
 */
public class IngestorBroker implements AutoCloseable {
    private static final Logger LOG = LoggerFactory.getLogger(IngestorBroker.class);

    private static final byte PROTOCOL_VERSION = 0x01;
    private static final byte MSG_TYPE_SUBMIT_REQUEST   = 0x10;
    private static final byte MSG_TYPE_WORKER_READY     = 0x20;
    private static final byte MSG_TYPE_WORK_COMPLETE    = 0x21;
    private static final byte MSG_TYPE_WORK_HEARTBEAT   = 0x22;
    private static final byte MSG_TYPE_WORK_REJECT      = 0x23;
    private static final byte MSG_TYPE_WORK_ASSIGN      = 0x01;
    private static final byte MSG_TYPE_WORK_STOP        = 0x25;

    /** Re-queue realtime job if no heartbeat received within this window (ms) */
    private static final long HEARTBEAT_TIMEOUT_MS = 25_000;
    /** Re-queue historical job if not completed within this window (ms) */
    private static final long HISTORICAL_TIMEOUT_MS = 120_000;

    private final ZmqChannelManager zmqManager;
    private volatile boolean running;
    private Thread brokerThread;

    // ── Slot tracking ─────────────────────────────────────────────────────────

    /**
     * Available slots, in LRU order (head = least recently used).
     * Each entry is one WorkerReady slot offer from an ingestor.
     */
    private final Deque<WorkerSlot> freeSlots = new ArrayDeque<>();

    /** Jobs waiting for a compatible free slot */
    private final Queue<DataRequest> pendingJobs = new ArrayDeque<>();

    /** Jobs currently executing on a slot */
    private final Map<String, ActiveJob> activeJobs = new ConcurrentHashMap<>();

    // ── Thread-safe inbound queue from RealtimeSubscriptionManager ───────────

    private final Queue<DataRequest> externalSubmissions = new ConcurrentLinkedQueue<>();

    public IngestorBroker(ZmqChannelManager zmqManager) {
        this.zmqManager = zmqManager;
    }

    public void start() {
        if (running) {
            LOG.warn("IngestorBroker already running");
            return;
        }
        running = true;
        brokerThread = new Thread(this::brokerLoop, "IngestorBroker-Thread");
        brokerThread.setDaemon(false);
        brokerThread.start();
        LOG.info("IngestorBroker started");
    }

    public void stop() {
        running = false;
        if (brokerThread != null) {
            brokerThread.interrupt();
            try {
                brokerThread.join(5000);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
        }
        LOG.info("IngestorBroker stopped");
    }

    /**
     * Submit a realtime data request from outside the broker thread (thread-safe).
     * Called by RealtimeSubscriptionManager when subscription ref count goes 0→1.
     */
    public void submitRealtimeRequest(String ticker) {
        String jobId = UUID.randomUUID().toString();
        DataRequest request = DataRequest.newBuilder()
                .setRequestId(jobId)
                .setJobId(jobId)
                .setType(DataRequest.RequestType.REALTIME_TICKS)
                .setTicker(ticker)
                .setRealtime(RealtimeParams.newBuilder()
                        .setIncludeTicks(true)
                        .setIncludeOhlc(false)
                        .build())
                .build();
        externalSubmissions.add(request);
        LOG.info("Enqueued realtime request: ticker={}, jobId={}", ticker, jobId);
    }

    /**
     * Stop all realtime jobs for a ticker (called when last subscriber leaves).
     * Thread-safe via ConcurrentHashMap.
     */
    public void stopRealtimeJobsForTicker(String ticker) {
        List<String> toStop = new ArrayList<>();
        for (Map.Entry<String, ActiveJob> entry : activeJobs.entrySet()) {
            if (entry.getValue().ticker.equals(ticker) &&
                entry.getValue().type == DataRequest.RequestType.REALTIME_TICKS) {
                toStop.add(entry.getKey());
            }
        }
        for (String jobId : toStop) {
            ActiveJob job = activeJobs.remove(jobId);
            if (job != null) {
                sendStop(job.workerIdentity, jobId);
                LOG.info("Sent STOP to ingestor: ticker={}, jobId={}", ticker, jobId);
            }
        }
    }

    // ── Broker loop ───────────────────────────────────────────────────────────

    private void brokerLoop() {
        ZMQ.Socket pullSocket   = zmqManager.getSocket(ZmqChannelManager.Channel.CLIENT_REQUEST);
        ZMQ.Socket routerSocket = zmqManager.getSocket(ZmqChannelManager.Channel.INGESTOR_BROKER);

        ZMQ.Poller poller = zmqManager.createPoller(2);
        poller.register(pullSocket, ZMQ.Poller.POLLIN);
        poller.register(routerSocket, ZMQ.Poller.POLLIN);

        LOG.info("IngestorBroker loop running");

        while (running) {
            try {
                // Drain external submissions (realtime requests from subscription manager)
                DataRequest ext;
                while ((ext = externalSubmissions.poll()) != null) {
                    enqueueJob(ext);
                }

                poller.poll(100);

                if (poller.pollin(0)) {
                    handleClientRequest(pullSocket);
                }
                if (poller.pollin(1)) {
                    handleWorkerMessage(routerSocket);
                }

                checkTimeouts();

            } catch (Exception e) {
                if (running) {
                    LOG.error("Error in broker loop", e);
                }
            }
        }

        LOG.info("IngestorBroker loop exited");
    }

    /** Receive a SubmitHistoricalRequest forwarded by relay and enqueue it. */
    private void handleClientRequest(ZMQ.Socket pullSocket) {
        byte[] versionFrame = pullSocket.recv(ZMQ.DONTWAIT);
        if (versionFrame == null) return;
        if (!pullSocket.hasReceiveMore()) return;
        byte[] messageFrame = pullSocket.recv(0);
        if (messageFrame == null || messageFrame.length < 2) return;

        if (versionFrame.length != 1 || versionFrame[0] != PROTOCOL_VERSION) {
            LOG.warn("Bad protocol version on PULL socket");
            return;
        }

        byte msgType = messageFrame[0];
        byte[] payload = Arrays.copyOfRange(messageFrame, 1, messageFrame.length);

        if (msgType != MSG_TYPE_SUBMIT_REQUEST) {
            LOG.warn("Unexpected message type on PULL socket: 0x{}", Integer.toHexString(msgType & 0xFF));
            return;
        }

        try {
            SubmitHistoricalRequest req = SubmitHistoricalRequest.parseFrom(payload);
            String jobId = UUID.randomUUID().toString();
            DataRequest dataRequest = DataRequest.newBuilder()
                    .setRequestId(req.getRequestId())
                    .setJobId(jobId)
                    .setType(DataRequest.RequestType.HISTORICAL_OHLC)
                    .setTicker(req.getTicker())
                    .setHistorical(com.dexorder.proto.HistoricalParams.newBuilder()
                            .setStartTime(req.getStartTime())
                            .setEndTime(req.getEndTime())
                            .setPeriodSeconds(req.getPeriodSeconds())
                            .build())
                    .setClientId(req.hasClientId() ? req.getClientId() : "")
                    .build();
            enqueueJob(dataRequest);
            LOG.info("Received historical request from relay: request_id={}, ticker={}",
                    req.getRequestId(), req.getTicker());
        } catch (Exception e) {
            LOG.error("Failed to parse SubmitHistoricalRequest from relay", e);
        }
    }

    /** Receive and dispatch a message from an ingestor DEALER. */
    private void handleWorkerMessage(ZMQ.Socket routerSocket) {
        // ROUTER frame layout: [identity][empty][version][type+payload]
        byte[] identity = routerSocket.recv(ZMQ.DONTWAIT);
        if (identity == null) return;
        if (!routerSocket.hasReceiveMore()) return;
        routerSocket.recv(0); // empty delimiter
        if (!routerSocket.hasReceiveMore()) return;
        byte[] versionFrame = routerSocket.recv(0);
        if (!routerSocket.hasReceiveMore()) return;
        byte[] messageFrame = routerSocket.recv(0);

        if (versionFrame == null || versionFrame.length != 1 || versionFrame[0] != PROTOCOL_VERSION) {
            LOG.warn("Bad protocol version from ingestor");
            return;
        }
        if (messageFrame == null || messageFrame.length < 1) return;

        byte msgType = messageFrame[0];
        byte[] payload = Arrays.copyOfRange(messageFrame, 1, messageFrame.length);
        String identityKey = bytesToHex(identity);

        try {
            switch (msgType & 0xFF) {
                case 0x20: handleWorkerReady(identity, identityKey, payload); break;
                case 0x21: handleWorkComplete(identityKey, payload); break;
                case 0x22: handleWorkHeartbeat(identityKey, payload); break;
                case 0x23: handleWorkReject(identityKey, payload); break;
                default:
                    LOG.warn("Unknown message type from ingestor: 0x{}", Integer.toHexString(msgType & 0xFF));
            }
        } catch (Exception e) {
            LOG.error("Error handling worker message type 0x{}", Integer.toHexString(msgType & 0xFF), e);
        }
    }

    /**
     * A WorkerReady message represents ONE slot offer for one exchange and job type.
     * Add it directly to freeSlots — no deduplication (multiple slots per ingestor are expected).
     */
    private void handleWorkerReady(byte[] identity, String identityKey, byte[] payload) throws Exception {
        WorkerReady ready = WorkerReady.parseFrom(payload);
        SlotType slotType = ready.getJobType();

        for (String exchange : ready.getExchangesList()) {
            WorkerSlot slot = new WorkerSlot(identity, identityKey, exchange.toUpperCase(), slotType);
            freeSlots.addLast(slot);
            LOG.info("Worker slot READY: id={}, exchange={}, type={}, totalFreeSlots={}",
                    identityKey, exchange, slotType, freeSlots.size());
        }

        dispatchPending();
    }

    /**
     * Historical job completed. Remove from activeJobs.
     * The ingestor will send a new typed WorkerReady to re-offer the slot.
     */
    private void handleWorkComplete(String identityKey, byte[] payload) throws Exception {
        WorkComplete complete = WorkComplete.parseFrom(payload);
        String jobId = complete.getJobId();

        ActiveJob job = activeJobs.remove(jobId);
        if (job == null) {
            LOG.warn("COMPLETE for unknown jobId={}", jobId);
        } else {
            LOG.info("Job COMPLETE: jobId={}, ticker={}, success={}", jobId, job.ticker, complete.getSuccess());
        }
        // Slot re-registration is driven by the ingestor via a new WorkerReady.
    }

    private void handleWorkHeartbeat(String identityKey, byte[] payload) throws Exception {
        WorkHeartbeat hb = WorkHeartbeat.parseFrom(payload);
        String jobId = hb.getJobId();

        ActiveJob job = activeJobs.get(jobId);
        if (job != null) {
            job.lastHeartbeat = System.currentTimeMillis();
        } else {
            LOG.warn("HEARTBEAT for unknown jobId={} from worker={}", jobId, identityKey);
        }
    }

    /**
     * Ingestor rejected the job. Re-queue it with a new ID.
     * The ingestor will send a new typed WorkerReady when it's ready again.
     */
    private void handleWorkReject(String identityKey, byte[] payload) throws Exception {
        WorkReject reject = WorkReject.parseFrom(payload);
        String jobId = reject.getJobId();
        LOG.warn("Job REJECTED by worker={}: jobId={}, reason={}", identityKey, jobId, reject.getReason());

        ActiveJob job = activeJobs.remove(jobId);
        if (job != null) {
            DataRequest requeued = job.request.toBuilder()
                    .setJobId(UUID.randomUUID().toString())
                    .build();
            pendingJobs.add(requeued);
        }
        // Slot re-registration is driven by the ingestor via a new WorkerReady.
    }

    // ── Dispatch ──────────────────────────────────────────────────────────────

    private void enqueueJob(DataRequest request) {
        WorkerSlot slot = findFreeSlot(exchangeOf(request.getTicker()), request.getType());
        if (slot != null) {
            dispatch(slot, request);
        } else {
            pendingJobs.add(request);
            LOG.debug("No free slot for {}, queued (pendingJobs={})", request.getTicker(), pendingJobs.size());
        }
    }

    private void dispatchPending() {
        Queue<DataRequest> remaining = new ArrayDeque<>();
        DataRequest job;
        while ((job = pendingJobs.poll()) != null) {
            WorkerSlot slot = findFreeSlot(exchangeOf(job.getTicker()), job.getType());
            if (slot != null) {
                dispatch(slot, job);
            } else {
                remaining.add(job);
            }
        }
        pendingJobs.addAll(remaining);
    }

    private void dispatch(WorkerSlot slot, DataRequest request) {
        try {
            byte[] protoBytes = request.toByteArray();
            boolean sent = zmqManager.sendToWorker(slot.identity, PROTOCOL_VERSION, MSG_TYPE_WORK_ASSIGN, protoBytes);
            if (!sent) {
                // ROUTER_MANDATORY: identity is disconnected — purge all stale slots for this
                // worker and re-queue the job so dispatchPending() can try a live slot.
                int purged = purgeWorkerSlots(slot.identityKey);
                LOG.warn("Worker {} unreachable, purged {} stale free slots, re-queuing job={}",
                        slot.identityKey, purged, request.getJobId());
                pendingJobs.add(request);
                return;
            }

            ActiveJob active = new ActiveJob(slot.identity, slot.identityKey,
                    request, request.getTicker(), request.getType());
            activeJobs.put(request.getJobId(), active);

            LOG.info("Dispatched job: jobId={}, ticker={}, type={}, worker={}, slotType={}",
                    request.getJobId(), request.getTicker(), request.getType(),
                    slot.identityKey, slot.slotType);
        } catch (Exception e) {
            LOG.error("Error dispatching job", e);
            freeSlots.addLast(slot);
        }
    }

    private void sendStop(byte[] workerIdentity, String jobId) {
        try {
            WorkStop stop = WorkStop.newBuilder().setJobId(jobId).build();
            zmqManager.sendToWorker(workerIdentity, PROTOCOL_VERSION, MSG_TYPE_WORK_STOP, stop.toByteArray());
        } catch (Exception e) {
            LOG.error("Error sending STOP for jobId={}", jobId, e);
        }
    }

    // ── Timeout checking ──────────────────────────────────────────────────────

    private void checkTimeouts() {
        long now = System.currentTimeMillis();
        List<String> timedOut = new ArrayList<>();

        for (Map.Entry<String, ActiveJob> entry : activeJobs.entrySet()) {
            ActiveJob job = entry.getValue();
            long timeout = job.type == DataRequest.RequestType.REALTIME_TICKS
                    ? HEARTBEAT_TIMEOUT_MS : HISTORICAL_TIMEOUT_MS;
            if (now - job.lastHeartbeat > timeout) {
                timedOut.add(entry.getKey());
            }
        }

        for (String jobId : timedOut) {
            ActiveJob job = activeJobs.remove(jobId);
            if (job == null) continue;
            LOG.warn("Job timed out: jobId={}, ticker={}, type={}, worker={}",
                    jobId, job.ticker, job.type, job.workerIdentityKey);

            DataRequest requeued = job.request.toBuilder()
                    .setJobId(UUID.randomUUID().toString())
                    .build();
            pendingJobs.add(requeued);
            dispatchPending();
        }
    }

    // ── Helpers ───────────────────────────────────────────────────────────────

    /** Extract exchange name from ticker, e.g. "BTC/USDT.BINANCE" → "BINANCE" */
    private static String exchangeOf(String ticker) {
        int dot = ticker.lastIndexOf('.');
        return dot >= 0 ? ticker.substring(dot + 1).toUpperCase() : "";
    }

    /**
     * Remove all free slots offered by a given worker identity.
     * Called when a dispatch to that identity fails (ROUTER_MANDATORY unreachable).
     * Returns the number of slots removed.
     */
    private int purgeWorkerSlots(String identityKey) {
        int before = freeSlots.size();
        freeSlots.removeIf(slot -> slot.identityKey.equals(identityKey));
        return before - freeSlots.size();
    }

    /**
     * Find and remove a free slot that supports the given exchange and request type.
     * A slot with SlotType.ANY matches any request type.
     */
    private WorkerSlot findFreeSlot(String exchange, DataRequest.RequestType requestType) {
        for (WorkerSlot slot : freeSlots) {
            boolean exchangeMatch = exchange.isEmpty() || slot.exchange.equals(exchange);
            boolean typeMatch = slot.slotType == SlotType.ANY
                    || (slot.slotType == SlotType.HISTORICAL
                            && requestType == DataRequest.RequestType.HISTORICAL_OHLC)
                    || (slot.slotType == SlotType.REALTIME
                            && requestType == DataRequest.RequestType.REALTIME_TICKS);
            if (exchangeMatch && typeMatch) {
                freeSlots.remove(slot);
                return slot;
            }
        }
        return null;
    }

    private static String bytesToHex(byte[] bytes) {
        StringBuilder sb = new StringBuilder();
        for (byte b : bytes) sb.append(String.format("%02x", b));
        return sb.toString();
    }

    @Override
    public void close() {
        stop();
    }

    // ── Inner types ───────────────────────────────────────────────────────────

    /** One available work slot offered by an ingestor via WorkerReady. */
    private static class WorkerSlot {
        final byte[] identity;
        final String identityKey;
        final String exchange;
        final SlotType slotType;

        WorkerSlot(byte[] identity, String identityKey, String exchange, SlotType slotType) {
            this.identity = identity;
            this.identityKey = identityKey;
            this.exchange = exchange;
            this.slotType = slotType;
        }
    }

    private static class ActiveJob {
        final byte[] workerIdentity;
        final String workerIdentityKey;
        final DataRequest request;
        final String ticker;
        final DataRequest.RequestType type;
        long lastHeartbeat;

        ActiveJob(byte[] workerIdentity, String workerIdentityKey,
                  DataRequest request, String ticker, DataRequest.RequestType type) {
            this.workerIdentity = workerIdentity;
            this.workerIdentityKey = workerIdentityKey;
            this.request = request;
            this.ticker = ticker;
            this.type = type;
            this.lastHeartbeat = System.currentTimeMillis();
        }
    }
}
