// benches/ws_benchmarks.rs
//
// Criterion benchmarks for the Qrusty WebSocket API.
//
// Each benchmark binds the server to a real OS-assigned TCP port and
// measures throughput for the most performance-sensitive WS operations:
//
//   ws/publish            – publish a message through the WS API
//   ws/subscribe_ack      – subscribe, receive a pre-seeded message, ack it
//   ws/round_trip         – full cycle: publish → deliver → ack over WS
//
// Run with:
//   cargo bench --bench ws_benchmarks

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

// ---------------------------------------------------------------------------
// Test server setup
// ---------------------------------------------------------------------------

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

fn setup_server(rt: &Runtime) -> TestServer {
    rt.block_on(async {
        let dir = TempDir::new().unwrap();
        let storage = Arc::new(Storage::new(dir.path().to_str().unwrap()).unwrap());
        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),
            _dir: dir,
        }
    })
}

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();
}

/// Read the next text frame and parse it as JSON; skip non-text frames.
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,
        }
    }
}

// ---------------------------------------------------------------------------
// Benchmarks
// ---------------------------------------------------------------------------

/// WS publish throughput: how fast can we push messages through the WS API?
fn bench_ws_publish(c: &mut Criterion) {
    let rt = Runtime::new().unwrap();
    let srv = setup_server(&rt);

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

    let mut group = c.benchmark_group("ws");
    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": "bench-pub",
                            "payload": format!("msg-{}", i),
                            "priority": 0
                        }),
                    )
                    .await;
                    let _resp = recv_json(&mut ws).await;
                }
                start.elapsed()
            })
        });
    });

    group.finish();
}

/// Subscribe-then-ack round-trip: how fast can we deliver and ack a message?
fn bench_ws_round_trip(c: &mut Criterion) {
    let rt = Runtime::new().unwrap();
    let srv = setup_server(&rt);

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

    let mut group = c.benchmark_group("ws");
    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;

                // Subscribe once for the whole run.
                send_text(&mut ws, json!({"type": "subscribe", "queue": "bench-rt"})).await;
                let _ = recv_json(&mut ws).await; // ok

                let mut total = std::time::Duration::ZERO;
                for i in 0..iters {
                    // Publish.
                    let pub_start = std::time::Instant::now();
                    send_text(
                        &mut ws,
                        json!({
                            "type": "publish",
                            "queue": "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();

                    // Wait for delivery.
                    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();

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

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

    group.finish();
}

/// Batch-publish throughput: publish N messages in one bench iteration.
fn bench_ws_batch_publish(c: &mut Criterion) {
    let rt = Runtime::new().unwrap();
    let srv = setup_server(&rt);

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

    let mut group = c.benchmark_group("ws");
    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": "bench-batch",
                                        "payload": format!("bp-{}-{}", iter, i),
                                        "priority": 0,
                                    }),
                                )
                                .await;
                            }
                            // Drain responses.
                            for _ in 0..n {
                                let _ = recv_json(&mut ws).await;
                            }
                        }
                        start.elapsed()
                    })
                });
            },
        );
    }

    group.finish();
}

criterion_group!(
    ws_benches,
    bench_ws_publish,
    bench_ws_round_trip,
    bench_ws_batch_publish
);
criterion_main!(ws_benches);
