package com.dexorder.flink.ingestor;

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

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * Monitors XPUB subscription events from the relay and manages realtime ingestor lifecycle.
 *
 * This class is the <em>sole owner</em> of the MARKET_DATA_PUB XPUB socket. All outbound
 * publishes from other threads (e.g., HistoryNotificationForwarder, RealtimeOHLCPublisher)
 * must go through {@link #enqueuePublish(byte[]...)} so they are sent from the single loop
 * thread — ZMQ sockets are not thread-safe.
 *
 * Topic format: {@code {ticker}|ohlc:{period_seconds}}
 * Example:      {@code BTC/USDT.BINANCE|ohlc:60}
 *
 * Reference counting:
 *   tickerRefs — across all periods for a ticker; 0→1 triggers ingestor activation
 *   topicRefs  — per (ticker, period); consulted by RealtimeOHLCPublisher to filter output
 */
public class RealtimeSubscriptionManager implements AutoCloseable {
    private static final Logger LOG = LoggerFactory.getLogger(RealtimeSubscriptionManager.class);

    private static final Pattern TOPIC_PATTERN = Pattern.compile("^(.+)\\|ohlc:(\\d+)$");

    private final ZmqChannelManager zmqManager;
    private final ZMQ.Socket xpubSocket;
    private final IngestorBroker broker;

    /** Per-ticker reference count (across all subscribed periods for that ticker) */
    private final Map<String, Integer> tickerRefs = new HashMap<>();

    /** Per-topic reference count (ticker|ohlc:period → subscriber count) */
    private final Map<String, Integer> topicRefs = new HashMap<>();

    /**
     * Thread-safe outbound publish queue.
     * Each entry is one multi-frame message: {@code byte[][] frames}.
     */
    private final ConcurrentLinkedQueue<byte[][]> publishQueue = new ConcurrentLinkedQueue<>();

    private volatile boolean running;
    private Thread thread;

    public RealtimeSubscriptionManager(ZmqChannelManager zmqManager, IngestorBroker broker) {
        this.zmqManager = zmqManager;
        this.xpubSocket = zmqManager.getSocket(ZmqChannelManager.Channel.MARKET_DATA_PUB);
        this.broker = broker;
    }

    /**
     * Queue a multi-frame message for publication on MARKET_DATA_PUB.
     * Thread-safe — may be called from any thread (HistoryNotificationForwarder,
     * RealtimeOHLCPublisher, etc.).
     */
    public void enqueuePublish(byte[]... frames) {
        publishQueue.add(frames);
    }

    /**
     * Returns the current subscriber count for a topic.
     * Thread-safe for reads (value is written only from the loop thread but read from others).
     */
    public int getTopicRefCount(String topic) {
        return topicRefs.getOrDefault(topic, 0);
    }

    public void start() {
        if (running) {
            LOG.warn("RealtimeSubscriptionManager already running");
            return;
        }
        running = true;
        thread = new Thread(this::subscriptionLoop, "RealtimeSubscriptionManager");
        thread.setDaemon(false);
        thread.start();
        LOG.info("RealtimeSubscriptionManager started");
    }

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

    private void subscriptionLoop() {
        // Build a poller so we can block-wait rather than busy-spin
        ZMQ.Poller poller = zmqManager.createPoller(1);
        poller.register(xpubSocket, ZMQ.Poller.POLLIN);

        LOG.info("RealtimeSubscriptionManager loop running");

        while (running) {
            try {
                // 1. Flush any queued outbound messages before blocking
                byte[][] frames;
                while ((frames = publishQueue.poll()) != null) {
                    sendFrames(frames);
                }

                // 2. Wait up to 50ms for a subscription event
                poller.poll(50);

                // 3. Drain all available subscription events
                if (poller.pollin(0)) {
                    byte[] event;
                    while ((event = xpubSocket.recv(ZMQ.DONTWAIT)) != null) {
                        if (event.length > 0) {
                            processSubscriptionEvent(event);
                        }
                    }
                }
            } catch (Exception e) {
                if (running) {
                    LOG.error("Error in subscription loop", e);
                }
            }
        }

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

    private void sendFrames(byte[][] frames) {
        for (int i = 0; i < frames.length; i++) {
            if (i < frames.length - 1) {
                xpubSocket.sendMore(frames[i]);
            } else {
                xpubSocket.send(frames[i], 0);
            }
        }
    }

    private void processSubscriptionEvent(byte[] event) {
        // XPUB subscription frame: first byte is 0x01 (subscribe) or 0x00 (unsubscribe);
        // remaining bytes are the raw topic string.
        boolean isSubscribe = event[0] == 0x01;
        String topic = new String(event, 1, event.length - 1, ZMQ.CHARSET);

        Matcher m = TOPIC_PATTERN.matcher(topic);
        if (!m.matches()) {
            // Not a realtime OHLC topic — e.g. RESPONSE: or HISTORY_READY: prefixes
            LOG.debug("Ignoring subscription event for non-realtime topic: action={}, topic={}",
                    isSubscribe ? "subscribe" : "unsubscribe", topic);
            return;
        }

        String ticker = m.group(1);
        LOG.info("Subscription event: action={}, topic={}", isSubscribe ? "subscribe" : "unsubscribe", topic);

        if (isSubscribe) {
            handleSubscribe(ticker, topic);
        } else {
            handleUnsubscribe(ticker, topic);
        }
    }

    private void handleSubscribe(String ticker, String topic) {
        int newTopicRef = topicRefs.merge(topic, 1, Integer::sum);
        LOG.debug("topicRefs[{}]={}", topic, newTopicRef);

        int newTickerRef = tickerRefs.merge(ticker, 1, Integer::sum);
        if (newTickerRef == 1) {
            LOG.info("First subscriber for ticker={} — submitting realtime request", ticker);
            broker.submitRealtimeRequest(ticker);
        }
        LOG.debug("tickerRefs[{}]={}", ticker, newTickerRef);
    }

    private void handleUnsubscribe(String ticker, String topic) {
        int newTopicRef = topicRefs.merge(topic, -1, Integer::sum);
        if (newTopicRef <= 0) {
            topicRefs.remove(topic);
        }
        LOG.debug("topicRefs[{}]={}", topic, newTopicRef);

        int newTickerRef = tickerRefs.merge(ticker, -1, Integer::sum);
        if (newTickerRef <= 0) {
            tickerRefs.remove(ticker);
            LOG.info("Last subscriber for ticker={} left — stopping realtime jobs", ticker);
            broker.stopRealtimeJobsForTicker(ticker);
        }
        LOG.debug("tickerRefs[{}]={}", ticker, newTickerRef);
    }

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