//! UDP 接收、发现、多路帧重组。
//!
//! # 多路
//!
//! 裁判终端同时收两块图传板，队伍终端也会收到两块（只是界面上只显示一块）。
//! 所以重组状态**必须按发送方分开**——把两块板的帧塞进同一个重组器，会让它们
//! 的 `frame_id` 互相顶掉，两路都出不来。
//!
//! # 发现与订阅
//!
//! 两个方向各司其职：
//!
//! - 板子每 2 秒广播一次 [`protocol::MsgType::Hello`]，客户端被动收到就知道有谁在线。
//! - 客户端广播 [`protocol::MsgType::Subscribe`]，板子收到后开始定向推流。
//!
//! 板子靠**源地址**记住客户端，所以客户端只要周期性续订就能维持推流，不存在
//! 任何超时淘汰的竞态。
//!
//! # 丢帧策略
//!
//! 缺了分片的帧整个丢弃、下一帧顶上，不重传也不尝试补洞。对实时画面这是对的：
//! 等重传补齐时，那一帧早就不新鲜了。这个取舍的代价是丢帧率必须摆在界面上，否则
//! "画面卡"没法区分是链路问题还是传感器问题。

use std::collections::HashMap;
use std::io::ErrorKind;
use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4, UdpSocket};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::{Duration, Instant};

use protocol::{Fragment, FrameAssembler, Subscribe, SubscribeRole, DEFAULT_PORT, HEADER_LEN};

/// 接受的最大数据报。高于协议的理论上限（1427），留出余量，这样一个被截断的读只可能
/// 来自网络，而不是来自过小的缓冲。
const RECV_BUFFER: usize = 2048;

/// 重新广播订阅的间隔。
///
/// 比板子判定"客户端掉线"的超时（6 秒）短得多，所以偶发丢一两个包不会中断推流。
const SUBSCRIBE_INTERVAL: Duration = Duration::from_millis(800);

/// 一张帧没收齐就放弃的时限。
///
/// 一帧在 22fps 下约 45ms 到，所以几百毫秒还没齐的永远不会齐了。
const FRAME_TIMEOUT: Duration = Duration::from_millis(500);

/// 一块已知的图传板。
#[derive(Debug, Clone)]
pub struct CameraInfo {
    pub id: u32,
    pub addr: SocketAddrV4,
    pub firmware: u16,
    /// 最近一次收到它广播的时刻。
    pub last_seen: Instant,
}

impl CameraInfo {
    /// 多久没有广播了。
    pub fn silence(&self) -> Duration {
        self.last_seen.elapsed()
    }

    /// 板子在线吗？
    ///
    /// 阈值是广播间隔的数倍，所以偶发丢包不会让界面上的在线状态闪��。
    pub fn is_online(&self) -> bool {
        self.silence() < Duration::from_secs(7)
    }
}

/// 一块板的运行时统计。
#[derive(Debug, Clone, Copy, Default)]
pub struct SourceStats {
    pub frames_shown: u64,
    /// 因分片缺失被放弃的帧。
    pub frames_dropped: u64,
    pub datagrams: u64,
    pub last_frame_len: usize,
    pub width: u16,
    pub height: u16,
    /// 上一帧到这一帧之间的间隔。
    pub last_interval: Option<Duration>,
    last_shown_at: Option<Instant>,
}

impl SourceStats {
    /// 丢弃的帧占已收到帧的比例。
    pub fn loss_percent(&self) -> f32 {
        let total = self.frames_shown + self.frames_dropped;
        if total == 0 {
            0.0
        } else {
            self.frames_dropped as f32 * 100.0 / total as f32
        }
    }

    fn on_frame(&mut self, width: u16, height: u16, len: usize) {
        self.last_interval = self.last_shown_at.map(|t| t.elapsed());
        self.last_shown_at = Some(Instant::now());
        self.frames_shown += 1;
        self.width = width;
        self.height = height;
        self.last_frame_len = len;
    }
}

/// 每一路的最新帧，等着被解码。
///
/// 每路只保留最新的一张：在 22fps 下把每一帧都排队意味着内存里压着几十张 JPEG，
/// 然后解码那些几毫秒后就被下一帧顶掉的画面——而这正是这条链路要消除的延迟。
#[derive(Default)]
pub struct LatestFrames {
    inner: Mutex<HashMap<u32, protocol::Frame>>,
}

impl LatestFrames {
    pub fn put(&self, frame: protocol::Frame) {
        self.inner
            .lock()
            .expect("latest-frames mutex poisoned")
            .insert(frame.sender_id, frame);
    }

    /// 取走某一路的待处理帧。
    pub fn take(&self, sender_id: u32) -> Option<protocol::Frame> {
        self.inner
            .lock()
            .expect("latest-frames mutex poisoned")
            .remove(&sender_id)
    }

