//! 通用协议与帧封装：控制消息（bincode + 长度前缀）与视频帧格式。

/// P2P 打洞信令协议（信令通道专用，不含中继；三端共享契约）
pub mod signaling;

use anyhow::Result;
use serde::{Deserialize, Serialize};
use tokio::io::{AsyncReadExt, AsyncWriteExt};

/// 媒体/控制协议版本。
///
/// **改动规则**：任何对 `ControlMsg` 或 `VideoFrameHeader` 字段的增删改都必须递增此值。
/// 理由：线格式是 bincode（按位置编码，无字段名），新增字段会让旧端解析失败或错位，
/// 且失败方式可能是静默的错误数据而非明确报错。递增后由 `Auth` 握手显式拒绝旧端，
/// 把"连上了但行为诡异"变成"连不上且提示升级"。
///
/// - v1：初版（无版本字段，隐式）
/// - v2：`Auth`/`AuthOk` 携带版本；帧头增加 `seq`/`keyframe`
/// - v3：新增 `InputEvent::Text`（文本注入通道，UTF-16 code unit）。
///   `InputEvent` 因携带 `Vec<u16>` 不再是 `Copy`；三端必须同步升级，
///   旧端由 `Auth` 握手显式拒绝。
pub const PROTO_VERSION: u16 = 3;

