//! A2A cross-replica fanout tests. Two `A2aDispatcher` instances
//! sharing one `Arc<dyn Pubsub>` simulate a multi-replica deployment;
//! publish on dispatcher A, subscribe on dispatcher B, assert delivery.

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

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

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 memory_bus_fans_out_task_events_across_dispatchers() {
    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());

    // Subscribe to the per-task subject via the shared pubsub directly,
    // simulating how dispatcher B would receive events published by A.
    let durable = DurableName::new("fanout-test-t-x");
    let mut stream = pubsub
        .subscribe("klieo.a2a.task.t-x", durable)
        .await
        .expect("subscribe");

    let sink = dispatcher_a.event_sink();
    let event =
        TaskEvent::new("t-x".to_string(), TaskStatus::Working, None, false).with_event_id(1);
    sink.send(event).await.unwrap();

    let msg = tokio::time::timeout(Duration::from_millis(500), stream.next())
        .await
        .expect("timeout")
        .expect("stream ended")
        .expect("bus error");
    let got: TaskEvent = serde_json::from_slice(&msg.payload).unwrap();
    msg.ack.ack().await.unwrap();

    assert_eq!(got.event_id, 1);
    assert_eq!(got.task_id, "t-x");
}

#[tokio::test(flavor = "current_thread")]
async fn event_sink_publish_failure_surfaces_typed_error() {
    use async_trait::async_trait;
    use klieo_core::{BusError, MsgStream};

    struct FailingPubsub;
    #[async_trait]
    impl Pubsub for FailingPubsub {
        async fn publish(
            &self,
            _subject: &str,
            _payload: bytes::Bytes,
            _headers: klieo_core::Headers,
        ) -> Result<(), BusError> {
            Err(BusError::Connection("simulated".into()))
        }
        async fn subscribe(
            &self,
            _subject: &str,
            _durable: DurableName,
        ) -> Result<MsgStream, BusError> {
            Err(BusError::Connection("simulated".into()))
        }
    }

    let sink = TaskEventSink::new(Arc::new(FailingPubsub));
    let event =
        TaskEvent::new("t-fail".to_string(), TaskStatus::Submitted, None, false).with_event_id(1);
    let err = sink.send(event).await.unwrap_err();
    use std::error::Error as _;
    assert!(
        err.source().is_some(),
        "Bus variant must preserve source chain"
    );
    match err {
        klieo_a2a::error::A2aError::Bus(e) => {
            assert!(matches!(e, klieo_core::BusError::Connection(_)));
        }
        other => panic!("expected A2aError::Bus, got {other:?}"),
    }
}

#[tokio::test(flavor = "current_thread")]
async fn event_sink_rejects_subject_metacharacter_in_task_id() {
    // Caller-controlled task_id MUST NOT be embeddable into the bus
    // subject namespace; per-task subjects are the cross-tenant boundary.
    // A task_id of `>` would, without validation, subscribe to every
    // tenant's klieo.a2a.task.* stream (CWE-74 subject injection).
    let bus = MemoryBus::new();
    let pubsub: Arc<dyn Pubsub> = bus.pubsub.clone();
    let dispatcher = build_dispatcher(pubsub.clone());
    let sink = dispatcher.event_sink();

    for bad in [">", "*", "a.b", "a b", "", "a\tb"] {
        let event =
            TaskEvent::new(bad.to_string(), TaskStatus::Submitted, None, false).with_event_id(1);
        let err = sink.send(event).await.unwrap_err();
        match err {
            klieo_a2a::error::A2aError::Bus(klieo_core::BusError::Invalid(_)) => {}
            other => panic!(
                "expected A2aError::Bus(BusError::Invalid) for task_id={bad:?}, got {other:?}"
            ),
        }
    }
}

