//! A2A oneshot resume tests. Exercise the Last-Event-ID header on
//! SubscribeToTask plus the no-header fall-through.

#![cfg(feature = "http")]

use axum::body::Body;
use axum::http::{Request, StatusCode};
use futures::StreamExt as _;
use http_body_util::BodyExt;
use klieo_a2a::handler::EchoHandler;
use klieo_a2a::http::A2aHttpServer;
use klieo_a2a::server::A2aDispatcher;
use klieo_a2a::task_store::{A2aTaskStore, DEFAULT_BUCKET};
use klieo_a2a::types::{Task, TaskStatus};
use klieo_auth_common::{AllowAnonymous, Authenticator};
use klieo_bus_memory::MemoryBus;
use klieo_core::resume::{KvResumeBuffer, ResumeBuffer};
use std::sync::Arc;
use std::time::Duration;
use tokio_util::sync::CancellationToken;
use tower::ServiceExt;

struct TestServer {
    server: Arc<A2aHttpServer>,
    buffer: Arc<KvResumeBuffer>,
}

fn build_server() -> TestServer {
    let handler: Arc<dyn klieo_a2a::handler::A2aHandler> = Arc::new(EchoHandler::default());
    let auth: Arc<dyn Authenticator> = Arc::new(AllowAnonymous);
    let dispatcher = Arc::new(A2aDispatcher::with_in_process_pubsub(handler, auth));
    let bus = Arc::new(MemoryBus::new());
    // Use a separate KV bucket for the resume buffer so it does not
    // conflict with the task store bucket.
    let buffer_bus = Arc::new(MemoryBus::new());
    let task_store = Arc::new(
        A2aTaskStore::new(bus.kv.clone(), DEFAULT_BUCKET.into())
            .with_event_sink(dispatcher.event_sink()),
    );
    let parent = CancellationToken::new();
    let buffer = Arc::new(KvResumeBuffer::new(buffer_bus.kv.clone()));
    let server = Arc::new(
        A2aHttpServer::new(dispatcher, task_store, parent)
            .with_resume_buffer(buffer.clone() as Arc<dyn klieo_core::resume::ResumeBuffer>),
    );
    TestServer { server, buffer }
}

fn subscribe_request(task_id: &str, last_event_id: Option<u64>) -> Request<Body> {
    let body = serde_json::to_vec(&serde_json::json!({
        "jsonrpc": "2.0",
        "id": 1,
        "method": "SubscribeToTask",
        "params": { "id": task_id }
    }))
    .unwrap();

    let mut builder = Request::post("/a2a").header("content-type", "application/json");
    if let Some(id) = last_event_id {
        builder = builder.header("last-event-id", id.to_string());
    }
    builder.body(Body::from(body)).unwrap()
}

#[tokio::test(flavor = "current_thread")]
async fn subscribe_with_last_event_id_replays_missed_events() {
    let TestServer { server, buffer } = build_server();

    for id in 1..=5u64 {
        let payload = bytes::Bytes::from(
            serde_json::to_vec(&serde_json::json!({
                "task_id": "t-1",
                "status": "working",
                "message": null,
                "final_event": false,
                "event_id": id
            }))
            .unwrap(),
        );
        buffer.record("t-1", id, payload).await.unwrap();
    }

    let req = subscribe_request("t-1", Some(2));
    let resp = server.router().oneshot(req).await.unwrap();
    assert_eq!(resp.status(), StatusCode::OK);

    // Collect SSE frames with a timeout. The replay emits events 3-5 then
    // the broadcast tail hangs (no live sender). Collect until timeout;
    // by then all three replay frames must have arrived.
    let mut data_stream = resp.into_body().into_data_stream();
    let mut text = String::new();
    while let Ok(Some(Ok(chunk))) =
        tokio::time::timeout(Duration::from_millis(200), data_stream.next()).await
    {
        text.push_str(std::str::from_utf8(&chunk).unwrap_or(""));
    }

    // SSE replay uses the SSE `id:` field for the event sequence number.
    assert!(
        text.contains("id: 3") || text.contains("id:3"),
        "missing SSE id 3 in replay: {text}"
    );
    assert!(
        text.contains("id: 4") || text.contains("id:4"),
        "missing SSE id 4 in replay: {text}"
    );
    assert!(
        text.contains("id: 5") || text.contains("id:5"),
        "missing SSE id 5 in replay: {text}"
    );
    assert!(
        !text.contains("\nid: 2\n") && !text.contains("\nid:2\n"),
        "should not replay since_id event itself: {text}"
    );
}

