// Wire-format constants and packet builders.
// All integers are little-endian. Must stay in sync with slm_protocol.h.

use lz4_flex::compress_prepend_size;

pub(crate) const MAGIC: [u8; 4]       = [b'S', b'L', b'M', 0x01];
pub(crate) const VERSION: u8           = 4;
pub(crate) const MCAST_PORT: u16       = 9878;
pub(crate) const MCAST_ADDR: [u8; 4]  = [239, 255, 77, 43];
// 1400 bytes fits inside a 1500-byte Ethernet MTU (IP+UDP headers ~28 bytes).
pub(crate) const MAX_PAYLOAD: usize    = 1400;

pub(crate) const PKT_ANNOUNCE: u8  = 0x01;
pub(crate) const PKT_HELLO: u8     = 0x02;
pub(crate) const PKT_CAPS: u8      = 0x03;
pub(crate) const PKT_CHALLENGE: u8 = 0x04; // sender -> receiver: address-validation nonce (v4)
pub(crate) const PKT_VIDEO: u8     = 0x10;
pub(crate) const PKT_AUDIO: u8     = 0x11;
pub(crate) const PKT_FEC: u8       = 0x12;
pub(crate) const PKT_KEEPALIVE: u8 = 0x20;
pub(crate) const PKT_TALLY: u8     = 0x30; // receiver -> sender: tally state (1-byte payload)
pub(crate) const PKT_MSG: u8       = 0x31; // receiver -> sender: UTF-8 operator message
pub(crate) const PKT_STAT: u8      = 0x32; // receiver -> sender: echoed sender_ts_ns for latency
pub(crate) const PKT_BYE: u8       = 0xFF;

// Tally states carried in the 1-byte PKT_TALLY payload.
pub const TALLY_OFF:     u8 = 0; // not in preview or program
pub const TALLY_PREVIEW: u8 = 1; // in the preview slot
pub const TALLY_PROGRAM: u8 = 2; // on air (program slot)

// PKT_MSG payload is a length-prefixed UTF-8 string, truncated to this many bytes.
pub const MSG_MAX_BYTES: usize = 255;

// Byte offset of the payload inside a PKT_VIDEO datagram.
// SLM header (12) + timestamp_ns u64 (8) + frame_seq u32 (4) +
// uncompressed_sz u32 (4) + compressed_sz u32 (4) +
// chunk_index u16 (2) + chunk_total u16 (2) + compression u8 (1) + pad[3] (3) +
// sender_ts_ns u64 (8) = 48.
const VIDEO_HDR_SIZE: usize = 48;

pub(crate) const HELLO_FLAG_SHM: u16   = 0x0001;
pub(crate) const SOURCE_AUDIO_ONLY: u8 = 0x01;

pub(crate) const FMT_NV12: u8      = 1;
pub(crate) const COMPRESS_NONE: u8 = 0;
pub(crate) const COMPRESS_LZ4: u8  = 1;

// Codec identifiers carried in slm_caps_t.codec (was _pad[0]).
pub const CODEC_RAW:  u8 = 0; // NV12 + optional LZ4 (current default)
pub const CODEC_H264: u8 = 1; // H.264 Annex B NAL units
pub const CODEC_H265: u8 = 2; // H.265 Annex B NAL units

// Frame type for PKT_VIDEO when using an encoded codec.
// For CODEC_RAW the value is always FRAME_TYPE_P.
// For CODEC_H264/H265: carried in the compressed_sz field (low byte of u32).
pub const FRAME_TYPE_P:   u8 = 0; // P-frame (reference frame)
pub const FRAME_TYPE_IDR: u8 = 1; // IDR - decoder can join from this frame

pub(crate) const SAMPLE_RATE: u32        = 48_000;
pub(crate) const CHANNELS: u8           = 2;
// 172 x 2ch x 4B = 1376B PCM + headers < 1500B MTU (no fragmentation).
pub(crate) const MAX_AUDIO_FRAMES: usize = 172;

fn w16(buf: &mut Vec<u8>, v: u16) { buf.extend_from_slice(&v.to_le_bytes()); }
fn w32(buf: &mut Vec<u8>, v: u32) { buf.extend_from_slice(&v.to_le_bytes()); }
fn w64(buf: &mut Vec<u8>, v: u64) { buf.extend_from_slice(&v.to_le_bytes()); }

fn header(pkt_type: u8, flags: u16, stream_id: u32) -> Vec<u8> {
    let mut buf = Vec::with_capacity(12);
    buf.extend_from_slice(&MAGIC);
    buf.push(pkt_type);
    buf.push(VERSION);
    w16(&mut buf, flags);
    w32(&mut buf, stream_id);
    buf
}

