// 独立协议对端，不调用 Dever HTTP/WebSocket；绑到另一个 CPU，避免双方争抢工作线程。
use futures_util::{SinkExt, StreamExt};
use std::{
    io::{BufRead, BufReader, Write},
    process::{Child, Command, Stdio},
};
use tokio::net::{TcpListener, TcpStream};
use tokio_tungstenite::{
    WebSocketStream,
    tungstenite::{Message, protocol::WebSocketConfig},
};

pub fn config() -> WebSocketConfig {
    WebSocketConfig::default()
        .read_buffer_size(4096)
        .write_buffer_size(0)
        .max_write_buffer_size(1048576 + 4096)
        .max_message_size(Some(1048576))
        .max_frame_size(Some(1048576))
}

pub struct Peer {
    child: Child,
    pub port: u16,
}

impl Peer {
    pub fn start(connections: usize) -> Self {
        let cpu = std::env::var("DEVER_WS_PEER_CPU").expect("set a separate peer CPU");
        let executable = std::env::var_os("DEVER_WS_PEER_BINARY")
            .map(std::path::PathBuf::from)
            .unwrap_or_else(|| std::env::current_exe().unwrap());
        let mut child = Command::new("taskset")
            .args(["-c", &cpu])
            .arg(executable)
            .arg("--peer")
            .env("DEVER_WS_PEER_CONNECTIONS", connections.to_string())
            .stdout(Stdio::piped())
            .spawn()
            .unwrap();
        let mut ready = String::new();
        BufReader::new(child.stdout.take().unwrap())
            .read_line(&mut ready)
            .unwrap();
        let port = ready.trim().parse().expect("peer did not report its port");
        Self { child, port }
    }

    pub fn finish(&mut self) {
        assert!(self.child.wait().unwrap().success());
    }
}

impl Drop for Peer {
    fn drop(&mut self) {
        // 断言失败也回收本次测试创建的子进程。
        let _ = self.child.kill();
        let _ = self.child.wait();
    }
}

async fn serve(mut socket: WebSocketStream<TcpStream>) {
    while let Some(command) = socket.next().await {
        let command = command.unwrap();
        if command.is_close() {
            socket.flush().await.unwrap();
            return;
        }
        let (mode, count) = command.to_text().unwrap().split_once(':').unwrap();
        let count: usize = count.parse().unwrap();
        assert!(mode == "batch" || mode == "roundtrip");
        let expected = socket.next().await.unwrap().unwrap();
        socket.send(Message::Text("ready".into())).await.unwrap();
        for _ in 0..count {
            let actual = socket.next().await.unwrap().unwrap();
            assert_eq!(actual, expected);
            if mode == "roundtrip" {
                socket.send(actual).await.unwrap();
            }
        }
        if mode == "batch" {
            socket.send(expected).await.unwrap();
        }
    }
}

pub fn run() {
    let connections: usize = std::env::var("DEVER_WS_PEER_CONNECTIONS")
        .expect("set expected peer connections")
        .parse()
        .unwrap();
    let runtime = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .unwrap();
    runtime.block_on(async {
        tokio::time::timeout(std::time::Duration::from_secs(120), async {
            let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap();
            println!("{}", listener.local_addr().unwrap().port());
            std::io::stdout().flush().unwrap();
            let mut sessions = Vec::new();
            for _ in 0..connections {
                let (stream, _) = listener.accept().await.unwrap();
                stream.set_nodelay(true).unwrap();
                sessions.push(tokio::spawn(async move {
                    let socket =
                        tokio_tungstenite::accept_async_with_config(stream, Some(config()))
                            .await
                            .unwrap();
                    serve(socket).await;
                }));
            }
            for session in sessions {
                session.await.unwrap();
            }
        })
        .await
        .expect("protocol peer timed out");
    });
}