#[tokio::test(flavor = "current_thread")]
async fn event_sink_encode_failure_maps_to_internal_not_parse_error() {
    // Server-side encode failures are an Internal class — they must
    // not surface as JSON-RPC PARSE_ERROR (-32700), which would
    // mis-classify a server bug as a client parse failure (per
    // feedback_error_class_mapping).
    //
    // TaskEvent contains only String/enum/u64/bool fields, so
    // `to_vec` cannot fail in practice. Instead, assert the
    // dispatch shape: `A2aError::Internal { source }` maps to
    // SERVER_ERROR with `source().is_some()`.
    use klieo_a2a::error::A2aError;
    use std::error::Error as _;
    let err: A2aError = A2aError::Internal {
        source: Box::new(std::io::Error::other("synthetic encode failure")),
    };
    let resp = err.to_json_rpc_error(serde_json::json!(1));
    assert_eq!(resp.error.as_ref().unwrap().code, -32000);
    assert!(
        err.source().is_some(),
        "Internal must preserve source chain"
    );
}

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

    // Use the task store + HTTP subscribe path so the terminal stream-close
    // logic in `dispatch_streaming` is exercised end-to-end.
    let task_store = Arc::new(
        A2aTaskStore::new(bus.kv.clone(), DEFAULT_BUCKET.into())
            .with_event_sink(dispatcher.event_sink()),
    );
    let server = Arc::new(A2aHttpServer::new(
        dispatcher.clone(),
        task_store.clone(),
        CancellationToken::new(),
    ));

    // Seed a Completed (terminal) task so subscribe_snapshot_then_tail
    // emits a final_event=true and closes the stream immediately.
    server
        .task_store()
        .put(&Task {
            id: "t-end".into(),
            contextId: "ctx-end".into(),
            status: TaskStatus::Completed,
            artifacts: vec![],
            history: vec![],
            metadata: None,
        })
        .await
        .unwrap();

    let body = serde_json::to_vec(&serde_json::json!({
        "jsonrpc": "2.0",
        "id": 1,
        "method": "SubscribeToTask",
        "params": { "id": "t-end" }
    }))
    .unwrap();
    let req = Request::post("/a2a")
        .header("content-type", "application/json")
        .body(Body::from(body))
        .unwrap();
    let resp = server.router().oneshot(req).await.unwrap();
    assert_eq!(resp.status(), StatusCode::OK);

    let bytes = tokio::time::timeout(Duration::from_millis(500), resp.into_body().collect())
        .await
        .expect("timeout")
        .unwrap()
        .to_bytes();
    let text = std::str::from_utf8(&bytes).unwrap();

    // The SSE stream must contain a terminal event and then close.
    // task_event_to_sse_frame serialises the result as:
    //   {"result":{"final":true,"task_id":"t-end","status":"completed",...}}
    assert!(
        text.contains("\"final\":true"),
        "expected final:true in SSE result; got: {text}"
    );
    assert!(
        text.contains("\"task_id\":\"t-end\""),
        "expected t-end task id; got: {text}"
    );
}

#[tokio::test(flavor = "current_thread")]
async fn replica_b_handle_streaming_sees_replica_a_publishes() {
    let bus = MemoryBus::new();
    let pubsub: Arc<dyn Pubsub> = bus.pubsub.clone();
    let buffer = Arc::new(klieo_core::resume::KvResumeBuffer::new(bus.kv.clone()));

    let dispatcher_a = build_dispatcher(pubsub.clone());
    let dispatcher_b = build_dispatcher(pubsub.clone());

    let task_store_b = Arc::new(
        A2aTaskStore::new(bus.kv.clone(), DEFAULT_BUCKET.into())
            .with_event_sink(dispatcher_b.event_sink()),
    );
    let server_b = Arc::new(
        A2aHttpServer::new(dispatcher_b.clone(), task_store_b, CancellationToken::new())
            .with_resume_buffer(buffer.clone() as Arc<dyn klieo_core::resume::ResumeBuffer>),
    );

    // Pre-seed task on replica B so SubscribeToTask snapshot path has a
    // task to return (the snapshot then chains the live bus tail).
    server_b
        .task_store()
        .put(&Task {
            id: "t-cross".into(),
            contextId: "ctx-cross".into(),
            status: TaskStatus::Submitted,
            artifacts: vec![],
            history: vec![],
            metadata: None,
        })
        .await
        .unwrap();

    // Replica A publishes a Working event 50ms later via its own task_store.
    let task_store_a = Arc::new(
        A2aTaskStore::new(bus.kv.clone(), DEFAULT_BUCKET.into())
            .with_event_sink(dispatcher_a.event_sink()),
    );
    tokio::spawn(async move {
        tokio::time::sleep(Duration::from_millis(50)).await;
        task_store_a
            .put(&Task {
                id: "t-cross".into(),
                contextId: "ctx-cross".into(),
                status: TaskStatus::Working,
                artifacts: vec![],
                history: vec![],
                metadata: None,
            })
            .await
            .ok();
        // Send a terminal so the SubscribeToTask body closes cleanly.
        tokio::time::sleep(Duration::from_millis(20)).await;
        task_store_a
            .put(&Task {
                id: "t-cross".into(),
                contextId: "ctx-cross".into(),
                status: TaskStatus::Completed,
                artifacts: vec![],
                history: vec![],
                metadata: None,
            })
            .await
            .ok();
    });

    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();
    assert_eq!(resp.status(), StatusCode::OK);
    let bytes = tokio::time::timeout(Duration::from_millis(1000), 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-cross\""),
        "expected t-cross events; got: {text}"
    );
    // Replica A's "working" update should arrive via the bus subscription.
    assert!(
        text.contains("\"status\":\"working\""),
        "expected working status from replica A's publish; got: {text}"
    );
}
