// 成功会话的独立协议参照。原共享运行库基准保留；这里不声称覆盖别名、取消和资源归属合同。
use dever_runtime::{
    http::websocket::{self as ws, Message, WebSocket},
    task::{Limits, Runtime},
};
use futures_util::{SinkExt, StreamExt};
use std::{
    hint::black_box,
    time::{Duration, Instant},
};
use tokio::net::TcpStream;
use tokio_tungstenite::{WebSocketStream, tungstenite::Message as Frame};

enum Client {
    Dever(WebSocket),
    Rust(WebSocketStream<TcpStream>),
    Boxed(WebSocketStream<layers::Boxed>),
    Locked(layers::Shared<false>),
    Cancelled(layers::Shared<true>),
    Shared {
        socket: shared::Socket,
        _owner: shared::Owner,
    },
}

impl Client {
    async fn send(&mut self, value: &MESSAGE_TYPE) {
        if let Self::Dever(socket) = self {
            match value {
                MESSAGE_TYPE::Variant0(text) => WS_SEND_TEXT(socket, text).await.unwrap(),
                MESSAGE_TYPE::Variant1(bytes) => WS_SEND_BINARY(socket, bytes).await.unwrap(),
                _ => unreachable!(),
            }
            return;
        }
        // 所有参照采用同样的借用输入、拥有消息及写期限。
        let frame = match value {
            MESSAGE_TYPE::Variant0(text) => Frame::Text(text.clone().into()),
            MESSAGE_TYPE::Variant1(bytes) => Frame::Binary(bytes.clone()),
            _ => unreachable!(),
        };
        if let Self::Shared { socket, .. } = self {
            socket.send(frame).await.unwrap();
            return;
        }
        assert!(frame.len() <= 1048576);
        tokio::time::timeout(Duration::from_millis(5000), async {
            match self {
                Self::Rust(socket) => socket.send(frame).await,
                Self::Boxed(socket) => socket.send(frame).await,
                Self::Locked(socket) => socket.send(frame).await,
                Self::Cancelled(socket) => socket.send(frame).await,
                Self::Dever(_) => unreachable!(),
                Self::Shared { .. } => unreachable!(),
            }
        })
        .await
        .unwrap()
        .unwrap();
    }

    async fn receive(&mut self) -> Message {
        if let Self::Dever(socket) = self {
            return socket.receive().await.unwrap().unwrap();
        }
        let frame = if let Self::Shared { socket, .. } = self {
            socket.receive().await.unwrap().unwrap()
        } else {
            tokio::time::timeout(Duration::from_millis(60000), async {
                match self {
                    Self::Rust(socket) => socket.next().await.transpose(),
                    Self::Boxed(socket) => socket.next().await.transpose(),
                    Self::Locked(socket) => socket.receive().await,
                    Self::Cancelled(socket) => socket.receive().await,
                    Self::Dever(_) => unreachable!(),
                    Self::Shared { .. } => unreachable!(),
                }
            })
            .await
            .unwrap()
            .unwrap()
            .unwrap()
        };
        match frame {
            Frame::Text(value) => Message::Text(value.to_string()),
            Frame::Binary(value) => Message::Binary(value),
            _ => panic!("unexpected protocol frame"),
        }
    }

    async fn close(&mut self) {
        match self {
            Self::Dever(socket) => socket.close(1000, "").await.unwrap(),
            Self::Rust(socket) => socket.close(None).await.unwrap(),
            Self::Boxed(socket) => socket.close(None).await.unwrap(),
            Self::Locked(socket) => socket.close().await,
            Self::Cancelled(socket) => socket.close().await,
            Self::Shared { socket, .. } => socket.close(1000, "").await.unwrap(),
        }
    }
}

fn check(actual: Message, expected: &MESSAGE_TYPE) {
    match (actual, expected) {
        (Message::Text(actual), MESSAGE_TYPE::Variant0(expected)) => assert_eq!(&actual, expected),
        (Message::Binary(actual), MESSAGE_TYPE::Variant1(expected)) => {
            assert_eq!(&actual, expected)
        }
        _ => panic!("payload changed"),
    }
}

