//! Real-TCP disconnect test for A2A. Binds 127.0.0.1:0, opens a hyper
//! client, reads one SSE frame, drops the connection, asserts the
//! handler's RequestContext.cancel fires within 2s.

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

use http_body_util::{BodyExt, Full};
use hyper::body::Bytes;
use hyper::http::Request;
use hyper_util::client::legacy::Client;
use hyper_util::rt::TokioExecutor;
use klieo_a2a::auth::RequestContext;
use klieo_a2a::error::A2aError;
use klieo_a2a::handler::A2aHandler;
use klieo_a2a::http::A2aHttpServer;
use klieo_a2a::server::{A2aDispatcher, TaskEvent, TaskEventStream};
use klieo_a2a::task_store::{A2aTaskStore, DEFAULT_BUCKET};
use klieo_a2a::types::{SendMessageParams, TaskStatus};
use klieo_auth_common::{AllowAnonymous, Authenticator};
use klieo_bus_memory::MemoryBus;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio_util::sync::CancellationToken;

/// Time to wait after `drop(body_stream)` for hyper to deliver the TCP FIN
/// through the kernel receive buffer and wake the server-side poll. A bare
/// `yield_now` is insufficient on a loaded multi-thread runtime; 100 ms was
/// chosen empirically as the smallest reliable value on CI hardware.
const HYPER_CLOSE_DETECTION_BUFFER: Duration = Duration::from_millis(100);

/// Maximum time we are willing to wait for the request-scoped cancel token
/// to flip after `drop(body_stream)`. Measured wall-clock is ~1.1 s; the
/// 2 s bound leaves headroom for slower CI hosts.
const CANCEL_OBSERVATION_DEADLINE: Duration = Duration::from_secs(2);

/// Tick interval for the cancel-token polling loop. Short enough to surface
/// the success path quickly, long enough to avoid burning CPU.
const POLL_TICK: Duration = Duration::from_millis(20);

/// How long the handler's stream parks before giving up. Must outlast the
/// `CANCEL_OBSERVATION_DEADLINE` so the success path is "DropGuard cancels
/// the stream's pending future"; this branch never fires on a green test.
const HANDLER_PARK: Duration = Duration::from_secs(60);

/// Teardown ceiling — bound on how long we wait for the spawned axum task
/// to exit after `parent.cancel()`.
const TEARDOWN_DEADLINE: Duration = Duration::from_secs(1);

/// Captures the request-scoped cancel token and yields one non-terminal
/// event so the SSE body stays open until the TCP connection drops.
struct TokenCapturingHandler {
    captured: Arc<Mutex<Option<CancellationToken>>>,
}

#[async_trait::async_trait]
impl A2aHandler for TokenCapturingHandler {
    async fn send_streaming_message(
        &self,
        ctx: &RequestContext,
        _params: SendMessageParams,
    ) -> Result<TaskEventStream, A2aError> {
        *self.captured.lock().unwrap() = Some(ctx.cancel.clone());

        let stream = async_stream::stream! {
            yield TaskEvent::new("t-tcp", TaskStatus::Submitted, None, false);
            // Park until the DropGuard fires when hyper closes the body.
            tokio::time::sleep(HANDLER_PARK).await;
        };
        Ok(Box::pin(stream))
    }
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn tcp_client_disconnect_fires_request_cancel() {
    let captured: Arc<Mutex<Option<CancellationToken>>> = Arc::new(Mutex::new(None));
    let handler: Arc<dyn A2aHandler> = Arc::new(TokenCapturingHandler {
        captured: captured.clone(),
    });
    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 task_store = Arc::new(
        A2aTaskStore::new(bus.kv.clone(), DEFAULT_BUCKET.into())
            .with_event_sink(dispatcher.event_sink()),
    );
    let parent = CancellationToken::new();
    let server = Arc::new(A2aHttpServer::new(dispatcher, task_store, parent.clone()));
    let router = server.router();

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

    let client = Client::builder(TokioExecutor::new()).build_http::<Full<Bytes>>();
    let body = serde_json::to_vec(&serde_json::json!({
        "jsonrpc": "2.0",
        "id": 1,
        "method": "SendStreamingMessage",
        "params": {
            "message": {
                "role": "user",
                "parts": [],
                "messageId": "m-tcp",
                "extensions": [],
                "referenceTaskIds": []
            }
        }
    }))
    .unwrap();

    let req = Request::post(format!("http://127.0.0.1:{port}/a2a"))
        .header("content-type", "application/json")
        .body(Full::new(Bytes::from(body)))
        .unwrap();

    let resp = client.request(req).await.unwrap();
    let mut body_stream = resp.into_body();

    // Read at least one frame to ensure the handler ran and the DropGuard
    // is armed before we drop the body.
    let first_frame_wait = tokio::time::timeout(Duration::from_secs(1), body_stream.frame()).await;
    let first_frame_opt = first_frame_wait.expect("timed out waiting for first SSE frame");
    let first_frame_result = first_frame_opt.expect("stream ended before first SSE frame");
    first_frame_result.expect("SSE frame contained an error");

    drop(body_stream); // close the TCP connection — DropGuard must fire.

    tokio::time::sleep(HYPER_CLOSE_DETECTION_BUFFER).await;

    let deadline = std::time::Instant::now() + CANCEL_OBSERVATION_DEADLINE;
    let mut fired = false;
    while std::time::Instant::now() < deadline {
        if let Some(t) = captured.lock().unwrap().clone() {
            if t.is_cancelled() {
                fired = true;
                break;
            }
        }
        tokio::time::sleep(POLL_TICK).await;
    }

    // Tear down BEFORE asserting so a failure here does not leak the
    // spawned axum task or the handler's parked stream.
    parent.cancel();
    let _ = tokio::time::timeout(TEARDOWN_DEADLINE, serve).await;

    assert!(
        fired,
        "cancel did not fire within {:?} after TCP close",
        CANCEL_OBSERVATION_DEADLINE
    );
}
