//! A2A cross-replica cancel propagation tests.

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

use klieo_a2a::handler::EchoHandler;
use klieo_a2a::server::A2aDispatcher;
use klieo_auth_common::{AllowAnonymous, Authenticator};
use klieo_bus_memory::MemoryBus;
use klieo_core::Pubsub;
use std::sync::Arc;

fn build_dispatcher(pubsub: Arc<dyn Pubsub>) -> Arc<A2aDispatcher> {
    let handler: Arc<dyn klieo_a2a::handler::A2aHandler> = Arc::new(EchoHandler::default());
    let auth: Arc<dyn Authenticator> = Arc::new(AllowAnonymous);
    Arc::new(A2aDispatcher::new(handler, auth, pubsub))
}

#[tokio::test(flavor = "current_thread")]
async fn publish_cancel_validates_subject_token() {
    let bus = MemoryBus::new();
    let pubsub: Arc<dyn Pubsub> = bus.pubsub.clone();
    let dispatcher = build_dispatcher(pubsub);

    for bad in [">", "*", "a.b", "a b", ""] {
        let err = dispatcher.publish_cancel(bad).await.unwrap_err();
        match err {
            klieo_a2a::error::A2aError::Bus(klieo_core::BusError::Invalid(_)) => {}
            other => panic!("expected Bus(Invalid) for {bad:?}, got {other:?}"),
        }
    }
}

#[tokio::test(flavor = "current_thread")]
async fn publish_cancel_publishes_on_per_task_subject() {
    use klieo_core::DurableName;
    use std::time::Duration;
    use tokio_stream::StreamExt as _;

    let bus = MemoryBus::new();
    let pubsub: Arc<dyn Pubsub> = bus.pubsub.clone();
    let dispatcher = build_dispatcher(pubsub.clone());

    // Subscribe to the cancel subject first so the publish lands.
    let mut stream = pubsub
        .subscribe("klieo.a2a.cancel.t-1", DurableName::new("cancel-test-eph"))
        .await
        .unwrap();
    dispatcher.publish_cancel("t-1").await.unwrap();

    let msg = tokio::time::timeout(Duration::from_millis(500), stream.next())
        .await
        .expect("publish_cancel did not deliver to bus subject within 500ms")
        .expect("stream ended")
        .expect("bus error");
    assert_eq!(msg.subject, "klieo.a2a.cancel.t-1");
    msg.ack.ack().await.unwrap();
}

use std::time::Duration;
use tokio_util::sync::CancellationToken;

// Exercises the `klieo.a2a.cancel.>` wildcard subscription path over
// the in-process bus. Relies on `klieo-bus-memory`'s NATS-compatible
// wildcard fanout (single `*`, greedy tail `>`), which mirrors the
// production NATS transport per the `klieo-core::Pubsub` trait
// docstring.
#[tokio::test(flavor = "current_thread")]
async fn cross_replica_cancel_fires_originating_replicas_token() {
    let bus = MemoryBus::new();
    let pubsub: Arc<dyn Pubsub> = bus.pubsub.clone();
    let dispatcher_a = build_dispatcher(pubsub.clone());
    let dispatcher_b = build_dispatcher(pubsub.clone());
    dispatcher_a.with_cancel_subscription();
    dispatcher_b.with_cancel_subscription();
    // Allow the wildcard subscribers to attach.
    tokio::time::sleep(Duration::from_millis(50)).await;

    // Replica A registers an in-flight invoke for t-1.
    let token = CancellationToken::new();
    dispatcher_a
        .cancel_registry()
        .register("t-1".into(), token.clone());

    // Replica B publishes a cancel for t-1 (no registry entry on B).
    dispatcher_b.publish_cancel("t-1").await.unwrap();

    // A's token must fire within 500ms.
    tokio::time::timeout(Duration::from_millis(500), token.cancelled())
        .await
        .expect("expected A's token to fire after B's cancel publish");
}

#[tokio::test(flavor = "current_thread")]
async fn body_drop_on_replica_b_fires_replica_a_token() {
    use axum::body::Body;
    use axum::http::Request;
    use klieo_a2a::http::A2aHttpServer;
    use klieo_a2a::task_store::{A2aTaskStore, DEFAULT_BUCKET};
    use klieo_a2a::types::{Task, TaskStatus};
    use tower::ServiceExt;

    let bus = MemoryBus::new();
    let pubsub: Arc<dyn Pubsub> = bus.pubsub.clone();
    let dispatcher_a = build_dispatcher(pubsub.clone());
    let dispatcher_b = build_dispatcher(pubsub.clone());
    dispatcher_a.with_cancel_subscription();
    dispatcher_b.with_cancel_subscription();
    tokio::time::sleep(Duration::from_millis(50)).await;

    // Replica A registers a long-running invoke for t-cross.
    let token = CancellationToken::new();
    dispatcher_a
        .cancel_registry()
        .register("t-cross".into(), token.clone());

    // Replica B opens an HTTP stream subscribed to t-cross, drops the response.
    let task_store_b = Arc::new(
        A2aTaskStore::new(bus.kv.clone(), DEFAULT_BUCKET.into())
            .with_event_sink(dispatcher_b.event_sink()),
    );
    task_store_b
        .put(&Task {
            id: "t-cross".into(),
            contextId: "ctx-cross".into(),
            status: TaskStatus::Submitted,
            artifacts: vec![],
            history: vec![],
            metadata: None,
        })
        .await
        .unwrap();
    let server_b = Arc::new(A2aHttpServer::new(
        dispatcher_b.clone(),
        task_store_b,
        tokio_util::sync::CancellationToken::new(),
    ));
    let body = serde_json::to_vec(&serde_json::json!({
        "jsonrpc": "2.0", "id": 1,
        "method": "SubscribeToTask",
        "params": {"id": "t-cross"}
    }))
    .unwrap();
    let req = Request::post("/a2a")
        .header("content-type", "application/json")
        .body(Body::from(body))
        .unwrap();
    let resp = server_b.router().oneshot(req).await.unwrap();
    drop(resp); // body drop → CancelOnDrop fires + publishes

    tokio::time::timeout(Duration::from_millis(1000), token.cancelled())
        .await
        .expect("body-drop on replica B should fire replica A's token");
}

#[tokio::test(flavor = "current_thread")]
async fn unknown_id_cancel_is_noop_on_replicas() {
    let bus = MemoryBus::new();
    let pubsub: Arc<dyn Pubsub> = bus.pubsub.clone();
    let dispatcher = build_dispatcher(pubsub.clone());
    dispatcher.with_cancel_subscription();
    tokio::time::sleep(Duration::from_millis(50)).await;
    dispatcher.publish_cancel("never-registered").await.unwrap();
    // Success = no panic + ack-on-receive completes.
    tokio::time::sleep(Duration::from_millis(50)).await;
}
