//! A minimal WireGuard *responder* used to prove end-to-end connectivity
//! locally, without needing a real VPN server. It listens on UDP, completes
//! the handshake with a connecting `magpie up` client, and then periodically
//! sends it a small encrypted IPv4 packet so the data path is exercised too.
//!
//!   magpie testserver <服务器私钥> <客户端公钥> [监听地址]

use base64::{engine::general_purpose::STANDARD, Engine};
use boringtun::noise::{Tunn, TunnResult};
use boringtun::x25519::{PublicKey, StaticSecret};
use std::net::{SocketAddr, UdpSocket};
use std::time::{Duration, Instant};

const BUF: usize = 65536;

fn flush() {
    use std::io::Write;
    let _ = std::io::stdout().flush();
}

fn decode_key(s: &str) -> Result<[u8; 32], String> {
    let v = STANDARD
        .decode(s.trim())
        .map_err(|e| format!("base64 解码失败: {e}"))?;
    if v.len() != 32 {
        return Err(format!("密钥长度 {} != 32", v.len()));
    }
    let mut k = [0u8; 32];
    k.copy_from_slice(&v);
    Ok(k)
}

/// Build a tiny valid IPv4 packet (src 10.7.0.1 -> dst 10.7.0.2) carrying a
/// short payload, so the client decapsulates it as real tunnel traffic.
fn probe_packet() -> Vec<u8> {
    let payload = b"magpie-server-probe";
    let mut pkt = vec![
        0x45, 0, 0, 0, 0, 0, 0, 0, 64, 0, 0, 0, 10, 7, 0, 1, 10, 7, 0, 2,
    ];
    pkt.extend_from_slice(payload);
    let total = (pkt.len() as u16).to_be_bytes();
    pkt[2] = total[0];
    pkt[3] = total[1];
    pkt
}

pub fn run(server_priv: &str, client_pub: &str, listen: &str) -> Result<(), String> {
    let sp = StaticSecret::from(decode_key(server_priv)?);
    let cp = PublicKey::from(decode_key(client_pub)?);
    let mut srv =
        Tunn::new(sp, cp, None, None, 0, None).map_err(|e| format!("初始化失败: {e}"))?;

    let sock = UdpSocket::bind(listen).map_err(|e| format!("绑定 {listen} 失败: {e}"))?;
    sock.set_read_timeout(Some(Duration::from_millis(250)))
        .map_err(|e| format!("设置超时失败: {e}"))?;
    println!("● 测试 WireGuard 服务器已监听 {listen}");
    println!("  等待客户端握手...\n");
    flush();

    let mut peer: Option<SocketAddr> = None;
    let mut announced = false;
    let mut last_probe = Instant::now();
    let mut buf = vec![0u8; BUF];

    loop {
        match sock.recv_from(&mut buf) {
            Ok((n, src)) => {
                peer = Some(src);
                let mut input: &[u8] = &buf[..n];
                loop {
                    let mut out = vec![0u8; BUF];
                    let res = srv.decapsulate(None, input, &mut out);
                    match res {
                        TunnResult::WriteToNetwork(b) => {
                            let _ = sock.send_to(b, src);
                            input = &[];
                        }
                        TunnResult::WriteToTunnelV4(b, _)
                        | TunnResult::WriteToTunnelV6(b, _) => {
                            println!("← 收到来自客户端 {src} 的隧道数据 {} 字节", b.len());
                            flush();
                            break;
                        }
                        TunnResult::Err(e) => {
                            eprintln!("解密错误: {e:?}");
                            break;
                        }
                        TunnResult::Done => break,
                    }
                }
            }
            Err(ref e)
                if e.kind() == std::io::ErrorKind::WouldBlock
                    || e.kind() == std::io::ErrorKind::TimedOut => {}
            Err(e) => return Err(format!("recv 错误: {e}")),
        }

        // Periodic timers (sends keepalive/rekey if needed).
        let mut out = vec![0u8; BUF];
        if let TunnResult::WriteToNetwork(b) = srv.update_timers(&mut out) {
            if let Some(p) = peer {
                let _ = sock.send_to(b, p);
            }
        }

        if srv.time_since_last_handshake().is_some() {
            if !announced {
                println!("✓ 与客户端握手成功，安全会话已建立。");
                flush();
                announced = true;
            }
            // Every second, push an encrypted probe packet down to the client.
            if last_probe.elapsed() > Duration::from_secs(1) {
                last_probe = Instant::now();
                if let Some(p) = peer {
                    let pkt = probe_packet();
                    let mut o = vec![0u8; BUF];
                    if let TunnResult::WriteToNetwork(b) = srv.encapsulate(&pkt, &mut o) {
                        let _ = sock.send_to(b, p);
                        println!("→ 向客户端发送加密探测包 ({} 字节密文)", b.len());
                        flush();
                    }
                }
            }
        }
    }
}