//! Projection socket coverage: snapshots replace state on reconnect and live
//! deltas patch rows without client-side event folding.

use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;

use axum::extract::ws::{Message, WebSocketUpgrade};
use axum::extract::State;
use axum::response::Response;
use axum::routing::any;
use axum::Router;
use maestro_protocol::{
    ReducedState, StructuredEvent, TranscriptDelta, TranscriptRow, TranscriptRowKind,
    TranscriptSnapshot,
};
use maestro_tui::session_task::SessionState;
use maestro_tui::structured_task::StructuredTask;
use serde_json::json;

async fn structured_handler(
    State(connects): State<Arc<AtomicUsize>>,
    ws: WebSocketUpgrade,
) -> Response {
    ws.on_upgrade(move |mut socket| async move {
        let attempt = connects.fetch_add(1, Ordering::SeqCst);
        let snapshot = if attempt == 0 {
            TranscriptSnapshot {
                cursor: 1,
                rows: vec![row("old", json!({"content":{"type":"text","text":"old"}}))],
            }
        } else {
            TranscriptSnapshot {
                cursor: 3,
                rows: vec![row("live", json!({"content":{"type":"text","text":"one"}}))],
            }
        };
        send(
            &mut socket,
            StructuredEvent::TranscriptSnapshot { snapshot },
        )
        .await;
        if attempt == 0 {
            let _ = socket.close().await;
            return;
        }
        send(
            &mut socket,
            StructuredEvent::ProjectionUpdate {
                transcript: Some(TranscriptDelta::Patch {
                    row: row("live", json!({"content":{"type":"text","text":"two"}})),
                }),
                state: ReducedState {
                    model: Some("test-model".into()),
                    ..ReducedState::default()
                },
            },
        )
        .await;
        tokio::time::sleep(Duration::from_secs(3)).await;
    })
}

fn row(id: &str, data: serde_json::Value) -> TranscriptRow {
    TranscriptRow {
        id: id.into(),
        kind: TranscriptRowKind::AgentText,
        data,
    }
}

async fn send(socket: &mut axum::extract::ws::WebSocket, event: StructuredEvent) {
    socket
        .send(Message::Text(serde_json::to_string(&event).unwrap()))
        .await
        .unwrap();
}

#[tokio::test(flavor = "multi_thread")]
async fn reconnect_snapshot_replaces_then_delta_patches_projection() {
    let connects = Arc::new(AtomicUsize::new(0));
    let app = Router::new()
        .route("/api/structured/x", any(structured_handler))
        .with_state(connects.clone());
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });

    let handle = StructuredTask::spawn(format!("ws://{addr}/api/structured/x"));
    for _ in 0..250 {
        tokio::time::sleep(Duration::from_millis(10)).await;
        let state = handle.state.lock().unwrap().clone();
        if state.reduced.model.as_deref() == Some("test-model") {
            assert_eq!(state.rows.len(), 1);
            assert_eq!(state.rows[0].id, "live");
            assert_eq!(state.rows[0].data["content"]["text"], "two");
            assert!(connects.load(Ordering::SeqCst) >= 2);
            handle.end();
            assert_eq!(handle.connection(), SessionState::Ended);
            return;
        }
    }
    panic!(
        "structured projection never reached reconnect delta; connection={:?}",
        handle.connection()
    );
}