#[tokio::test(flavor = "current_thread")]
async fn resume_expired_window_returns_typed_error() {
    let TestServer { server, buffer } = build_server();

    // RESUME_RING_CAP = 256. Record 263 events (ids 1..=263) so that the
    // oldest retained is 8 (263 - 256 + 1 = 8), making since_id=5 expired.
    for id in 1..=263u64 {
        let payload = bytes::Bytes::from(
            serde_json::to_vec(&serde_json::json!({
                "task_id": "t-2",
                "status": "working",
                "message": null,
                "final_event": false,
                "event_id": id
            }))
            .unwrap(),
        );
        buffer.record("t-2", id, payload).await.unwrap();
    }

    let req = subscribe_request("t-2", Some(5));
    let resp = server.router().oneshot(req).await.unwrap();

    let bytes = resp.into_body().collect().await.unwrap().to_bytes();
    let envelope: serde_json::Value = serde_json::from_slice(&bytes).unwrap();

    assert_eq!(
        envelope["error"]["code"],
        serde_json::json!(-32010),
        "expected -32010, got: {envelope}"
    );
    let msg = envelope["error"]["message"].as_str().unwrap_or("<missing>");
    assert!(
        msg.contains("resume window expired"),
        "error message should say 'resume window expired', got: {msg}"
    );
}

#[tokio::test(flavor = "current_thread")]
async fn no_last_event_id_falls_through_to_snapshot() {
    let TestServer { server, .. } = build_server();

    // Use a terminal task status so the snapshot-then-tail stream closes
    // immediately after emitting the synthetic replay event, without
    // waiting for a live broadcast.
    let task = Task {
        id: "t-3".into(),
        contextId: "ctx-t-3".into(),
        status: TaskStatus::Completed,
        artifacts: vec![],
        history: vec![],
        metadata: None,
    };
    server.task_store().put(&task).await.unwrap();

    let req = subscribe_request("t-3", None);
    let resp = server.router().oneshot(req).await.unwrap();
    assert_eq!(resp.status(), StatusCode::OK);

    let bytes = resp.into_body().collect().await.unwrap().to_bytes();
    let text = std::str::from_utf8(&bytes).unwrap();

    // Snapshot path emits the synthetic event with the current task state.
    assert!(
        text.contains("\"task_id\":\"t-3\""),
        "expected task_id t-3 in snapshot response: {text}"
    );
}

// ---- TCP integration ----

use http_body_util::Full;
use hyper::body::Bytes as HyperBytes;
use hyper_util::client::legacy::Client;
use hyper_util::rt::TokioExecutor;

/// Time to wait after frame read for hyper to finish flushing.
const FIRST_FRAME_TIMEOUT: Duration = Duration::from_secs(1);

/// Teardown bound for the spawned axum task.
const TEARDOWN_DEADLINE: Duration = Duration::from_secs(1);

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn tcp_disconnect_then_resume_recovers_stream() {
    let handler: Arc<dyn klieo_a2a::handler::A2aHandler> = Arc::new(EchoHandler::default());
    let auth: Arc<dyn Authenticator> = Arc::new(AllowAnonymous);
    let dispatcher = Arc::new(A2aDispatcher::with_in_process_pubsub(handler, auth));
    let bus = Arc::new(MemoryBus::new());
    let buffer_bus = Arc::new(MemoryBus::new());
    let task_store = Arc::new(
        A2aTaskStore::new(bus.kv.clone(), DEFAULT_BUCKET.into())
            .with_event_sink(dispatcher.event_sink()),
    );
    let parent = CancellationToken::new();
    let buffer = Arc::new(KvResumeBuffer::new(buffer_bus.kv.clone()));
    let server = Arc::new(
        A2aHttpServer::new(dispatcher, task_store, parent)
            .with_resume_buffer(buffer.clone() as Arc<dyn klieo_core::resume::ResumeBuffer>),
    );

    // Pre-populate buffer with 5 events for task t-tcp.
    for id in 1..=5u64 {
        buffer
            .record(
                "t-tcp",
                id,
                bytes::Bytes::from(format!(
                    r#"{{"task_id":"t-tcp","status":"working","message":null,"final_event":false,"event_id":{id}}}"#
                )),
            )
            .await
            .unwrap();
    }

    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let port = listener.local_addr().unwrap().port();
    let router = server.router();
    let serve = tokio::spawn(async move {
        let _ = axum::serve(listener, router).await;
    });

    let client = Client::builder(TokioExecutor::new()).build_http::<Full<HyperBytes>>();
    let body = serde_json::to_vec(&serde_json::json!({
        "jsonrpc": "2.0",
        "id": 1,
        "method": "SubscribeToTask",
        "params": { "id": "t-tcp" }
    }))
    .unwrap();
    let req = hyper::http::Request::post(format!("http://127.0.0.1:{port}/a2a"))
        .header("content-type", "application/json")
        .header("last-event-id", "2")
        .body(Full::new(HyperBytes::from(body)))
        .unwrap();
    let resp = client.request(req).await.unwrap();
    let mut body_stream = resp.into_body();

    // Read the first frame — expected to be event_id=3.
    let first = tokio::time::timeout(FIRST_FRAME_TIMEOUT, body_stream.frame())
        .await
        .expect("first frame timeout")
        .expect("body ended early")
        .expect("frame error")
        .into_data()
        .expect("data frame");
    let s = std::str::from_utf8(&first).unwrap();

    drop(body_stream);
    serve.abort();
    let _ = tokio::time::timeout(TEARDOWN_DEADLINE, serve).await;

    // Per T15: replay events appear as `id: N` SSE field, not `"event_id":N` in JSON.
    assert!(
        s.contains("id: 3") || s.contains("id:3") || s.contains("\"event_id\":3"),
        "expected event_id 3 in first frame; got: {s}"
    );
}

