package com.dexorder.flink.ingestor;

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

/**
 * Manages the ingestor control channel.
 * Broadcasts control messages to all ingestor workers via ZMQ PUB socket.
 */
public class IngestorControlChannel {
    private static final Logger LOG = LoggerFactory.getLogger(IngestorControlChannel.class);
    private static final byte PROTOCOL_VERSION = 0x01;
    private static final byte MSG_TYPE_INGESTOR_CONTROL = 0x02;

    private final ZmqChannelManager zmqManager;

    public IngestorControlChannel(ZmqChannelManager zmqManager) {
        this.zmqManager = zmqManager;
    }

    /**
     * Cancel a specific data request.
     */
    public void cancelRequest(String requestId) {
        IngestorControlMessage msg = IngestorControlMessage.cancel(requestId);
        broadcastControlMessage(msg);
        LOG.info("Sent CANCEL control message for request: {}", requestId);
    }

    /**
     * Send shutdown signal to all ingestors.
     */
    public void shutdown() {
        IngestorControlMessage msg = IngestorControlMessage.shutdown();
        broadcastControlMessage(msg);
        LOG.info("Sent SHUTDOWN control message to all ingestors");
    }

    /**
     * Update ingestor configuration.
     */
    public void updateConfig(IngestorConfig config) {
        IngestorControlMessage msg = IngestorControlMessage.configUpdate(config);
        broadcastControlMessage(msg);
        LOG.info("Sent CONFIG_UPDATE control message to all ingestors");
    }

    /**
     * Send heartbeat to ingestors.
     */
    public void sendHeartbeat() {
        IngestorControlMessage msg = IngestorControlMessage.heartbeat();
        broadcastControlMessage(msg);
        LOG.debug("Sent HEARTBEAT control message to all ingestors");
    }

    /**
     * Broadcast a control message to all ingestors.
     */
    private void broadcastControlMessage(IngestorControlMessage message) {
        try {
            byte[] protobufData = message.toProtobuf();

            boolean sent = zmqManager.sendMessage(
                ZmqChannelManager.Channel.INGESTOR_CONTROL,
                PROTOCOL_VERSION,
                MSG_TYPE_INGESTOR_CONTROL,
                protobufData
            );

            if (!sent) {
                LOG.error("Failed to send control message: action={}", message.getAction());
            }
        } catch (Exception e) {
            LOG.error("Error broadcasting control message: action={}", message.getAction(), e);
        }
    }

    /**
     * Control message wrapper.
     */
    public static class IngestorControlMessage {
        private final ControlAction action;
        private final String requestId;
        private final IngestorConfig config;

        public enum ControlAction {
            CANCEL,
            SHUTDOWN,
            CONFIG_UPDATE,
            HEARTBEAT
        }

        private IngestorControlMessage(ControlAction action, String requestId, IngestorConfig config) {
            this.action = action;
            this.requestId = requestId;
            this.config = config;
        }

        public static IngestorControlMessage cancel(String requestId) {
            return new IngestorControlMessage(ControlAction.CANCEL, requestId, null);
        }

        public static IngestorControlMessage shutdown() {
            return new IngestorControlMessage(ControlAction.SHUTDOWN, null, null);
        }

        public static IngestorControlMessage configUpdate(IngestorConfig config) {
            return new IngestorControlMessage(ControlAction.CONFIG_UPDATE, null, config);
        }

        public static IngestorControlMessage heartbeat() {
            return new IngestorControlMessage(ControlAction.HEARTBEAT, null, null);
        }

        public ControlAction getAction() {
            return action;
        }

        public String getRequestId() {
            return requestId;
        }

        public IngestorConfig getConfig() {
            return config;
        }

        /**
         * Serialize to protobuf bytes.
         * TODO: Replace with actual generated protobuf serialization
         */
        public byte[] toProtobuf() {
            // Placeholder - will be replaced with actual protobuf serialization
            return new byte[0];
        }
    }

    /**
     * Ingestor configuration.
     */
    public static class IngestorConfig {
        private final Integer maxConcurrent;
        private final Integer timeoutSeconds;
        private final String kafkaTopic;

        public IngestorConfig(Integer maxConcurrent, Integer timeoutSeconds, String kafkaTopic) {
            this.maxConcurrent = maxConcurrent;
            this.timeoutSeconds = timeoutSeconds;
            this.kafkaTopic = kafkaTopic;
        }

        public Integer getMaxConcurrent() {
            return maxConcurrent;
        }

        public Integer getTimeoutSeconds() {
            return timeoutSeconds;
        }

        public String getKafkaTopic() {
            return kafkaTopic;
        }
    }
}
