package com.dexorder.flink.publisher;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.zeromq.SocketType;
import org.zeromq.ZContext;
import org.zeromq.ZMQ;

/**
 * Runs in the job manager. Pulls notifications from task managers (via PUSH/PULL)
 * and republishes them on the MARKET_DATA_PUB socket that the relay subscribes to.
 *
 * Flow:
 *   Task manager HistoryNotificationPublisher  →  PUSH
 *                                               ↓
 *   Job manager HistoryNotificationForwarder   PULL  →  MARKET_DATA_PUB
 *                                                        ↓
 *   Relay (XSUB) → Relay (XPUB) → Clients
 */
public class HistoryNotificationForwarder implements AutoCloseable {
    private static final Logger LOG = LoggerFactory.getLogger(HistoryNotificationForwarder.class);

    private final ZMQ.Socket pullSocket;
    private final ZMQ.Socket pubSocket;
    private final ZContext context;
    private volatile boolean running = true;
    private Thread thread;

    /**
     * @param pullPort   Port to bind PULL socket on (task managers connect PUSH here)
     * @param pubSocket  Existing MARKET_DATA_PUB socket from ZmqChannelManager
     */
    public HistoryNotificationForwarder(int pullPort, ZMQ.Socket pubSocket) {
        this.pubSocket = pubSocket;
        this.context = new ZContext();
        this.pullSocket = context.createSocket(SocketType.PULL);
        this.pullSocket.setRcvHWM(10000);

        String endpoint = "tcp://*:" + pullPort;
        this.pullSocket.bind(endpoint);
        LOG.info("HistoryNotificationForwarder PULL socket bound to {}", endpoint);
    }

    public void start() {
        thread = new Thread(this::forwardLoop, "notification-forwarder");
        thread.setDaemon(true);
        thread.start();
        LOG.info("HistoryNotificationForwarder started");
    }

    private void forwardLoop() {
        LOG.info("Notification forwarder loop running");
        pullSocket.setReceiveTimeOut(200); // ms, so we can check running flag

        while (running) {
            // Receive all frames of a multi-part message and forward to PUB
            byte[] frame = pullSocket.recv(0);
            if (frame == null) {
                continue; // timeout, check running flag
            }

            boolean more = pullSocket.hasReceiveMore();
            if (more) {
                pubSocket.sendMore(frame);
            } else {
                pubSocket.send(frame, 0);
                continue;
            }

            // Receive remaining frames
            while (more) {
                frame = pullSocket.recv(0);
                more = pullSocket.hasReceiveMore();
                if (more) {
                    pubSocket.sendMore(frame);
                } else {
                    pubSocket.send(frame, 0);
                }
            }

            LOG.debug("Forwarded notification to MARKET_DATA_PUB");
        }

        LOG.info("Notification forwarder loop stopped");
    }

    @Override
    public void close() {
        running = false;
        if (thread != null) {
            try {
                thread.join(2000);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
        }
        pullSocket.close();
        context.close();
    }
}