pub(crate) fn build_announce(stream_id: u32, name: &str, ctrl_port: u16, audio_only: bool) -> Vec<u8> {
    let mut buf = header(PKT_ANNOUNCE, 0, stream_id);
    w16(&mut buf, ctrl_port);
    let mut name_bytes = [0u8; 64];
    let n = name.len().min(63);
    name_bytes[..n].copy_from_slice(&name.as_bytes()[..n]);
    buf.extend_from_slice(&name_bytes);
    buf.push(if audio_only { SOURCE_AUDIO_ONLY } else { 0 });
    buf
}

pub(crate) fn build_caps(stream_id: u32, width: u16, height: u16, fps_num: u32, fps_den: u32, codec: u8) -> Vec<u8> {
    let mut buf = header(PKT_CAPS, 0, stream_id);
    w16(&mut buf, width);
    w16(&mut buf, height);
    w32(&mut buf, fps_num);
    w32(&mut buf, fps_den);
    buf.push(FMT_NV12);
    buf.push(COMPRESS_LZ4); // offered compression
    w16(&mut buf, SAMPLE_RATE as u16);
    buf.push(CHANNELS);
    buf.push(codec);          // slm_codec_t; was _pad[0]
    buf.extend_from_slice(&[0u8; 2]); // _pad[1..2]
    buf
}

/// Serialize a video frame into one or more UDP datagrams.
/// When compress=true, tries LZ4; falls back to raw if compression doesn't help.
///
/// `frame_type` is FRAME_TYPE_P for all CODEC_RAW frames (value = 0, no wire change).
/// For CODEC_H264/H265, `compress` must be false and `frame_type` is written into
/// `compressed_sz` so the receiver can distinguish IDR frames from P-frames without
/// fully reassembling the NAL stream.
pub(crate) fn build_video_chunks(
    stream_id:    u32,
    frame_seq:    u32,
    timestamp_ns: u64,
    payload:      &[u8],
    compress:     bool,
    frame_type:   u8,
) -> Vec<Vec<u8>> {
    let (send_payload, compressed_sz, compression) = if compress {
        // strip the 4-byte size prefix compress_prepend_size adds: the C consumer's
        // LZ4_decompress_safe expects a bare block, and compressed_sz must be its
        // length (see 6dad052)
        let with_prefix = compress_prepend_size(payload);
        let block = &with_prefix[4..];
        if block.len() < payload.len() {
            (block.to_vec(), block.len() as u32, COMPRESS_LZ4)
        } else {
            // LZ4 didn't help; send raw.  frame_type is 0 for CODEC_RAW so
            // compressed_sz stays 0, preserving v1 wire semantics.
            (payload.to_vec(), frame_type as u32, COMPRESS_NONE)
        }
    } else {
        // No LZ4.  For CODEC_RAW: frame_type == 0, so compressed_sz == 0 (unchanged).
        // For CODEC_H264/H265: frame_type carries IDR vs P-frame distinction.
        (payload.to_vec(), frame_type as u32, COMPRESS_NONE)
    };

    let chunk_total     = send_payload.len().div_ceil(MAX_PAYLOAD) as u16;
    let uncompressed_sz = payload.len() as u32;

    let ts_send = {
        use std::time::{SystemTime, UNIX_EPOCH};
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map(|d| d.as_nanos() as u64)
            .unwrap_or(0)
    };

    send_payload
        .chunks(MAX_PAYLOAD)
        .enumerate()
        .map(|(i, chunk)| {
            let mut buf = header(PKT_VIDEO, 0, stream_id);
            w64(&mut buf, timestamp_ns);
            w32(&mut buf, frame_seq);
            w32(&mut buf, uncompressed_sz);
            w32(&mut buf, compressed_sz);
            w16(&mut buf, i as u16);
            w16(&mut buf, chunk_total);
            buf.push(compression);
            buf.extend_from_slice(&[0u8; 3]); // padding
            w64(&mut buf, ts_send);            // sender_ts_ns: UNIX epoch ns at send time
            buf.extend_from_slice(chunk);
            buf
        })
        .collect()
}

fn build_audio_packet(stream_id: u32, frame_seq: u32, timestamp_ns: u64, samples: &[f32]) -> Vec<u8> {
    let num_frames = (samples.len() / CHANNELS as usize) as u16;
    let ts_send = {
        use std::time::{SystemTime, UNIX_EPOCH};
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map(|d| d.as_nanos() as u64)
            .unwrap_or(0)
    };
    let mut buf = header(PKT_AUDIO, 0, stream_id);
    w64(&mut buf, timestamp_ns);
    w32(&mut buf, frame_seq);
    w16(&mut buf, num_frames);
    buf.push(CHANNELS);
    buf.push(COMPRESS_NONE);
    w64(&mut buf, ts_send);            // sender_ts_ns: UNIX epoch ns at send time (v3+)
    for s in samples { buf.extend_from_slice(&s.to_le_bytes()); }
    buf
}

