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 java.io.Closeable;
import java.util.HashMap;
import java.util.Map;

/**
 * Manages all ZeroMQ channels for the Flink application.
 * Each channel is bound to a specific port and socket type.
 */
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 {
        INGESTOR_WORK_QUEUE,
        INGESTOR_RESPONSE,
        INGESTOR_CONTROL,
        MARKET_DATA_PUB,
        CLIENT_REQUEST,
        CEP_WEBHOOK
    }

    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. Ingestor Work Queue - PUB socket for topic-based work distribution (exchange prefix filtering)
        createAndBind(
            Channel.INGESTOR_WORK_QUEUE,
            SocketType.PUB,
            bindAddress + ":" + config.getIngestorWorkQueuePort(),
            "Ingestor Work Queue (PUB)"
        );

        // 2. Ingestor Response - ROUTER socket for receiving historical data responses
        createAndBind(
            Channel.INGESTOR_RESPONSE,
            SocketType.ROUTER,
            bindAddress + ":" + config.getIngestorResponsePort(),
            "Ingestor Response (ROUTER)"
        );

        // 3. Ingestor Control - PUB socket for broadcast control messages
        createAndBind(
            Channel.INGESTOR_CONTROL,
            SocketType.PUB,
            bindAddress + ":" + config.getIngestorControlPort(),
            "Ingestor Control (PUB)"
        );

        // 4. Market Data Publication - PUB socket for market data streaming
        createAndBind(
            Channel.MARKET_DATA_PUB,
            SocketType.PUB,
            bindAddress + ":" + config.getMarketDataPubPort(),
            "Market Data Publication (PUB)"
        );

        // 5. Client Request - REP socket for request-response
        createAndBind(
            Channel.CLIENT_REQUEST,
            SocketType.REP,
            bindAddress + ":" + config.getClientRequestPort(),
            "Client Request (REP)"
        );

        // 6. CEP Webhook - ROUTER socket for async callbacks
        createAndBind(
            Channel.CEP_WEBHOOK,
            SocketType.ROUTER,
            bindAddress + ":" + config.getCepWebhookPort(),
            "CEP Webhook (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);

            // Set socket options
            socket.setLinger(1000); // 1 second linger on close
            socket.setSndHWM(10000); // High water mark for outbound messages
            socket.setRcvHWM(10000); // High water mark for inbound messages

            // Bind the socket
            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);
        }
    }

    /**
     * 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 the specified channel.
     *
     * @param channel The channel to send on
     * @param versionByte Protocol version byte
     * @param messageTypeByte Message type ID byte
     * @param protobufData Serialized protobuf message
     * @return true if sent successfully
     */
    public boolean sendMessage(Channel channel, byte versionByte, byte messageTypeByte, byte[] protobufData) {
        ZMQ.Socket socket = getSocket(channel);

        // Send as two frames: [version byte] [type byte + protobuf data]
        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 message with a topic prefix (for PUB sockets).
     *
     * @param channel The channel to send on
     * @param topic Topic string for subscription filtering
     * @param versionByte Protocol version byte
     * @param messageTypeByte Message type ID byte
     * @param protobufData Serialized protobuf message
     * @return true if sent successfully
     */
    public boolean sendTopicMessage(Channel channel, String topic, byte versionByte, byte messageTypeByte, byte[] protobufData) {
        ZMQ.Socket socket = getSocket(channel);

        // Send as three frames: [topic] [version byte] [type byte + protobuf data]
        boolean sentTopic = socket.send(topic.getBytes(ZMQ.CHARSET), ZMQ.SNDMORE);
        if (!sentTopic) {
            LOG.error("Failed to send topic frame on channel {}", channel);
            return false;
        }

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

    /**
     * Receive a message from a ROUTER socket.
     * Returns a ReceivedMessage containing the identity, version, type, and payload.
     *
     * @param channel The channel to receive from (must be ROUTER)
     * @param timeout Timeout in milliseconds (0 for non-blocking, -1 for blocking)
     * @return ReceivedMessage or null if no message available
     */
    public ReceivedMessage receiveRouterMessage(Channel channel, int timeout) {
        ZMQ.Socket socket = getSocket(channel);

        // Set receive timeout
        if (timeout >= 0) {
            socket.setReceiveTimeOut(timeout);
        }

        // Receive identity frame
        byte[] identity = socket.recv(0);
        if (identity == null) {
            return null;
        }

        // Receive version frame
        byte[] versionFrame = socket.recv(0);
        if (versionFrame == null || versionFrame.length != 1) {
            LOG.error("Invalid version frame received on channel {}", channel);
            return null;
        }

        // Receive message frame (type byte + protobuf data)
        byte[] messageFrame = socket.recv(0);
        if (messageFrame == null || messageFrame.length < 1) {
            LOG.error("Invalid message frame received on channel {}", channel);
            return null;
        }

        byte versionByte = versionFrame[0];
        byte messageTypeByte = messageFrame[0];
        byte[] protobufData = new byte[messageFrame.length - 1];
        System.arraycopy(messageFrame, 1, protobufData, 0, protobufData.length);

        return new ReceivedMessage(identity, versionByte, messageTypeByte, protobufData);
    }

    /**
     * Represents a received message from a ROUTER socket.
     */
    public static class ReceivedMessage {
        private final byte[] identity;
        private final byte version;
        private final byte messageType;
        private final byte[] protobufData;

        public ReceivedMessage(byte[] identity, byte version, byte messageType, byte[] protobufData) {
            this.identity = identity;
            this.version = version;
            this.messageType = messageType;
            this.protobufData = protobufData;
        }

        public byte[] getIdentity() {
            return identity;
        }

        public byte getVersion() {
            return version;
        }

        public byte getMessageType() {
            return messageType;
        }

        public byte[] getProtobufData() {
            return protobufData;
        }
    }

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