    /// 列出所有有待处理帧的板子 id。
    pub fn pending_senders(&self) -> Vec<u32> {
        self.inner
            .lock()
            .expect("latest-frames mutex poisoned")
            .keys()
            .copied()
            .collect()
    }
}

/// 所有板子的共享状态，跨线程只读。
pub struct Discovery {
    cameras: Mutex<HashMap<u32, CameraInfo>>,
    stats: Mutex<HashMap<u32, SourceStats>>,
}

impl Discovery {
    fn new() -> Self {
        Self {
            cameras: Mutex::new(HashMap::new()),
            stats: Mutex::new(HashMap::new()),
        }
    }

    /// 当前已知的板子，按 id 排序，让界面上���列表顺序稳定。
    pub fn cameras(&self) -> Vec<CameraInfo> {
        let mut list: Vec<CameraInfo> = self
            .cameras
            .lock()
            .expect("cameras mutex poisoned")
            .values()
            .cloned()
            .collect();

        list.sort_by_key(|c| c.id);
        list
    }

    /// 在线且刚刚广播过的板子。
    pub fn online_cameras(&self) -> Vec<CameraInfo> {
        self.cameras().into_iter().filter(CameraInfo::is_online).collect()
    }

    /// 某一路的统计。
    pub fn stats(&self, sender_id: u32) -> SourceStats {
        self.stats
            .lock()
            .expect("stats mutex poisoned")
            .get(&sender_id)
            .copied()
            .unwrap_or_default()
    }

    /// 只刷新存活时间戳，不改地址或固件版本。
    ///
    /// 收到视频帧时调用：正在推流本身就是最可靠的存活证据。
    fn touch(&self, id: u32) {
        let mut cameras = self.cameras.lock().expect("cameras mutex poisoned");

        if let Some(info) = cameras.get_mut(&id) {
            info.last_seen = Instant::now();
        }
    }

    fn note_hello(&self, id: u32, addr: SocketAddrV4, firmware: u16) {
        let mut cameras = self.cameras.lock().expect("cameras mutex poisoned");

        cameras
            .entry(id)
            .and_modify(|c| {
                c.addr = addr;
                c.last_seen = Instant::now();
                c.firmware = firmware;
            })
            .or_insert(CameraInfo {
                id,
                addr,
                firmware,
                last_seen: Instant::now(),
            });
    }

    fn note_frame(&self, id: u32, width: u16, height: u16, len: usize) {
        let mut stats = self.stats.lock().expect("stats mutex poisoned");
        stats.entry(id).or_default().on_frame(width, height, len);
    }

    fn count(&self, id: u32, f: impl FnOnce(&mut SourceStats)) {
        let mut stats = self.stats.lock().expect("stats mutex poisoned");
        f(stats.entry(id).or_default());
    }

    /// 超过 `max_idle` 没有广播的板子从列表里移除。
    ///
    /// 界面上不会留下一堆永远离线的板子。
    pub fn forget_silent_cameras(&self, max_idle: Duration) {
        self.cameras
            .lock()
            .expect("cameras mutex poisoned")
            .retain(|_, c| c.silence() < max_idle);
    }
}

/// 客户端身份。
///
/// 常量而非 MAC 派生：队伍客户端和裁判终端各自唯一，进程内也不需要区分自己，
/// 而一个确定的常量让日志更好读。
pub const CLIENT_ID: u32 = 0x0C11_0001;

/// 接收器。
///
/// socket 在这里创建而不是在接收线程里，这样"端口被占用"会立刻报错，而不是被线程
/// 吞掉、表现为一个永远不出画面的空窗口。
pub struct Receiver {
    discovery: Arc<Discovery>,
    frames: Arc<LatestFrames>,
    role: SubscribeRole,
    camera: Ipv4Addr,
    /// 只处理这一块板子的数据，`None` 表示全都要（裁判端）。
    ///
    /// 队伍端在 release 下设成 `Some(bound)`，这是一道**接收端**的闸门而不只是界面
    /// 的：对手的板子即使把流推到了本机端口，也不会被重组、不会进解码队列、不会
    /// 出现在任何统计里。
    ///
    /// 为什么必须在这一层而不是靠界面隐藏：`.slint` 只控制显示，而数据已经收下并
    /// 解码了。隐藏一行文字不阻止任何事。
    ///
    /// 用 `Mutex` 而不是 `Cell`/`AtomicU32`：接收线程要读它，而绑定回调在主线程上
    /// 写它。订阅循环本身已经是串行的，所以锁的持有时间极短。
    only_sender: Mutex<Option<u32>>,
    stop: Arc<AtomicBool>,
}

