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;

import java.util.ArrayList;
import java.util.List;
import java.util.function.Consumer;

/**
 * Runs in the job manager. Pulls notifications from task managers (via PUSH/PULL)
 * and enqueues them for publication on MARKET_DATA_PUB via the provided publish callback.
 *
 * The publish callback must be thread-safe (e.g., RealtimeSubscriptionManager.enqueuePublish).
 * Direct socket access is avoided here because the MARKET_DATA_PUB XPUB socket is owned
 * exclusively by RealtimeSubscriptionManager to satisfy ZMQ's single-thread-per-socket rule.
 *
 * Flow:
 *   Task manager HistoryNotificationPublisher  →  PUSH
 *                                               ↓
 *   Job manager HistoryNotificationForwarder   PULL  →  publishCallback (queue)
 *                                                        ↓ (RealtimeSubscriptionManager loop)
 *                                                        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 Consumer<byte[][]> publishCallback;
    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 publishCallback Thread-safe callback to enqueue outbound multi-frame messages
     */
    public HistoryNotificationForwarder(int pullPort, Consumer<byte[][]> publishCallback) {
        this.publishCallback = publishCallback;
        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) {
            byte[] frame = pullSocket.recv(0);
            if (frame == null) {
                continue; // timeout — check running flag
            }

            // Collect all frames of the multi-part message, then enqueue atomically
            List<byte[]> frames = new ArrayList<>();
            frames.add(frame);

            while (pullSocket.hasReceiveMore()) {
                byte[] next = pullSocket.recv(0);
                if (next != null) {
                    frames.add(next);
                }
            }

            publishCallback.accept(frames.toArray(new byte[0][]));
            LOG.debug("Enqueued notification ({} frames) for MARKET_DATA_PUB", frames.size());
        }

        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();
    }
}