async fn measure(client: &mut Client, input: &MESSAGE_TYPE, mode: &str, count: usize) -> u128 {
    client
        .send(&MESSAGE_TYPE::Variant0(format!("{mode}:{count}")))
        .await;
    client.send(input).await;
    check(
        client.receive().await,
        &MESSAGE_TYPE::Variant0("ready".to_owned()),
    );
    if mode == "batch" {
        // 批量发送后单独收确认，避免把每条消息的回声等待计入发送成本。
        let start = Instant::now();
        for _ in 0..count {
            client.send(black_box(input)).await;
        }
        let elapsed = start.elapsed().as_nanos();
        check(client.receive().await, input);
        elapsed
    } else {
        let mut elapsed = 0;
        for _ in 0..count {
            let start = Instant::now();
            client.send(black_box(input)).await;
            let actual = client.receive().await;
            elapsed += start.elapsed().as_nanos();
            check(actual, input);
        }
        elapsed
    }
}

fn main() {
    if std::env::args().any(|argument| argument == "--peer") {
        peer::run();
        return;
    }
    let mut peer = peer::Peer::start(2);
    let reference = std::env::var("DEVER_WS_REFERENCE").unwrap_or_else(|_| "raw".to_owned());
    let label = match reference.as_str() {
        "raw" => "rust_protocol",
        "boxed" => "rust_boxed_diagnostic",
        "locked" => "rust_locked_diagnostic",
        "cancelled" => "rust_cancelled_diagnostic",
        "shared" => "rust_shared",
        _ => panic!("unknown WebSocket reference"),
    };
    let runtime = Runtime::new(Limits {
        threads: 1,
        workers: 1,
        tasks: 128,
    })
    .unwrap();
    println!("case,sample,implementation,elapsed_ns,checksum");
    runtime
        .run(async |scope| {
            tokio::time::timeout(Duration::from_secs(120), async {
                let socket = ws::connect(
                    "127.0.0.1",
                    i64::from(peer.port),
                    "/bench",
                    &ws::Limits::default(),
                    scope,
                )
                .await
                .unwrap();
                let stream = TcpStream::connect(("127.0.0.1", peer.port)).await.unwrap();
                stream.set_nodelay(true).unwrap();
                let url = format!("ws://127.0.0.1:{}/bench", peer.port);
                let control = if reference == "raw" || reference == "shared" {
                    let (wire, _) = tokio_tungstenite::client_async_with_config(
                        url,
                        stream,
                        Some(peer::config()),
                    )
                    .await
                    .unwrap();
                    if reference == "shared" {
                        let (socket, owner) = shared::Socket::new(
                            wire,
                            shared::Limits {
                                message: 1048576,
                                idle: Duration::from_millis(60000),
                                write: Duration::from_millis(5000),
                            },
                        );
                        Client::Shared {
                            socket,
                            _owner: owner,
                        }
                    } else {
                        Client::Rust(wire)
                    }
                } else {
                    let stream: layers::Boxed = Box::new(stream);
                    let (wire, _) = tokio_tungstenite::client_async_with_config(
                        url,
                        stream,
                        Some(peer::config()),
                    )
                    .await
                    .unwrap();
                    match reference.as_str() {
                        "boxed" => Client::Boxed(wire),
                        "locked" => Client::Locked(layers::Shared::new(wire)),
                        "cancelled" => Client::Cancelled(layers::Shared::new(wire)),
                        _ => unreachable!(),
                    }
                };
                let mut clients = [Client::Dever(socket), control];
                for bytes in [64, 65536] {
                    for binary in [false, true] {
                        let input = if binary {
                            MESSAGE_TYPE::Variant1(
                                (0..bytes)
                                    .map(|index| (index % 251) as u8)
                                    .collect::<Vec<_>>()
                                    .into(),
                            )
                        } else {
                            MESSAGE_TYPE::Variant0("雪a".repeat(bytes / 4))
                        };
                        for mode in ["batch", "roundtrip"] {
                            let name = format!(
                                "ws_protocol_{}_{bytes}_{mode}",
                                if binary { "binary" } else { "text" }
                            );
                            for client in &mut clients {
                                black_box(measure(client, &input, mode, 20).await);
                            }
                            for sample in 0..15 {
                                let mut order = [("dever", 0), (label, 1), ("control", 0)];
                                order.rotate_left(sample % 3);
                                if sample / 3 % 2 == 1 {
                                    order.swap(0, 1);
                                }
                                for (implementation, index) in order {
                                    let elapsed =
                                        measure(&mut clients[index], &input, mode, 200).await;
                                    println!(
                                        "{name},{sample},{implementation},{elapsed},{}",
                                        bytes * 200
                                    );
                                }
                            }
                        }
                    }
                }
                for client in &mut clients {
                    client.close().await;
                }
            })
            .await
            .expect("protocol comparison timed out");
        })
        .unwrap();
    peer.finish();
}