// ZeroMQ DEALER client connecting to Flink IngestorBroker (ROUTER, port 5567)
import * as zmq from 'zeromq';
import os from 'os';
import {
  DataRequest,
  WorkerReady, WorkComplete, WorkHeartbeat, WorkReject, WorkStop,
  MessageTypeId, PROTOCOL_VERSION
} from './proto/messages.js';

const PROTOCOL_VERSION_BUF = Buffer.from([PROTOCOL_VERSION]);

/**
 * Encodes a broker protocol message for sending over DEALER.
 * Frame layout (DEALER → ROUTER):
 *   Frame 0: empty delimiter (required for ROUTER peering)
 *   Frame 1: [0x01] version
 *   Frame 2: [typeId][protobuf bytes]
 */
function encodeBrokerMessage(typeId, messageData, MessageType) {
  const protoBytes = MessageType.encode(MessageType.create(messageData)).finish();
  const frame2 = Buffer.concat([Buffer.from([typeId]), Buffer.from(protoBytes)]);
  return [Buffer.alloc(0), PROTOCOL_VERSION_BUF, frame2];
}

export class ZmqClient {
  constructor(config, logger) {
    this.config = config;
    this.logger = logger;

    this.dealerSocket = null;
    this.isShutdown = false;

    // Callbacks set by IngestorWorker / SlotPool
    this.onWorkAssign = null;  // (DataRequest) => void
    this.onWorkStop = null;    // (jobId) => void
    this.onConnected = null;   // async () => void — fires on initial connect AND reconnect
  }

  /**
   * Connect DEALER socket to Flink IngestorBroker (ROUTER).
   * Fires onConnected on every TCP (re)connect so SlotPool can re-offer slots.
   */
  async connect() {
    const { flink_hostname, ingestor_broker_port = 5567 } = this.config;

    this.dealerSocket = new zmq.Dealer();

    // Set a stable routing ID so the ROUTER always recognises this peer, even
    // after a TCP reconnect.  Without this, ZMQ assigns a new random identity
    // on every reconnect, leaving dead slots in the broker's queue → Errno 65.
    const stableId = os.hostname();
    this.dealerSocket.routingId = stableId;
    this.logger.info({ routingId: stableId }, 'Set stable DEALER routing ID');

    // Subscribe to connection events BEFORE calling connect() so we catch the
    // initial establishment.  The 'connect' event fires on initial TCP handshake
    // and again after every ZMQ reconnect (e.g. Flink restart).
    this.dealerSocket.events.on('connect', ({ address }) => {
      this.logger.info({ address }, 'DEALER connected to broker');
      if (this.onConnected) {
        this.onConnected().catch(err =>
          this.logger.error({ error: err.message }, 'onConnected callback failed'));
      }
    });

    const endpoint = `tcp://${flink_hostname}:${ingestor_broker_port}`;
    this.dealerSocket.connect(endpoint);
    this.logger.info(`Connecting DEALER to Flink IngestorBroker at ${endpoint}`);

    // Start receiving work in background
    this._receiveLoop();
  }

  /**
   * Send one typed WorkerReady slot offer.
   * @param {string} exchange - Exchange name (e.g. 'BINANCE')
   * @param {number} slotType - SlotType enum value (0=ANY, 1=HISTORICAL, 2=REALTIME)
   */
  async sendTypedReady(exchange, slotType) {
    const frames = encodeBrokerMessage(
      MessageTypeId.WORKER_READY,
      { exchanges: [exchange], jobType: slotType },
      WorkerReady
    );
    await this.dealerSocket.send(frames);
    this.logger.debug({ exchange, slotType }, 'Sent WorkerReady slot offer');
  }

  /**
   * Send WorkComplete after a historical job finishes.
   * Slot re-registration is handled by SlotPool after this call.
   */
  async sendComplete(jobId, success, errorMessage) {
    const frames = encodeBrokerMessage(
      MessageTypeId.WORK_COMPLETE,
      {
        jobId,
        success,
        ...(errorMessage ? { errorMessage } : {})
      },
      WorkComplete
    );
    await this.dealerSocket.send(frames);
    this.logger.info({ jobId, success }, 'Sent WorkComplete');
  }

  /**
   * Send WorkHeartbeat for an active realtime job.
   */
  async sendHeartbeat(jobId) {
    const frames = encodeBrokerMessage(
      MessageTypeId.WORK_HEARTBEAT,
      { jobId },
      WorkHeartbeat
    );
    await this.dealerSocket.send(frames);
    this.logger.debug({ jobId }, 'Sent WorkHeartbeat');
  }

  /**
   * Send WorkReject if we cannot handle the dispatched job.
   */
  async sendReject(jobId, reason) {
    const frames = encodeBrokerMessage(
      MessageTypeId.WORK_REJECT,
      { jobId, reason },
      WorkReject
    );
    await this.dealerSocket.send(frames);
    this.logger.warn({ jobId, reason }, 'Sent WorkReject');
  }

  /**
   * Background loop: receive WorkAssign (DataRequest) or WorkStop from Flink.
   * ROUTER→DEALER frame layout: [empty][version][typeId+payload]
   */
  async _receiveLoop() {
    try {
      for await (const frames of this.dealerSocket) {
        if (this.isShutdown) break;

        try {
          // frames[0] = empty delimiter, frames[1] = version, frames[2] = type+payload
          if (frames.length < 3) {
            this.logger.warn({ frameCount: frames.length }, 'Unexpected frame count from broker');
            continue;
          }

          const versionByte = frames[1][0];
          if (versionByte !== PROTOCOL_VERSION) {
            this.logger.warn({ versionByte }, 'Unexpected protocol version from broker');
            continue;
          }

          const typeId = frames[2][0];
          const payload = frames[2].slice(1);

          if (typeId === MessageTypeId.WORK_ASSIGN) {
            const request = DataRequest.decode(payload);
            const req = DataRequest.toObject(request, {
              longs: String, enums: String, bytes: Buffer
            });
            this.logger.info(
              { jobId: req.jobId, requestId: req.requestId, type: req.type, ticker: req.ticker },
              'Received WorkAssign from broker'
            );
            if (this.onWorkAssign) {
              this.onWorkAssign(req);
            }

          } else if (typeId === MessageTypeId.WORK_STOP) {
            const stop = WorkStop.decode(payload);
            const { jobId } = WorkStop.toObject(stop);
            this.logger.info({ jobId }, 'Received WorkStop from broker');
            if (this.onWorkStop) {
              this.onWorkStop(jobId);
            }

          } else {
            this.logger.warn({ typeId: `0x${typeId.toString(16)}` }, 'Unknown message type from broker');
          }

        } catch (err) {
          this.logger.error({ error: err.message }, 'Error processing broker message');
        }
      }
    } catch (err) {
      if (!this.isShutdown) {
        this.logger.error({ error: err.message }, 'DEALER receive loop error');
      }
    }
  }

  async shutdown() {
    this.isShutdown = true;
    this.logger.info('Shutting down ZMQ DEALER connection');
    if (this.dealerSocket) {
      this.dealerSocket.close();
    }
  }
}