/// Split samples into MAX_AUDIO_FRAMES-sized chunks, one UDP packet each.
pub(crate) fn build_audio_packets(
    stream_id:    u32,
    frame_seq:    u32,
    timestamp_ns: u64,
    samples:      &[f32],
) -> Vec<Vec<u8>> {
    let chunk_size = MAX_AUDIO_FRAMES * CHANNELS as usize;
    let step_ns    = 1_000_000_000u64 * MAX_AUDIO_FRAMES as u64 / SAMPLE_RATE as u64;
    samples
        .chunks(chunk_size)
        .enumerate()
        .map(|(i, chunk)| {
            build_audio_packet(stream_id, frame_seq, timestamp_ns + i as u64 * step_ns, chunk)
        })
        .collect()
}

/// Compute XOR-parity FEC packets for `chunks` from `build_video_chunks`.
///
/// One PKT_FEC packet is emitted per group of `fec_n` data chunks.  The
/// payload is the XOR of those chunks' payloads, zero-padded to MAX_PAYLOAD.
/// A receiver that loses exactly one chunk per group can reconstruct it.
///
/// Returns empty when `fec_n < 2` (FEC disabled) or `chunks` is empty.
pub(crate) fn append_fec_packets(
    chunks:       &[Vec<u8>],
    fec_n:        u8,
    stream_id:    u32,
    frame_seq:    u32,
    timestamp_ns: u64,
) -> Vec<Vec<u8>> {
    if fec_n < 2 || chunks.is_empty() { return vec![]; }
    let chunk_total = chunks.len() as u16;
    chunks
        .chunks(fec_n as usize)
        .enumerate()
        .map(|(g, group)| {
            let mut parity = [0u8; MAX_PAYLOAD];
            for chunk_pkt in group {
                let payload = chunk_pkt.get(VIDEO_HDR_SIZE..).unwrap_or(&[]);
                for (i, &b) in payload.iter().enumerate() {
                    parity[i] ^= b;
                }
            }
            let mut buf = header(PKT_FEC, 0, stream_id);
            w64(&mut buf, timestamp_ns);
            w32(&mut buf, frame_seq);
            w16(&mut buf, chunk_total);
            buf.push(fec_n);
            buf.push(0u8);             // _pad1
            w16(&mut buf, g as u16);   // fec_group
            buf.extend_from_slice(&[0u8; 2]); // _pad2
            buf.extend_from_slice(&parity);
            buf
        })
        .collect()
}

pub(crate) fn build_bye(stream_id: u32) -> Vec<u8> {
    header(PKT_BYE, 0, stream_id)
}

/// Sender -> receiver address-validation challenge (v4). The receiver must echo
/// `nonce` in the body of its HELLO (the 8 bytes after the 12-byte header)
/// before the sender will stream media to it. This proves the receiver actually
/// receives at the address it claimed, defeating UDP source-address spoofing
/// (a spoofer never sees the challenge, so can never echo the right nonce).
pub(crate) fn build_challenge(stream_id: u32, nonce: u64) -> Vec<u8> {
    let mut buf = header(PKT_CHALLENGE, 0, stream_id);
    w64(&mut buf, nonce);
    buf
}

/// Receiver -> sender tally update. Payload is a single byte: `TALLY_OFF`,
/// `TALLY_PREVIEW`, or `TALLY_PROGRAM`.
pub fn build_tally(stream_id: u32, state: u8) -> Vec<u8> {
    let mut buf = header(PKT_TALLY, 0, stream_id);
    buf.push(state);
    buf
}

/// Receiver -> sender operator message. Payload is a 1-byte length followed by
/// up to `MSG_MAX_BYTES` of UTF-8 text. The text is truncated on a byte
/// boundary, which may split a multi-byte char; the receiver decodes lossily.
pub fn build_msg(stream_id: u32, text: &str) -> Vec<u8> {
    let mut buf = header(PKT_MSG, 0, stream_id);
    let n = text.len().min(MSG_MAX_BYTES);
    buf.push(n as u8);
    buf.extend_from_slice(&text.as_bytes()[..n]);
    buf
}

/// Parse a receiver -> sender PKT_STAT. The 8-byte payload is the `sender_ts_ns`
/// the sender stamped on the most recently consumed frame, echoed back verbatim.
/// Returns it so the sender can compute elapsed time on its own clock.
pub(crate) fn parse_stat(buf: &[u8]) -> Option<u64> {
    // 12-byte SLM header + u64 echoed timestamp.
    if buf.len() < 20 { return None; }
    Some(u64::from_le_bytes(buf[12..20].try_into().ok()?))
}

/// Parse a receiver -> sender PKT_MSG: 12-byte header, a 1-byte length, then that
/// many UTF-8 bytes (see `build_msg`). Decoded lossily so a truncated multi-byte
/// char never rejects the whole message. Returns the operator text.
pub(crate) fn parse_msg(buf: &[u8]) -> Option<String> {
    if buf.len() < 13 { return None; }
    let n = buf[12] as usize;
    let end = 13 + n;
    if buf.len() < end { return None; }
    Some(String::from_utf8_lossy(&buf[13..end]).into_owned())
}