impl Receiver {
    /// 绑定视频端口并准备好发现。
    ///
    /// 绑定固定端口（而非临时端口）是必须的：板子是从订阅的**源端口**学到推流目标的，
    /// 临时端口也能工作，但固定端口让三台客户端行为一致，也方便抓包排查。
    ///
    /// `only_sender` 为 `Some` 时只接受那一块板子的数据，见 [`Receiver`] 的字段说明。
    pub fn start(
        camera: Ipv4Addr,
        role: SubscribeRole,
        only_sender: Option<u32>,
    ) -> std::io::Result<Self> {
        let socket =
            UdpSocket::bind(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, DEFAULT_PORT))?;
        socket.set_read_timeout(Some(Duration::from_millis(100)))?;

        let broadcast = socket.set_broadcast(true).is_ok();
        if !broadcast {
            log::warn!("could not enable broadcast on the socket");
        }

        Ok(Self {
            discovery: Arc::new(Discovery::new()),
            frames: Arc::new(LatestFrames::default()),
            role,
            camera,
            only_sender: Mutex::new(only_sender),
            stop: Arc::new(AtomicBool::new(false)),
        })
    }

    /// 绑定之后切换限制目标，或解除限制。
    ///
    /// 队伍端在选择界面里点定一块板子时调用：此前要能看到所有候选（否则无从选
    /// 起），此后只看这一块。
    pub fn restrict_to(&self, only_sender: Option<u32>) {
        *self.only_sender.lock().expect("only_sender mutex poisoned") = only_sender;
    }

    /// 当前限制到哪一块板子。
    fn sender_limit(&self) -> Option<u32> {
        *self.only_sender.lock().expect("only_sender mutex poisoned")
    }

    pub fn discovery(&self) -> Arc<Discovery> {
        Arc::clone(&self.discovery)
    }

    pub fn frames(&self) -> Arc<LatestFrames> {
        Arc::clone(&self.frames)
    }

    /// 运行接收循环，直到 [`Receiver::shutdown`]。
    ///
    /// 阻塞，所以放在单独的线程上。
    pub fn run(&self, socket: UdpSocket) {
        let target = SocketAddrV4::new(self.camera, DEFAULT_PORT);
        let subscribe = Subscribe {
            role: self.role,
            client_id: CLIENT_ID,
        };
        let subscribe_bytes = subscribe.encode();

        // 每个发送方一个重组器：两块板的 frame_id 各自独立，混在一起会互相顶掉。
        let mut assemblers: HashMap<u32, FrameAssembler> = HashMap::new();
        let mut buf = [0u8; RECV_BUFFER];
        let mut last_subscribe = Instant::now() - SUBSCRIBE_INTERVAL;
        let mut last_sweep = Instant::now();

        log::info!(
            "listening on UDP {DEFAULT_PORT}, announcing {:?} to {target}",
            self.role
        );

        while !self.stop.load(Ordering::Relaxed) {
            // 周期性续订就是保活：板子在 6 秒收不到就认为我们走了。
            if last_subscribe.elapsed() >= SUBSCRIBE_INTERVAL {
                if let Err(err) = socket.send_to(&subscribe_bytes, target) {
                    log::warn!("could not subscribe at {target}: {err}");
                }
                last_subscribe = Instant::now();
            }

            match socket.recv_from(&mut buf) {
                Ok((n, from)) => {
                    self.handle_datagram(&buf[..n], from, &mut assemblers);
                }
                Err(err)
                    if err.kind() == ErrorKind::WouldBlock
                        || err.kind() == ErrorKind::TimedOut =>
                {
                    // 帧之间的常态。读超时正是让退出能及时响应的原因。
                }
                Err(err) => {
                    log::warn!("recv failed: {err}");
                    thread::sleep(Duration::from_millis(100));
                }
            }

            // 放弃等太久的半截帧。否则板子在帧中途死掉时，重组器会一直留着它。
            for (id, asm) in assemblers.iter_mut() {
                if asm.expire().is_some() {
                    self.discovery.count(*id, |s| s.frames_dropped += 1);
                }
            }

            if last_sweep.elapsed() >= Duration::from_secs(5) {
                self.discovery.forget_silent_cameras(Duration::from_secs(60));
                last_sweep = Instant::now();
            }
        }

        log::info!("receive loop stopped");
    }

    /// 处理一个数据报。
    fn handle_datagram(
        &self,
        bytes: &[u8],
        from: std::net::SocketAddr,
        assemblers: &mut HashMap<u32, FrameAssembler>,
    ) {
        if bytes.len() < HEADER_LEN {
            return;
        }

        let limit = self.sender_limit();

        // 板子的在线通告。注意这里不看来源端口：广播来的源端口没有意义，身份由
        // 头里的 sender_id 决定。
        if let Some(hello) = protocol::Hello::decode(bytes) {
            // 限制生效时不记录其他板子。理由不只是"不显示"：如果把它们记进
            // discovery，`--list` 式的输出、槽位排序、界面上任何一处遍历 cameras()
            // 的地方都会间接泄露场上还有谁。
            if let (Some(only), SocketAddr::V4(v4)) = (limit, from) {
                if hello.sender_id != only {
                    return;
                }
                self.discovery.note_hello(hello.sender_id, v4, hello.firmware);
                return;
            }

            if let SocketAddr::V4(v4) = from {
                let is_new = !self.discovery.cameras().iter().any(|c| c.id == hello.sender_id);
                self.discovery.note_hello(hello.sender_id, v4, hello.firmware);

                if is_new {
                    log::info!(
                        "found camera 0x{:08X} at {v4}",
                        hello.sender_id
                    );
                }
            }
            return;
        }

        let Some(frag) = Fragment::decode(bytes) else {
            // 既不是 Hello 也不是合法分片：将来的其它消息，或一次损坏。
            return;
        };

        let id = frag.header.sender_id;

        // 视频同样要过闸门。对手的板子一旦把流推到这个端口（IP 已被广播泄露，
        // 端口是公开常量），它的数据不会进任何重组器。
        if let Some(only) = limit {
            if id != only {
                return;
            }
        }

        // 视频帧也算存活证据。
        //
        // `is_online` 看的是 `last_seen`，而它此前只由 Hello 更新。于是出现一个
        // 荒唐的状态：画面在每秒刷 22 帧，状态栏却显示"离线"（橙色）。真板子每 2
        // 秒广播一次所以不易撞上，但它本身就是设计错误——正在推流的板子比广播更
        // 可靠地证明自己活着。
        self.discovery.touch(id);

        match assemblers
            .entry(id)
            .or_insert_with(|| FrameAssembler::new(FRAME_TIMEOUT))
            .push(frag)
        {
            Ok(frame) => {
                self.discovery.count(id, |s| s.datagrams += 1);
                self.discovery
                    .note_frame(id, frame.width, frame.height, frame.data.len());
                self.frames.put(frame);
            }
            Err(protocol::DropReason::Superseded) => {
                // 被更新的帧追上了。这不是丢帧：半截的那张只是被超越，在健康链路上
                // 这是帧被放弃的**正常**方式。
                self.discovery.count(id, |s| s.datagrams += 1);
            }
            Err(protocol::DropReason::TimedOut) => {
                self.discovery.count(id, |s| {
                    s.datagrams += 1;
                    s.frames_dropped += 1;
                });
            }
        }
    }

    pub fn shutdown(&self) {
        self.stop.store(true, Ordering::Relaxed);
    }

}

