// benches/ws_memory_benchmarks.rs
//
// Criterion benchmarks for the Qrusty WebSocket API backed by MemoryStorage
// (no disk I/O, no fsync).
//
// Run with:
//   cargo bench --bench ws_memory_benchmarks
//
// Compare with ws_benchmarks (RocksDB) to isolate the disk / fsync overhead
// on WebSocket operations.

use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
use futures_util::{SinkExt, StreamExt};
use qrusty::{api::ApiServer, memory_storage::MemoryStorage};
use serde_json::{json, Value};
use std::sync::Arc;
use tokio::net::TcpListener;
use tokio::runtime::Runtime;
use tokio_tungstenite::{connect_async, tungstenite::Message as TMsg};

struct TestServer {
    ws_url: String,
    http_url: String,
}

fn setup_server_memory(rt: &Runtime) -> TestServer {
    rt.block_on(async {
        let storage = Arc::new(MemoryStorage::new());
        let api = ApiServer::new(storage);
        let router = api
            .router()
            .layer(tower_http::cors::CorsLayer::permissive());
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let port = listener.local_addr().unwrap().port();
        tokio::spawn(async move {
            axum::serve(listener, router).await.unwrap();
        });
        TestServer {
            ws_url: format!("ws://127.0.0.1:{}", port),
            http_url: format!("http://127.0.0.1:{}", port),
        }
    })
}

type WsStream =
    tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>;

async fn ws_connect(ws_url: &str) -> WsStream {
    let (ws, _) = connect_async(format!("{}/ws", ws_url)).await.unwrap();
    ws
}

async fn send_text(ws: &mut WsStream, v: Value) {
    ws.send(TMsg::Text(v.to_string().into())).await.unwrap();
}

async fn recv_json(ws: &mut WsStream) -> Value {
    loop {
        match ws.next().await.unwrap().unwrap() {
            TMsg::Text(t) => return serde_json::from_str(&t).unwrap(),
            _ => continue,
        }
    }
}

fn bench_ws_mem_publish(c: &mut Criterion) {
    let rt = Runtime::new().unwrap();
    let srv = setup_server_memory(&rt);

    rt.block_on(async {
        let _ = reqwest::Client::new()
            .post(format!("{}/create-queue", srv.http_url))
            .json(&json!({"name": "mem-bench-pub", "config": {"ordering": "MaxFirst"}}))
            .send()
            .await;
    });

    let mut group = c.benchmark_group("ws_memory");
    group.sample_size(50);

    group.bench_function("publish", |b| {
        b.iter_custom(|iters| {
            rt.block_on(async {
                let mut ws = ws_connect(&srv.ws_url).await;
                let start = std::time::Instant::now();
                for i in 0..iters {
                    send_text(
                        &mut ws,
                        json!({
                            "type": "publish",
                            "queue": "mem-bench-pub",
                            "payload": format!("msg-{}", i),
                            "priority": 0
                        }),
                    )
                    .await;
                    let _resp = recv_json(&mut ws).await;
                }
                start.elapsed()
            })
        });
    });

    group.finish();
}

fn bench_ws_mem_round_trip(c: &mut Criterion) {
    let rt = Runtime::new().unwrap();
    let srv = setup_server_memory(&rt);

    rt.block_on(async {
        let _ = reqwest::Client::new()
            .post(format!("{}/create-queue", srv.http_url))
            .json(&json!({"name": "mem-bench-rt", "config": {"ordering": "MaxFirst"}}))
            .send()
            .await;
    });

    let mut group = c.benchmark_group("ws_memory");
    group.sample_size(20);

    group.bench_function("round_trip", |b| {
        b.iter_custom(|iters| {
            rt.block_on(async {
                let mut ws = ws_connect(&srv.ws_url).await;

                send_text(
                    &mut ws,
                    json!({"type": "subscribe", "queue": "mem-bench-rt"}),
                )
                .await;
                let _ = recv_json(&mut ws).await;

                let mut total = std::time::Duration::ZERO;
                for i in 0..iters {
                    let pub_start = std::time::Instant::now();
                    send_text(
                        &mut ws,
                        json!({
                            "type": "publish",
                            "queue": "mem-bench-rt",
                            "payload": format!("rt-{}", i),
                            "priority": 0,
                        }),
                    )
                    .await;
                    let pub_resp = recv_json(&mut ws).await;
                    let _pub_id = pub_resp["id"].as_str().unwrap().to_owned();

                    let delivery = loop {
                        let f = recv_json(&mut ws).await;
                        if f["type"] == "deliver" {
                            break f;
                        }
                    };
                    let msg_id = delivery["id"].as_str().unwrap().to_owned();

                    send_text(
                        &mut ws,
                        json!({"type": "ack", "queue": "mem-bench-rt", "id": msg_id}),
                    )
                    .await;
                    let _ = recv_json(&mut ws).await;

                    total += pub_start.elapsed();
                }
                total
            })
        });
    });

    group.finish();
}

fn bench_ws_mem_batch_publish(c: &mut Criterion) {
    let rt = Runtime::new().unwrap();
    let srv = setup_server_memory(&rt);

    rt.block_on(async {
        let _ = reqwest::Client::new()
            .post(format!("{}/create-queue", srv.http_url))
            .json(&json!({"name": "mem-bench-batch", "config": {"ordering": "MaxFirst"}}))
            .send()
            .await;
    });

    let mut group = c.benchmark_group("ws_memory");
    group.sample_size(20);

    for batch_size in [10usize, 50, 100] {
        group.bench_with_input(
            BenchmarkId::new("batch_publish", batch_size),
            &batch_size,
            |b, &n| {
                b.iter_custom(|iters| {
                    rt.block_on(async {
                        let mut ws = ws_connect(&srv.ws_url).await;
                        let start = std::time::Instant::now();
                        for iter in 0..iters {
                            for i in 0..n {
                                send_text(
                                    &mut ws,
                                    json!({
                                        "type": "publish",
                                        "queue": "mem-bench-batch",
                                        "payload": format!("bp-{}-{}", iter, i),
                                        "priority": 0,
                                    }),
                                )
                                .await;
                            }
                            for _ in 0..n {
                                let _ = recv_json(&mut ws).await;
                            }
                        }
                        start.elapsed()
                    })
                });
            },
        );
    }

    group.finish();
}

criterion_group!(
    ws_memory_benches,
    bench_ws_mem_publish,
    bench_ws_mem_round_trip,
    bench_ws_mem_batch_publish
);
criterion_main!(ws_memory_benches);
