package com.dexorder.flink.zmq;

import com.dexorder.flink.config.AppConfig;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.zeromq.SocketType;
import org.zeromq.ZContext;
import org.zeromq.ZMQ;
import org.zeromq.ZMQException;

import java.io.Closeable;
import java.util.HashMap;
import java.util.Map;

/**
 * Manages all ZeroMQ channels for the Flink application.
 *
 * Port layout:
 *   5558  XPUB  MARKET_DATA_PUB   — market data + notifications to clients (via relay XSUB)
 *                                   XPUB exposes subscription frames so Flink can detect
 *                                   which realtime topics clients are interested in.
 *   5561  PULL  (internal)         — task manager → job manager notifications (unchanged)
 *   5566  PULL  CLIENT_REQUEST     — receives forwarded SubmitHistoricalRequest from relay PUSH
 *   5567  ROUTER INGESTOR_BROKER   — exclusive work queue; ingestors connect with DEALER
 */
public class ZmqChannelManager implements Closeable {
    private static final Logger LOG = LoggerFactory.getLogger(ZmqChannelManager.class);

    private final ZContext context;
    private final Map<String, ZMQ.Socket> sockets;
    private final AppConfig config;

    public enum Channel {
        MARKET_DATA_PUB,
        CLIENT_REQUEST,
        INGESTOR_BROKER,
    }

    public ZmqChannelManager(AppConfig config) {
        this.config = config;
        this.context = new ZContext();
        this.sockets = new HashMap<>();
    }

    /**
     * Initialize and bind all ZMQ channels.
     */
    public void initializeChannels() {
        String bindAddress = config.getBindAddress();

        LOG.info("Initializing ZeroMQ channels on {}", bindAddress);

        // 1. Market Data Publication — XPUB so subscription events are visible to Flink
        //    Relay's XSUB connects here to proxy data to clients.
        //    Subscription frames from relay (forwarded from clients) arrive as readable messages.
        ZMQ.Socket marketDataSocket = context.createSocket(SocketType.XPUB);
        marketDataSocket.setXpubVerbose(true);   // emit every sub/unsub, not just first/last
        marketDataSocket.setLinger(1000);
        marketDataSocket.setSndHWM(10000);
        marketDataSocket.setRcvHWM(10000);
        String marketDataEndpoint = bindAddress + ":" + config.getMarketDataPubPort();
        marketDataSocket.bind(marketDataEndpoint);
        sockets.put(Channel.MARKET_DATA_PUB.name(), marketDataSocket);
        LOG.info("Bound Market Data Publication (XPUB) to {}", marketDataEndpoint);

        // 2. Client Request Pull — receives SubmitHistoricalRequest forwarded by relay PUSH
        createAndBind(
            Channel.CLIENT_REQUEST,
            SocketType.PULL,
            bindAddress + ":" + config.getFlinkRequestPullPort(),
            "Client Request (PULL)"
        );

        // 3. Ingestor Broker — ROUTER for exclusive work dispatch to ingestor DEALER workers
        createAndBind(
            Channel.INGESTOR_BROKER,
            SocketType.ROUTER,
            bindAddress + ":" + config.getIngestorBrokerPort(),
            "Ingestor Broker (ROUTER)"
        );

        LOG.info("All ZeroMQ channels initialized successfully");
    }

    private void createAndBind(Channel channel, SocketType socketType, String endpoint, String description) {
        try {
            ZMQ.Socket socket = context.createSocket(socketType);
            socket.setLinger(1000);
            socket.setSndHWM(10000);
            socket.setRcvHWM(10000);
            if (socketType == SocketType.ROUTER) {
                // Return false (EHOSTUNREACH) instead of silently dropping messages to
                // unknown/disconnected peer identities. Enables immediate stale-slot detection.
                socket.setRouterMandatory(true);
            }
            socket.bind(endpoint);
            sockets.put(channel.name(), socket);
            LOG.info("Bound {} to {}", description, endpoint);
        } catch (Exception e) {
            LOG.error("Failed to bind {} to {}", description, endpoint, e);
            throw new RuntimeException("Failed to initialize ZMQ channel: " + channel, e);
        }
    }

