//! Disconnect-cancel unit tests for A2A HTTP transport (oneshot path).
//! Asserts that dropping the SSE response body fires the cancellation
//! token surfaced via RequestContext.cancel.

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

use axum::body::Body;
use axum::http::{Request, StatusCode};
use http_body_util::BodyExt;
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::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio_util::sync::CancellationToken;
use tower::ServiceExt;

/// Captures the request-scoped cancel token from the handler context
/// and yields one non-terminal event so the SSE body stays open.
struct TokenCapturingHandler {
    captured_token: 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> {
        // Capture the token so the test can poll it after body drop.
        let mut guard = self.captured_token.lock().unwrap();
        *guard = Some(ctx.cancel.clone());
        drop(guard);

        let stream = async_stream::stream! {
            // Yield one non-terminal event — keeps the SSE body open.
            yield TaskEvent::new("t-1", TaskStatus::Submitted, None, false);
            // Park until the test drops the body (which fires the DropGuard).
            futures::future::pending::<()>().await;
        };
        Ok(Box::pin(stream))
    }
}

#[tokio::test(flavor = "current_thread")]
async fn dropping_sse_body_fires_request_cancel() {
    let captured_token: Arc<Mutex<Option<CancellationToken>>> = Arc::new(Mutex::new(None));
    let handler: Arc<dyn A2aHandler> = Arc::new(TokenCapturingHandler {
        captured_token: captured_token.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));
    let router = server.router();

    let body = serde_json::to_vec(&serde_json::json!({
        "jsonrpc":"2.0","id":1,"method":"SendStreamingMessage",
        "params":{"message":{
            "role":"user","parts":[],
            "messageId":"m-1","extensions":[],"referenceTaskIds":[]
        }}
    }))
    .unwrap();

    let req = Request::post("/a2a")
        .header("content-type", "application/json")
        .body(Body::from(body))
        .unwrap();

    let resp = router.oneshot(req).await.unwrap();
    assert_eq!(resp.status(), StatusCode::OK);

    let mut body_stream = resp.into_body().into_data_stream();
    use futures::StreamExt;
    // Poll one frame so the handler runs and the token is captured.
    let first = tokio::time::timeout(Duration::from_millis(200), body_stream.next())
        .await
        .expect("timed out waiting for first SSE frame")
        .expect("stream closed before first frame");
    drop(first);

    // Snapshot the captured token before dropping the body.
    let token = captured_token
        .lock()
        .unwrap()
        .clone()
        .expect("handler must have populated the token");
    assert!(!token.is_cancelled(), "token must be live before body drop");

    drop(body_stream); // simulate client disconnect — DropGuard fires here.

    // Give the runtime one tick to process the drop.
    tokio::task::yield_now().await;

    assert!(
        token.is_cancelled(),
        "request_cancel must be set after SSE body is dropped"
    );
}

/// Negative case: handler completes on its own before any client disconnect.
/// The `observed_cancel_mid_run` flag must stay `false`, proving the
/// DropGuard does NOT spuriously cancel work that has already finished.
struct FastHandler {
    /// Set if the handler observes cancellation while still running.
    observed_cancel_mid_run: Arc<AtomicBool>,
}

#[async_trait::async_trait]
impl A2aHandler for FastHandler {
    async fn send_streaming_message(
        &self,
        ctx: &RequestContext,
        _params: SendMessageParams,
    ) -> Result<TaskEventStream, A2aError> {
        if ctx.cancel.is_cancelled() {
            self.observed_cancel_mid_run.store(true, Ordering::SeqCst);
        }
        let stream = futures::stream::iter(vec![TaskEvent::new(
            "t-fast",
            TaskStatus::Completed,
            None,
            true,
        )]);
        Ok(Box::pin(stream))
    }
}

#[tokio::test(flavor = "current_thread")]
async fn full_consumption_does_not_cancel_during_handler_work() {
    let observed_mid_run = Arc::new(AtomicBool::new(false));
    let handler: Arc<dyn A2aHandler> = Arc::new(FastHandler {
        observed_cancel_mid_run: observed_mid_run.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));
    let router = server.router();

    let body = serde_json::to_vec(&serde_json::json!({
        "jsonrpc":"2.0","id":1,"method":"SendStreamingMessage",
        "params":{"message":{
            "role":"user","parts":[],
            "messageId":"m-fast","extensions":[],"referenceTaskIds":[]
        }}
    }))
    .unwrap();

    let req = Request::post("/a2a")
        .header("content-type", "application/json")
        .body(Body::from(body))
        .unwrap();

    let resp = router.oneshot(req).await.unwrap();
    assert_eq!(resp.status(), StatusCode::OK);

    // Drain the body fully — do not drop early.
    let bytes = resp.into_body().collect().await.unwrap().to_bytes();
    let text = std::str::from_utf8(&bytes).unwrap();
    // Match the SSE wire token, not a bare substring, so a task_id like
    // `"my-task-update-1"` cannot satisfy the assertion by accident.
    assert!(
        text.contains("event: task-update\n"),
        "expected `event: task-update` SSE field; got: {text}"
    );

    assert!(
        !observed_mid_run.load(Ordering::SeqCst),
        "cancel must NOT fire while handler is still doing work for a successful natural completion"
    );
}