/// 解码后的画面。
pub struct DecodedFrame {
    pub rgba: Vec<u8>,
    pub width: u32,
    pub height: u32,
}

/// 解码一张 JPEG。
///
/// 返回 `None` 表示字节不是可解码的图像——协议层挡不住这种：它只检查分片结构，
/// 不检查 JPEG 本身是否有效。
pub fn decode_frame(jpeg: &[u8]) -> Option<DecodedFrame> {
    let img = image::load_from_memory_with_format(jpeg, image::ImageFormat::Jpeg).ok()?;
    let rgba = img.to_rgba8();

    Some(DecodedFrame {
        width: rgba.width(),
        height: rgba.height(),
        rgba: rgba.into_raw(),
    })
}

/// 装甲板血量的来源。
///
/// **预留，当前没有实现也没有数据源。** 血量显示为占位符。
///
/// 接口和端口都已就位（[`protocol::ARMOR_PORT`]、[`protocol::ArmorHealth`]）：装甲板
/// 接上以后往那个端口发消息即可，裁判界面与其它代码都不用改。
pub trait ArmorSource: Send + Sync {
    /// 某一边的血量。
    ///
    /// 返回 4 个装甲板的当前血量，按 index 排列。没有数据时返回 `None`，界面据此
    /// 显示占位符而不是 0——0 是一个看起来像真值的假值。
    fn health(&self, team: u8) -> Option<[u16; 4]>;
}

/// 一边的装甲板数量。
///
/// RoboMaster 规则是每队 4 块。界面按这个数量画格子，所以数据还没接上时也能看到
/// 位置——空着比画 0 个格子更让人以为布局坏了。
pub const ARMOR_BOARDS: usize = 4;

/// 目前唯一的实现：什么都不提供。
///
/// 存在的意义是让界面代码从一开始就对着一个真实接口开发，而不是等装甲板就绪后再改
/// 一次结构。
pub struct NoArmorYet;

impl ArmorSource for NoArmorYet {
    fn health(&self, _team: u8) -> Option<[u16; 4]> {
        None
    }
}