    /**
     * Create a ZMQ Poller backed by this manager's context.
     */
    public ZMQ.Poller createPoller(int size) {
        return context.getContext().poller(size);
    }

    /**
     * Get a socket by channel type.
     */
    public ZMQ.Socket getSocket(Channel channel) {
        ZMQ.Socket socket = sockets.get(channel.name());
        if (socket == null) {
            throw new IllegalStateException("Socket not initialized: " + channel);
        }
        return socket;
    }

    /**
     * Send a message on a channel (no topic prefix — for PULL/PUSH or direct sends).
     */
    public boolean sendMessage(Channel channel, byte versionByte, byte messageTypeByte, byte[] protobufData) {
        ZMQ.Socket socket = getSocket(channel);

        boolean sentFrame1 = socket.send(new byte[]{versionByte}, ZMQ.SNDMORE);
        if (!sentFrame1) {
            LOG.error("Failed to send version frame on channel {}", channel);
            return false;
        }

        byte[] frame2 = new byte[1 + protobufData.length];
        frame2[0] = messageTypeByte;
        System.arraycopy(protobufData, 0, frame2, 1, protobufData.length);

        boolean sentFrame2 = socket.send(frame2, 0);
        if (!sentFrame2) {
            LOG.error("Failed to send message frame on channel {}", channel);
            return false;
        }

        return true;
    }

    /**
     * Send a topic-prefixed message (for XPUB market data publishing).
     * Frame layout: [topic][version][type+payload]
     */
    public boolean sendTopicMessage(Channel channel, String topic, byte versionByte, byte messageTypeByte, byte[] protobufData) {
        ZMQ.Socket socket = getSocket(channel);

        if (!socket.send(topic.getBytes(ZMQ.CHARSET), ZMQ.SNDMORE)) {
            LOG.error("Failed to send topic frame on channel {}", channel);
            return false;
        }

        if (!socket.send(new byte[]{versionByte}, ZMQ.SNDMORE)) {
            LOG.error("Failed to send version frame on channel {}", channel);
            return false;
        }

        byte[] frame2 = new byte[1 + protobufData.length];
        frame2[0] = messageTypeByte;
        System.arraycopy(protobufData, 0, frame2, 1, protobufData.length);

        if (!socket.send(frame2, 0)) {
            LOG.error("Failed to send message frame on channel {}", channel);
            return false;
        }

        return true;
    }

    /**
     * Send a targeted message to a specific DEALER worker via ROUTER.
     * Frame layout: [identity][empty][version][type+payload]
     */
    public boolean sendToWorker(byte[] identity, byte versionByte, byte messageTypeByte, byte[] protobufData) {
        ZMQ.Socket socket = getSocket(Channel.INGESTOR_BROKER);

        try {
            if (!socket.send(identity, ZMQ.SNDMORE)) return false;
            if (!socket.send(new byte[0], ZMQ.SNDMORE)) return false;
            if (!socket.send(new byte[]{versionByte}, ZMQ.SNDMORE)) return false;

            byte[] frame = new byte[1 + protobufData.length];
            frame[0] = messageTypeByte;
            System.arraycopy(protobufData, 0, frame, 1, protobufData.length);

            return socket.send(frame, 0);
        } catch (ZMQException e) {
            // ROUTER_MANDATORY throws (not returns false) for unreachable peer identities.
            // Returning false lets the caller detect and purge the stale slot.
            LOG.debug("sendToWorker failed (errno={}): {}", e.getErrorCode(), e.getMessage());
            return false;
        }
    }

    @Override
    public void close() {
        LOG.info("Closing ZeroMQ channels");
        sockets.values().forEach(ZMQ.Socket::close);
        sockets.clear();
        context.close();
        LOG.info("ZeroMQ context closed");
    }
}