/// 控制通道消息（双向）
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ControlMsg {
    /// 主控端认证（临时验证码，MVP 明文走 QUIC-TLS 加密通道）。
    /// `proto_version` 必须等于 [`PROTO_VERSION`]，否则被控端回 `AuthFail` 并拒绝。
    Auth { code: String, proto_version: u16 },
    /// 认证通过。`proto_version` 回显被控端的版本，供主控端确认双方一致。
    AuthOk {
        width: u32,
        height: u32,
        proto_version: u16,
    },
    AuthFail(String),
    KeyFrameReq,
    Ping(u64),
    /// t=主控端发送时刻（主控时钟），agent_ts=被控端收到时刻（被控时钟）
    /// 两者配合可估计时钟偏移，进而用帧头 ts_ms 计算端到端延迟
    /// lost/sent = agent 发送方向（视频下行）的 QUIC 丢包/发包累计计数（conn.stats().path）
    Pong { t: u64, agent_ts: u64, lost: u64, sent: u64 },
    Input(InputEvent),
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum MouseButton {
    Left,
    Right,
    Middle,
}

/// 输入事件。鼠标坐标为**被采集显示器**（主屏）归一化 0..=65535
/// （协议 v2 语义，v3 未改；`agent/src/input.rs` 的注入空间与此一致，见个人层 lessons R-4）；
/// 按键为 Set-1 扫描码，extended 标记扩展键（方向键/Home 等，需要 EXTENDEDKEY 标志）。
///
/// v3 起本类型不再 `Copy`：`Text` 携带堆上 `Vec<u16>`。需要复制的地方用 `.clone()`。
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum InputEvent {
    MouseMove { x: u16, y: u16 },
    MouseDown(MouseButton),
    MouseUp(MouseButton),
    Wheel { dy: i32 },
    KeyDown { scancode: u16, extended: bool },
    KeyUp { scancode: u16, extended: bool },
    /// v3 新增：**文本注入通道**。`units` 为 UTF-16 code unit 序列。
    ///
    /// 为什么不是 `String`：Windows 注入侧（`KEYEVENTF_UNICODE`）与 Java 字符串的
    /// 原生表示都是 UTF-16 code unit；用 `String` 要走 UTF-8 往返，遇到 IME 提交边界
    /// 产生的落单代理项会**有损**。`Vec<u16>` 忠实透传，由接收端按 Win32 语义处理。
    ///
    /// 为什么整段携带、而不是逐 code unit 一条消息：一次 IME 候选确认常是多个汉字；
    /// 代理对（emoji 等增补平面字符）**必须在同一次注入批次里成对提交**才能被目标
    /// 应用合成成一个字符，拆开注入必现乱码。
    ///
    /// 上限 [`MAX_TEXT_UNITS`]：超长输入截断且两侧打警告日志（不静默丢弃）。
    /// 另一道防线是 [`read_control`] 的 1 MiB 消息长度闸。
    ///
    /// 日志纪律：中间层与注入端只允许打印 `units.len()`，**不得**打印用户输入内容。
    Text { units: Vec<u16> },
}

/// [`InputEvent::Text`] 单次携带的 code unit 上限。
///
/// 与 `read_control` 的 1 MiB 消息闸是两道独立防线：前者防恶意对端把接收端分配打爆，
/// 后者防合法但异常长的 IME 提交把注入队列刷爆。4096 unit ≈ 8 KiB，
/// 远超正常 IME 一次提交（数十 unit），又足够小使单次 `SendInput` 批量仍可接受。
pub const MAX_TEXT_UNITS: usize = 4096;

/// 把文本输入截断到 [`MAX_TEXT_UNITS`]，返回 `(可注入的 units, 是否发生截断)`。
///
/// 调用方**必须**在截断为真时打印日志（可见后果）；本函数只负责机械截断，
/// 不自己打印——因为 `common` 同时被主控端与被控端复用，日志口径由调用方决定。
pub fn clamp_text_units(units: Vec<u16>) -> (Vec<u16>, bool) {
    if units.len() > MAX_TEXT_UNITS {
        let mut trimmed = units;
        trimmed.truncate(MAX_TEXT_UNITS);
        (trimmed, true)
    } else {
        (units, false)
    }
}

/// 视频帧头。载荷为 H.264 Annex-B 码流（含 00 00 00 01 起始码）。
///
/// 字段语义：
/// - `width` / `height`：被控端桌面尺寸（逻辑尺寸，非编码对齐后的尺寸）
/// - `ts_ms`：被控端采集时刻（被控时钟），配合 Ping/Pong 估计的时钟偏移算端到端延迟
/// - `seq`：帧序号，从 0 单调递增，每发一帧 +1（含静止桌面的重复帧）。
///   用于丢帧检测与弱网统计；注意它统计的是"发送序"，不是"显示序"
/// - `keyframe`：本帧是否 IDR。解码端（尤其 MediaCodec）需要它决定是否重置解码器，
///   不必再自己扫 NAL 起始码判断
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VideoFrameHeader {
    pub width: u32,
    pub height: u32,
    pub ts_ms: u64,
    pub seq: u64,
    pub keyframe: bool,
}

async fn read_u32le<R: AsyncReadExt + Unpin>(r: &mut R) -> Result<u32> {
    let mut b = [0u8; 4];
    r.read_exact(&mut b).await?;
    Ok(u32::from_le_bytes(b))
}

/// 编码一条控制消息为「长度前缀 + bincode」
pub fn encode_control(msg: &ControlMsg) -> Vec<u8> {
    let body = bincode::serialize(msg).expect("serialize control");
    let mut out = Vec::with_capacity(body.len() + 4);
    out.extend_from_slice(&(body.len() as u32).to_le_bytes());
    out.extend_from_slice(&body);
    out
}

pub async fn write_control<W: AsyncWriteExt + Unpin>(
    w: &mut W,
    msg: &ControlMsg,
) -> Result<()> {
    w.write_all(&encode_control(msg)).await?;
    Ok(())
}

pub async fn read_control<R: AsyncReadExt + Unpin>(r: &mut R) -> Result<ControlMsg> {
    let len = read_u32le(r).await? as usize;
    anyhow::ensure!(len < 1 << 20, "control frame too large: {len}");
    let mut buf = vec![0u8; len];
    r.read_exact(&mut buf).await?;
    Ok(bincode::deserialize(&buf)?)
}

/// 编码视频帧：[总长 u32][头长 u32][bincode 头][H.264 载荷]
///
/// `总长` 不含自身这 4 字节，等于 `4 + 头长 + 载荷长`。
pub fn encode_video_frame(header: &VideoFrameHeader, payload: &[u8]) -> Vec<u8> {
    let hb = bincode::serialize(header).expect("serialize header");
    let total = 4 + hb.len() + payload.len();
    let mut out = Vec::with_capacity(total + 4);
    out.extend_from_slice(&(total as u32).to_le_bytes());
    out.extend_from_slice(&(hb.len() as u32).to_le_bytes());
    out.extend_from_slice(&hb);
    out.extend_from_slice(payload);
    out
}

/// 单帧总长上限（256 MiB）。超过即视为协议错误或恶意输入。
const MAX_VIDEO_FRAME: usize = 256 << 20;

/// 判断一段 Annex-B H.264 码流里是否含 IDR 片（NAL 类型 5）。
///
/// 用途：填充 [`VideoFrameHeader::keyframe`]。
/// 为什么不问编码器要这个标志——NVENC 与 openh264 的 API 都能间接推出，
/// 但两者规则不同（首帧强制、`force_intra_frame`、GOP 自动 IDR 交织），
/// 而**实际输出**是唯一可信来源。直接看码流，两种编码器一套逻辑。
///
/// 只识别 3/4 字节起始码；`start-code emulation prevention` 由编码器保证，
/// 因此不会在 SPS/PPS 载荷里误判出起始码。
pub fn h264_contains_idr(payload: &[u8]) -> bool {
    let mut i = 0usize;
    while i + 3 < payload.len() {
        let sc = if payload[i..].starts_with(&[0, 0, 0, 1]) {
            4
        } else if payload[i..].starts_with(&[0, 0, 1]) {
            3
        } else {
            i += 1;
            continue;
        };
        match payload.get(i + sc) {
            // nal_unit_type 占首字节低 5 位
            Some(&b) if b & 0x1F == 5 => return true,
            Some(_) => i += sc + 1,
            None => break,
        }
    }
    false
}

/// 读一个视频帧。
///
/// 对端数据一律视为不可信输入：`total` 与 `hlen` 都必须做**上下界**校验，
/// 否则构造一个 `total < 4` 的帧就能让接收端在切片处 panic（进程级拒绝服务）。
pub async fn read_video_frame<R: AsyncReadExt + Unpin>(
    r: &mut R,
) -> Result<(VideoFrameHeader, Vec<u8>)> {
    let total = read_u32le(r).await? as usize;
    // 下界：至少要容得下「头长字段(4) + 一个非空的 bincode 头」。
    // 上界：防止对端声明超大长度导致巨量分配。
    anyhow::ensure!(
        (4..MAX_VIDEO_FRAME).contains(&total),
        "video frame length out of range: {total}"
    );

    let mut buf = vec![0u8; total];
    r.read_exact(&mut buf).await?;

    let hlen = u32::from_le_bytes(buf[0..4].try_into().expect("已校验 total >= 4")) as usize;
    // hlen 必须落在 [1, total-4] 内：0 无法反序列化，超出则切片越界
    anyhow::ensure!(
        hlen > 0 && hlen <= total - 4,
        "video header length invalid: hlen={hlen}, total={total}"
    );

    let header: VideoFrameHeader = bincode::deserialize(&buf[4..4 + hlen])?;
    Ok((header, buf[4 + hlen..].to_vec()))
}

/// 把一个已编码的完整帧拆回 `(帧头, 载荷)`。
///
/// 用途：静止桌面重发时，被控端需要复用上一帧的**载荷**但重建**帧头**
/// （`seq` 必须递增、`ts_ms` 必须是本次发送时刻）。直接 `clone` 整个帧会让
/// 这两个字段停在上一次发送的时刻，破坏 `seq` 单调递增契约。
///
/// 只接受本模块 [`encode_video_frame`] 产出的字节；格式不符返回 `Err`。
pub fn split_video_frame(frame: &[u8]) -> Result<(VideoFrameHeader, &[u8])> {
    anyhow::ensure!(frame.len() >= 8, "frame too short: {}", frame.len());
    let hlen = u32::from_le_bytes(frame[4..8].try_into().expect("已校验 len >= 8")) as usize;
    anyhow::ensure!(
        hlen > 0 && hlen <= frame.len() - 8,
        "header length invalid: hlen={hlen}, len={}",
        frame.len()
    );
    let header: VideoFrameHeader = bincode::deserialize(&frame[8..8 + hlen])?;
    Ok((header, &frame[8 + hlen..]))
}

#[cfg(test)]
mod tests {
    use super::*;
    use tokio::io::duplex;

    fn hdr(seq: u64, keyframe: bool) -> VideoFrameHeader {
        VideoFrameHeader {
            width: 1920,
            height: 1080,
            ts_ms: 12345,
            seq,
            keyframe,
        }
    }

    // ---- 控制消息 ----

    #[test]
    fn control_roundtrip_preserves_version() {
        let m = ControlMsg::Auth {
            code: "123456".into(),
            proto_version: PROTO_VERSION,
        };
        let bytes = encode_control(&m);
        // 长度前缀 + 载荷
        let len = u32::from_le_bytes(bytes[0..4].try_into().unwrap()) as usize;
        assert_eq!(len, bytes.len() - 4);
        let back: ControlMsg = bincode::deserialize(&bytes[4..]).unwrap();
        match back {
            ControlMsg::Auth { code, proto_version } => {
                assert_eq!(code, "123456");
                assert_eq!(proto_version, PROTO_VERSION);
            }
            other => panic!("期望 Auth，实际 {other:?}"),
        }
    }

    #[tokio::test]
    async fn control_write_read_roundtrip() {
        let (mut a, mut b) = duplex(64 * 1024);
        let msgs = [
            ControlMsg::Auth {
                code: "x".into(),
                proto_version: PROTO_VERSION,
            },
            ControlMsg::AuthOk {
                width: 800,
                height: 600,
                proto_version: PROTO_VERSION,
            },
            ControlMsg::KeyFrameReq,
            ControlMsg::Ping(42),
            ControlMsg::Pong {
                t: 1,
                agent_ts: 2,
                lost: 3,
                sent: 4,
            },
        ];
        for m in &msgs {
            write_control(&mut a, m).await.unwrap();
            let got = read_control(&mut b).await.unwrap();
            assert_eq!(format!("{got:?}"), format!("{m:?}"));
        }
    }

    // ---- v3 文本注入通道 ----

    /// 版本守卫：v3 是本轮契约变更的登记点，改动必须显式经过这里。
    #[test]
    fn proto_version_is_three_after_text_event() {
        assert_eq!(PROTO_VERSION, 3, "新增 InputEvent::Text 后必须递增到 3");
    }

    #[test]
    fn text_event_roundtrips_utf16_units() {
        // 「你好」+ 一个 emoji 的代理对（U+1F600 = D83D DE00）
        let units: Vec<u16> = vec![0x4F60, 0x597D, 0xD83D, 0xDE00];
        let m = ControlMsg::Input(InputEvent::Text {
            units: units.clone(),
        });
        let bytes = encode_control(&m);
        let back: ControlMsg = bincode::deserialize(&bytes[4..]).unwrap();
        match back {
            ControlMsg::Input(InputEvent::Text { units: got }) => assert_eq!(got, units),
            other => panic!("期望 Input::Text，实际 {other:?}"),
        }
    }

    /// 落单代理项必须原样往返：IME 提交边界可能给出孤立高/低代理项，
    /// 用 `String` 承载会在这里有损（这正是选 `Vec<u16>` 的理由）。
    #[test]
    fn text_event_keeps_lone_surrogate() {
        let units = vec![0xD83Du16];
        let m = ControlMsg::Input(InputEvent::Text { units: units.clone() });
        let bytes = encode_control(&m);
        let back: ControlMsg = bincode::deserialize(&bytes[4..]).unwrap();
        match back {
            ControlMsg::Input(InputEvent::Text { units: got }) => assert_eq!(got, units),
            other => panic!("期望 Input::Text，实际 {other:?}"),
        }
    }

    #[test]
    fn text_units_clamped_at_boundary() {
        // 边界：恰好上限不截断；超一个才截断
        let exact = vec![0x61u16; MAX_TEXT_UNITS];
        let (kept, truncated) = clamp_text_units(exact);
        assert_eq!(kept.len(), MAX_TEXT_UNITS);
        assert!(!truncated);

        let over = vec![0x61u16; MAX_TEXT_UNITS + 1];
        let (kept, truncated) = clamp_text_units(over);
        assert_eq!(kept.len(), MAX_TEXT_UNITS);
        assert!(truncated, "超限必须返回截断标志，让调用方打印警告");

        let (empty, truncated) = clamp_text_units(vec![]);
        assert!(empty.is_empty());
        assert!(!truncated);
    }

    /// 空 `Text` 也必须能往返（IME 可能提交空串；调用方负责决定忽略还是报错）。
    #[test]
    fn text_event_empty_roundtrips() {
        let m = ControlMsg::Input(InputEvent::Text { units: vec![] });
        let bytes = encode_control(&m);
        let back: ControlMsg = bincode::deserialize(&bytes[4..]).unwrap();
        match back {
            ControlMsg::Input(InputEvent::Text { units }) => assert!(units.is_empty()),
            other => panic!("期望 Input::Text，实际 {other:?}"),
        }
    }

    /// 上限内的 `Text` 必须能通过 `read_control` 的长度闸；
    /// 反向验证：伪造一个超限的声明长度必须在分配前被拒。
    #[tokio::test]
    async fn text_within_cap_passes_control_length_gate() {
        let (mut a, mut b) = duplex(1 << 20);
        let msg = ControlMsg::Input(InputEvent::Text {
            units: vec![0x4F60; MAX_TEXT_UNITS],
        });
        write_control(&mut a, &msg).await.unwrap();
        match read_control(&mut b).await.unwrap() {
            ControlMsg::Input(InputEvent::Text { units }) => assert_eq!(units.len(), MAX_TEXT_UNITS),
            other => panic!("期望 Input::Text，实际 {other:?}"),
        }
    }

    // ---- 视频帧头 ----

    #[test]
    fn video_header_roundtrip_keeps_seq_and_keyframe() {
        for (seq, kf) in [(0u64, true), (1, false), (u64::MAX, false)] {
            let h = hdr(seq, kf);
            let payload = b"\x00\x00\x00\x01\x65fake-nal";
            let bytes = encode_video_frame(&h, payload);

            let total = u32::from_le_bytes(bytes[0..4].try_into().unwrap()) as usize;
            assert_eq!(total, bytes.len() - 4, "总长字段应不含自身");

            let hlen = u32::from_le_bytes(bytes[4..8].try_into().unwrap()) as usize;
            let back: VideoFrameHeader = bincode::deserialize(&bytes[8..8 + hlen]).unwrap();
            assert_eq!(back.seq, seq);
            assert_eq!(back.keyframe, kf);
            assert_eq!(&bytes[8 + hlen..], payload);
        }
    }

    #[tokio::test]
    async fn video_frame_write_read_roundtrip() {
        let (mut a, mut b) = duplex(256 * 1024);
        let h = hdr(7, true);
        let payload = vec![0xABu8; 4096];
        let bytes = encode_video_frame(&h, &payload);
        tokio::io::AsyncWriteExt::write_all(&mut a, &bytes).await.unwrap();

        let (h2, p2) = read_video_frame(&mut b).await.unwrap();
        assert_eq!(h2.seq, 7);
        assert!(h2.keyframe);
        assert_eq!((h2.width, h2.height), (1920, 1080));
        assert_eq!(p2, payload);
    }

    // ---- 不可信输入（P3 的核心回归） ----

    /// `total < 4` 必须报错而不是 panic：旧实现在这里 `buf[0..4]` 会越界 panic。
    #[tokio::test]
    async fn rejects_total_below_minimum() {
        for total in [0u32, 1, 3] {
            let (mut a, mut b) = duplex(64);
            tokio::io::AsyncWriteExt::write_all(&mut a, &total.to_le_bytes())
                .await
                .unwrap();
            let err = read_video_frame(&mut b).await.unwrap_err().to_string();
            assert!(
                err.contains("out of range"),
                "total={total} 应被拒绝，实际错误: {err}"
            );
        }
    }

    /// `total == 4` 但没有头数据：hlen 读出来是 0，必须拒绝。
    #[tokio::test]
    async fn rejects_zero_header_length() {
        let (mut a, mut b) = duplex(64);
        let mut bytes = 4u32.to_le_bytes().to_vec();
        bytes.extend_from_slice(&0u32.to_le_bytes());
        tokio::io::AsyncWriteExt::write_all(&mut a, &bytes)
            .await
            .unwrap();
        let err = read_video_frame(&mut b).await.unwrap_err().to_string();
        assert!(err.contains("header length invalid"), "实际错误: {err}");
    }

    /// hlen 声明超过 total-4：旧实现会切片越界 panic。
    #[tokio::test]
    async fn rejects_header_length_beyond_frame() {
        let (mut a, mut b) = duplex(64);
        // total = 8 → body 必须正好 8 字节（hlen 字段 4 + 4 字节填充）
        let mut bytes = 8u32.to_le_bytes().to_vec();
        bytes.extend_from_slice(&999u32.to_le_bytes()); // hlen = 999 > total-4 = 4
        bytes.extend_from_slice(&[0u8; 4]);
        tokio::io::AsyncWriteExt::write_all(&mut a, &bytes)
            .await
            .unwrap();
        let err = read_video_frame(&mut b).await.unwrap_err().to_string();
        assert!(err.contains("header length invalid"), "实际错误: {err}");
    }

    /// 超过 256 MiB 上限必须在**分配前**拒绝，不能先 vec! 再检查。
    #[tokio::test]
    async fn rejects_oversized_frame_without_allocating() {
        let (mut a, mut b) = duplex(64);
        let total = (MAX_VIDEO_FRAME as u32) + 1;
        tokio::io::AsyncWriteExt::write_all(&mut a, &total.to_le_bytes())
            .await
            .unwrap();
        let err = read_video_frame(&mut b).await.unwrap_err().to_string();
        assert!(err.contains("out of range"), "实际错误: {err}");
    }

    /// 头部字节损坏（非合法 bincode）应返回 Err 而非 panic。
    #[tokio::test]
    async fn rejects_corrupt_header_bytes() {
        let (mut a, mut b) = duplex(64);
        let garbage = [0xFFu8; 8];
        let total = (4 + garbage.len()) as u32;
        let mut bytes = total.to_le_bytes().to_vec();
        bytes.extend_from_slice(&(garbage.len() as u32).to_le_bytes());
        bytes.extend_from_slice(&garbage);
        tokio::io::AsyncWriteExt::write_all(&mut a, &bytes)
            .await
            .unwrap();
        assert!(read_video_frame(&mut b).await.is_err());
    }

    // ---- IDR 识别（P2 的 keyframe 标志来源） ----

    #[test]
    fn detects_idr_in_annex_b() {
        // SPS(7) + PPS(8) + IDR 片(5)，4 字节起始码
        let idr = [
            0, 0, 0, 1, 0x67, 0xAA, // SPS
            0, 0, 0, 1, 0x68, 0xBB, // PPS
            0, 0, 0, 1, 0x65, 0xCC, // IDR
        ];
        assert!(h264_contains_idr(&idr));

        // 3 字节起始码也应识别
        let idr3 = [0, 0, 1, 0x65, 0xCC];
        assert!(h264_contains_idr(&idr3));
    }

    #[test]
    fn rejects_non_idr_payloads() {
        // 纯 P 片（NAL 类型 1）
        let p_frame = [0, 0, 0, 1, 0x41, 0xDD];
        assert!(!h264_contains_idr(&p_frame));

        // 只有 SPS/PPS，没有片（关键帧边界可能落在下一个包，此时不应谎报 keyframe）
        let params_only = [
            0, 0, 0, 1, 0x67, 0xAA,
            0, 0, 0, 1, 0x68, 0xBB,
        ];
        assert!(!h264_contains_idr(&params_only));

        assert!(!h264_contains_idr(&[]));
        assert!(!h264_contains_idr(&[0, 0, 1])); // 起始码后无 NAL 字节
    }
}