// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0

//! Shared ZMQ wire-format helpers for the `state::kv_events` component
//! tests. Encodes events in the same msgspec layout SGLang emits, builds
//! the two-frame `[seq, payload]` ZMQ message a real publisher sends, and
//! binds a loopback PUB socket on an OS-assigned port.

#![allow(dead_code)]

use bytes::Bytes;
use rmp::encode as mp;
use zeromq::{Endpoint, PubSocket, Socket, ZmqMessage};

/// Bind a PUB socket to an OS-assigned 127.0.0.1 port. Returns
/// `(socket, port)`.
pub async fn make_pub_bound() -> (PubSocket, u16) {
    let mut sock = PubSocket::new();
    let endpoint = sock
        .bind("tcp://127.0.0.1:0")
        .await
        .expect("bind PUB socket");
    let port = match endpoint {
        Endpoint::Tcp(_, p) => p,
        other => panic!("unexpected endpoint: {other:?}"),
    };
    (sock, port)
}

/// Encode a single `BlockStored` event in the wire format msgspec emits: a
/// tagged map `{"type": "BlockStored", "block_hashes", "parent_block_hash",
/// "token_ids", "block_size", "lora_id", "medium"}`.
pub fn encode_block_stored_event(
    block_hashes: &[i64],
    parent: Option<i64>,
    token_ids: &[u32],
    block_size: u32,
) -> Vec<u8> {
    let mut buf = Vec::new();
    mp::write_map_len(&mut buf, 7).unwrap();
    mp::write_str(&mut buf, "type").unwrap();
    mp::write_str(&mut buf, "BlockStored").unwrap();
    mp::write_str(&mut buf, "block_hashes").unwrap();
    mp::write_array_len(&mut buf, block_hashes.len() as u32).unwrap();
    for v in block_hashes {
        mp::write_sint(&mut buf, *v).unwrap();
    }
    mp::write_str(&mut buf, "parent_block_hash").unwrap();
    match parent {
        Some(v) => {
            mp::write_sint(&mut buf, v).unwrap();
        }
        None => mp::write_nil(&mut buf).unwrap(),
    }
    mp::write_str(&mut buf, "token_ids").unwrap();
    mp::write_array_len(&mut buf, token_ids.len() as u32).unwrap();
    for v in token_ids {
        mp::write_uint(&mut buf, *v as u64).unwrap();
    }
    mp::write_str(&mut buf, "block_size").unwrap();
    mp::write_uint(&mut buf, block_size as u64).unwrap();
    mp::write_str(&mut buf, "lora_id").unwrap();
    mp::write_nil(&mut buf).unwrap();
    mp::write_str(&mut buf, "medium").unwrap();
    mp::write_str(&mut buf, "GPU").unwrap();
    buf
}

/// Wrap one or more pre-encoded events into a KVEventBatch with
/// timestamp + optional dp-rank.
pub fn encode_event_batch(ts: f64, events: Vec<Vec<u8>>, attn_dp_rank: Option<u32>) -> Vec<u8> {
    let mut buf = Vec::new();
    mp::write_array_len(&mut buf, 3).unwrap();
    mp::write_f64(&mut buf, ts).unwrap();
    mp::write_array_len(&mut buf, events.len() as u32).unwrap();
    for ev in events {
        buf.extend_from_slice(&ev);
    }
    match attn_dp_rank {
        Some(v) => {
            mp::write_uint(&mut buf, v as u64).unwrap();
        }
        None => mp::write_nil(&mut buf).unwrap(),
    }
    buf
}

/// Build the two-frame ZMQ message a real KV publisher sends:
/// `[seq (big-endian i64), payload]`.
pub fn build_multipart(seq: i64, payload: Vec<u8>) -> ZmqMessage {
    let mut msg = ZmqMessage::from(Bytes::new());
    msg.push_back(Bytes::copy_from_slice(&seq.to_be_bytes()));
    msg.push_back(Bytes::from(payload));
    msg
}