// ---- Live emit → record → replay seam ----

/// Drive the full pipeline end-to-end without manually pre-seeding the
/// resume buffer: seed a task, run a `SubscribeToTask` (no Last-Event-ID)
/// through the SSE pipeline, then verify the buffer was populated through
/// the live `build_sse_response` → `tokio::spawn(buffer.record)` path
/// rather than via test-only `buffer.record()` seeding. This drives the
/// T11 write-through seam end-to-end, which the seed-based tests bypass.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn live_emit_through_sse_records_to_buffer() {
    let TestServer { server, buffer } = build_server();
    let task_id = "t-live-seam";

    // Seed task in store. `task_store.put` writes to the broadcast
    // channel but the per-frame buffer record only fires when an SSE
    // response is actively forwarding frames (`build_sse_response`'s
    // mapped closure).
    let task = Task {
        id: task_id.into(),
        contextId: "ctx-live-seam".into(),
        status: TaskStatus::Completed,
        artifacts: vec![],
        history: vec![],
        metadata: None,
    };
    server.task_store().put(&task).await.unwrap();

    // SubscribeToTask without Last-Event-ID → snapshot+tail path runs
    // through `build_sse_response`, which spawns one
    // `buffer.record(task_id, event_id, payload)` per yielded frame.
    let resp = server
        .router()
        .oneshot(subscribe_request(task_id, None))
        .await
        .unwrap();
    assert_eq!(resp.status(), StatusCode::OK);
    let bytes = tokio::time::timeout(Duration::from_millis(500), resp.into_body().collect())
        .await
        .expect("collect timeout")
        .unwrap()
        .to_bytes();
    let text = std::str::from_utf8(&bytes).unwrap();
    assert!(
        text.contains("\"task_id\":\"t-live-seam\""),
        "snapshot path did not deliver the task: {text}"
    );

    // The per-frame `tokio::spawn(buffer.record(...))` may still be
    // draining when `oneshot.await` returns. Give the spawn pool a tick
    // to land the write before the replay check.
    tokio::time::sleep(Duration::from_millis(100)).await;

    // The buffer should now hold at least one event for this task,
    // recorded through the SSE pipeline without any test-only
    // `buffer.record(...)` calls. We don't care which event id landed —
    // only that meta exists. An empty buffer returns `NotFound`; any
    // other outcome (Ok stream OR `Expired` because we asked beyond
    // oldest_id) means the meta is populated and the seam fired.
    // Use a since-id large enough to exhaust the stream cleanly without
    // overflowing the `since_id + 1` boundary check in `replay`.
    let probe_since = 1_000_000u64;
    match buffer.replay(task_id, probe_since).await {
        Err(klieo_core::resume::ResumeError::NotFound(_)) => {
            panic!("buffer not populated by SSE pipeline — write-through seam regressed");
        }
        Ok(_) | Err(_) => {
            // Meta exists: write-through fired.
        }
    }
}
