package com.dexorder.flink.ingestor;

import com.dexorder.flink.zmq.ZmqChannelManager;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CompletableFuture;
import java.util.function.Consumer;

/**
 * Listens for DataResponse messages from ingestors on the ROUTER socket.
 * Matches responses to pending requests and delivers them to waiting handlers.
 */
public class IngestorResponseListener {
    private static final Logger LOG = LoggerFactory.getLogger(IngestorResponseListener.class);
    private static final byte PROTOCOL_VERSION = 0x01;
    private static final byte MSG_TYPE_DATA_RESPONSE = 0x02;

    private final ZmqChannelManager zmqManager;
    private final Map<String, CompletableFuture<DataResponseMessage>> pendingRequests;
    private volatile boolean running;
    private Thread listenerThread;

    public IngestorResponseListener(ZmqChannelManager zmqManager) {
        this.zmqManager = zmqManager;
        this.pendingRequests = new ConcurrentHashMap<>();
        this.running = false;
    }

    /**
     * Start the response listener thread.
     */
    public void start() {
        if (running) {
            LOG.warn("IngestorResponseListener already running");
            return;
        }

        running = true;
        listenerThread = new Thread(this::listenLoop, "IngestorResponseListener-Thread");
        listenerThread.setDaemon(false);
        listenerThread.start();
        LOG.info("IngestorResponseListener started");
    }

    /**
     * Stop the response listener.
     */
    public void stop() {
        if (!running) {
            return;
        }

        running = false;
        if (listenerThread != null) {
            listenerThread.interrupt();
            try {
                listenerThread.join(5000);
            } catch (InterruptedException e) {
                LOG.warn("Interrupted while waiting for listener thread to stop", e);
                Thread.currentThread().interrupt();
            }
        }

        // Cancel all pending requests
        pendingRequests.values().forEach(future ->
            future.completeExceptionally(new Exception("Listener stopped"))
        );
        pendingRequests.clear();

        LOG.info("IngestorResponseListener stopped");
    }

    /**
     * Register a request and return a CompletableFuture that will be completed
     * when the response arrives.
     */
    public CompletableFuture<DataResponseMessage> registerRequest(String requestId) {
        CompletableFuture<DataResponseMessage> future = new CompletableFuture<>();
        pendingRequests.put(requestId, future);
        LOG.debug("Registered pending request: {}", requestId);
        return future;
    }

    /**
     * Cancel a pending request.
     */
    public void cancelRequest(String requestId) {
        CompletableFuture<DataResponseMessage> future = pendingRequests.remove(requestId);
        if (future != null) {
            future.completeExceptionally(new Exception("Request cancelled"));
            LOG.debug("Cancelled pending request: {}", requestId);
        }
    }

    /**
     * Main listener loop - receives and processes DataResponse messages.
     */
    private void listenLoop() {
        LOG.info("IngestorResponseListener loop started");

        while (running) {
            try {
                // Receive message from ROUTER socket with 1 second timeout
                ZmqChannelManager.ReceivedMessage receivedMsg = zmqManager.receiveRouterMessage(
                    ZmqChannelManager.Channel.INGESTOR_RESPONSE,
                    1000
                );

                if (receivedMsg == null) {
                    continue;
                }

                // Verify protocol version and message type
                if (receivedMsg.getVersion() != PROTOCOL_VERSION) {
                    LOG.warn("Received message with unsupported protocol version: {}",
                        receivedMsg.getVersion());
                    continue;
                }

                if (receivedMsg.getMessageType() != MSG_TYPE_DATA_RESPONSE) {
                    LOG.warn("Received unexpected message type: {}",
                        receivedMsg.getMessageType());
                    continue;
                }

                // Parse the DataResponse
                DataResponseMessage response = DataResponseMessage.fromProtobuf(
                    receivedMsg.getProtobufData()
                );

                processResponse(response);

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

        LOG.info("IngestorResponseListener loop stopped");
    }

    /**
     * Process a received DataResponse message.
     */
    private void processResponse(DataResponseMessage response) {
        String requestId = response.getRequestId();

        CompletableFuture<DataResponseMessage> future = pendingRequests.remove(requestId);
        if (future == null) {
            LOG.warn("Received response for unknown request: {}", requestId);
            return;
        }

        LOG.info("Received response for request: {}, status={}, records={}",
            requestId, response.getStatus(), response.getTotalRecords());

        // Complete the future with the response
        future.complete(response);
    }

    public boolean isRunning() {
        return running;
    }

    public int getPendingRequestCount() {
        return pendingRequests.size();
    }
}
