//! A2A cluster-0.20 leader election + orphan detection tests.
//!
//! Three replicas-on-one-bus scenarios exercise the
//! `A2aDispatcherBuilder::with_leader_election` wiring T2 added and
//! the orphan-write path T3 added to
//! `A2aDispatcher::dispatch_subscribe_to_task`. ADR-020.

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

use async_trait::async_trait;
use axum::body::Body;
use axum::http::{Request, StatusCode};
use bytes::Bytes;
use http_body_util::BodyExt;
use klieo_auth_common::{AllowAnonymous, Authenticator};
use klieo_a2a::handler::{A2aHandler, EchoHandler};
use klieo_a2a::http::A2aHttpServer;
use klieo_a2a::server::{A2aDispatcher, LEADER_TTL};
use klieo_a2a::task_store::{A2aTaskStore, DEFAULT_BUCKET};
use klieo_a2a::types::{Task, TaskStatus};
use klieo_bus_memory::MemoryBus;
use klieo_core::resume::{KvResumeBuffer, ResumeBuffer};
use klieo_core::{BusError, KvEntry, KvStore, Lease, Pubsub, Revision};
use std::sync::Arc;
use std::time::Duration;
use tokio_util::sync::CancellationToken;
use tower::ServiceExt;

const TASK_ID: &str = "t-1";
const LEADER_BUCKET: &str = "klieo-leaders";

fn build_dispatcher_with_leader(
    pubsub: Arc<dyn Pubsub>,
    kv: Arc<dyn KvStore>,
) -> Arc<A2aDispatcher> {
    let handler: Arc<dyn A2aHandler> = Arc::new(EchoHandler::default());
    let auth: Arc<dyn Authenticator> = Arc::new(AllowAnonymous);
    A2aDispatcher::builder()
        .handler(handler)
        .authenticator(auth)
        .pubsub(pubsub)
        .with_leader_election(kv)
        .build_arc()
        .expect("test builder")
}

fn seed_task(_store: &A2aTaskStore) -> Task {
    Task {
        id: TASK_ID.into(),
        contextId: "ctx-1".into(),
        status: TaskStatus::Submitted,
        artifacts: vec![],
        history: vec![],
        metadata: None,
    }
}

fn subscribe_body() -> Vec<u8> {
    serde_json::to_vec(&serde_json::json!({
        "jsonrpc": "2.0",
        "id": 1,
        "method": "SubscribeToTask",
        "params": { "id": TASK_ID }
    }))
    .unwrap()
}

async fn collect_body_text(resp: axum::response::Response) -> String {
    let bytes = tokio::time::timeout(Duration::from_millis(1500), resp.into_body().collect())
        .await
        .expect("body collect timeout")
        .unwrap()
        .to_bytes();
    String::from_utf8(bytes.to_vec()).unwrap()
}

#[tokio::test(flavor = "current_thread")]
async fn leader_alive_during_invoke_subscribe_tails_normally() {
    let bus = MemoryBus::new();
    let pubsub: Arc<dyn Pubsub> = bus.pubsub.clone();
    let kv: Arc<dyn KvStore> = bus.kv.clone();
    let dispatcher_a = build_dispatcher_with_leader(pubsub.clone(), kv.clone());
    let dispatcher_b = build_dispatcher_with_leader(pubsub.clone(), kv.clone());
    let buffer = Arc::new(KvResumeBuffer::new(kv.clone()));

    // Replica A claims leadership for the task — simulates the invoke
    // path of dispatch_send_streaming holding a LeaderHandle.
    let registry_a = dispatcher_a.leader_registry().expect("registry wired");
    let _handle = registry_a
        .claim(format!("a2a.{TASK_ID}"), LEADER_TTL, None, None)
        .await
        .expect("claim");

    // Pre-seed task on replica B's store for the snapshot path.
    let task_store_b = Arc::new(
        A2aTaskStore::new(kv.clone(), DEFAULT_BUCKET.into())
            .with_event_sink(dispatcher_b.event_sink()),
    );
    task_store_b.put(&seed_task(&task_store_b)).await.unwrap();

    // Pre-seed the resume buffer with one event (proof the leader has
    // written) so an orphan check would have something to anchor on.
    buffer
        .record(TASK_ID, 1, Bytes::from_static(b"seeded"))
        .await
        .unwrap();

    // Replica A pushes a terminal Completed transition shortly after
    // subscribe begins so the SSE body closes cleanly.
    let task_store_a = Arc::new(
        A2aTaskStore::new(kv.clone(), DEFAULT_BUCKET.into())
            .with_event_sink(dispatcher_a.event_sink()),
    );
    tokio::spawn(async move {
        tokio::time::sleep(Duration::from_millis(40)).await;
        task_store_a
            .put(&Task {
                id: TASK_ID.into(),
                contextId: "ctx-1".into(),
                status: TaskStatus::Completed,
                artifacts: vec![],
                history: vec![],
                metadata: None,
            })
            .await
            .ok();
    });

    let server_b = Arc::new(
        A2aHttpServer::new(dispatcher_b.clone(), task_store_b, CancellationToken::new())
            .with_resume_buffer(buffer.clone() as Arc<dyn ResumeBuffer>),
    );
    let req = Request::post("/a2a")
        .header("content-type", "application/json")
        .body(Body::from(subscribe_body()))
        .unwrap();
    let resp = server_b.router().oneshot(req).await.unwrap();
    assert_eq!(resp.status(), StatusCode::OK);

    let text = collect_body_text(resp).await;
    assert!(
        !text.contains("\"code\":-32099"),
        "leader alive must NOT emit LEADER_DIED; body was: {text}"
    );
}

#[tokio::test(flavor = "current_thread")]
async fn leader_death_writes_terminal_frame_to_resume_buffer() {
    let bus = MemoryBus::new();
    let pubsub: Arc<dyn Pubsub> = bus.pubsub.clone();
    let kv: Arc<dyn KvStore> = bus.kv.clone();
    let dispatcher_b = build_dispatcher_with_leader(pubsub.clone(), kv.clone());
    let buffer = Arc::new(KvResumeBuffer::new(kv.clone()));

    // Pre-seed task on replica B's store. The snapshot+tail path
    // requires task_store.get to return a hit.
    let task_store_b = Arc::new(
        A2aTaskStore::new(kv.clone(), DEFAULT_BUCKET.into())
            .with_event_sink(dispatcher_b.event_sink()),
    );
    task_store_b.put(&seed_task(&task_store_b)).await.unwrap();

    // Pre-seed resume buffer with one event so the orphan check
    // anchors on the buffer + writes the terminal frame at max+1.
    buffer
        .record(TASK_ID, 1, Bytes::from_static(b"alive-1"))
        .await
        .unwrap();

    // Simulate leader claim then explicit drop to guarantee the
    // klieo-leaders entry is deleted by the time the subscribe
    // call probes for it (Drop::drop spawns a best-effort delete).
    let dispatcher_a = build_dispatcher_with_leader(pubsub.clone(), kv.clone());
    let registry_a = dispatcher_a.leader_registry().expect("registry");
    {
        let _handle = registry_a
            .claim(format!("a2a.{TASK_ID}"), LEADER_TTL, None, None)
            .await
            .unwrap();
    }
    // Yield for the spawned delete to settle.
    tokio::time::sleep(Duration::from_millis(50)).await;

    let server_b = Arc::new(
        A2aHttpServer::new(dispatcher_b.clone(), task_store_b, CancellationToken::new())
            .with_resume_buffer(buffer.clone() as Arc<dyn ResumeBuffer>),
    );
    let req = Request::post("/a2a")
        .header("content-type", "application/json")
        .body(Body::from(subscribe_body()))
        .unwrap();
    let resp = server_b.router().oneshot(req).await.unwrap();
    assert_eq!(resp.status(), StatusCode::OK);

    let text = collect_body_text(resp).await;
    assert!(
        text.contains("\"code\":-32099"),
        "expected LEADER_DIED code in response body; got: {text}"
    );
    assert!(
        text.contains("stream leader died"),
        "expected leader died message; got: {text}"
    );
    assert!(
        buffer.is_terminal(TASK_ID).await.unwrap(),
        "resume buffer must be marked terminal after orphan write",
    );
}

#[tokio::test(flavor = "current_thread")]
async fn kv_unavailable_at_claim_degrades_to_pre_cluster_behavior() {
    // Simulate a KV that fails every put/get/delete. Claims and
    // is_alive probes return Err; per ADR-020 the dispatcher must
    // fail open (treat as alive) so the SubscribeToTask response is
    // the regular snapshot+tail path, not a LEADER_DIED envelope.
    let bus = MemoryBus::new();
    let pubsub: Arc<dyn Pubsub> = bus.pubsub.clone();
    let failing_kv: Arc<dyn KvStore> = Arc::new(FailingKv);
    let dispatcher = build_dispatcher_with_leader(pubsub.clone(), failing_kv.clone());
    // Resume buffer is backed by the working in-mem KV so the body
    // collects without unrelated backend errors.
    let buffer = Arc::new(KvResumeBuffer::new(bus.kv.clone()));

    let task_store = Arc::new(
        A2aTaskStore::new(bus.kv.clone(), DEFAULT_BUCKET.into())
            .with_event_sink(dispatcher.event_sink()),
    );
    task_store.put(&seed_task(&task_store)).await.unwrap();

    // Spawn a terminal transition so the snapshot+tail body closes.
    let task_store_publisher = task_store.clone();
    tokio::spawn(async move {
        tokio::time::sleep(Duration::from_millis(30)).await;
        task_store_publisher
            .put(&Task {
                id: TASK_ID.into(),
                contextId: "ctx-1".into(),
                status: TaskStatus::Completed,
                artifacts: vec![],
                history: vec![],
                metadata: None,
            })
            .await
            .ok();
    });

    let server = Arc::new(
        A2aHttpServer::new(dispatcher.clone(), task_store, CancellationToken::new())
            .with_resume_buffer(buffer.clone() as Arc<dyn ResumeBuffer>),
    );
    let req = Request::post("/a2a")
        .header("content-type", "application/json")
        .body(Body::from(subscribe_body()))
        .unwrap();
    let resp = server.router().oneshot(req).await.unwrap();
    assert_eq!(resp.status(), StatusCode::OK);
    let text = collect_body_text(resp).await;
    assert!(
        !text.contains("\"code\":-32099"),
        "KV failure must fail open; LEADER_DIED must NOT surface. Body: {text}"
    );
}

/// `KvStore` test double whose every operation fails with
/// `BusError::Connection`. Used to verify ADR-020 fail-open on
/// the leader-claim + is_alive paths.
struct FailingKv;

#[async_trait]
impl KvStore for FailingKv {
    async fn get(&self, _bucket: &str, _key: &str) -> Result<Option<KvEntry>, BusError> {
        Err(BusError::Connection("simulated".into()))
    }
    async fn put(&self, _bucket: &str, _key: &str, _value: Bytes) -> Result<Revision, BusError> {
        Err(BusError::Connection("simulated".into()))
    }
    async fn cas(
        &self,
        _bucket: &str,
        _key: &str,
        _value: Bytes,
        _expected: Option<Revision>,
    ) -> Result<Revision, BusError> {
        Err(BusError::Connection("simulated".into()))
    }
    async fn delete(&self, _bucket: &str, _key: &str) -> Result<(), BusError> {
        Err(BusError::Connection("simulated".into()))
    }
    async fn lease(&self, _bucket: &str, _key: &str, _ttl: Duration) -> Result<Lease, BusError> {
        Err(BusError::Connection("simulated".into()))
    }
}

/// Sanity guard against silent renames of the leader-bucket name —
/// keeps the operator-runbook value (klieo-leaders) in lockstep
/// with the constant the builder uses.
#[test]
fn leader_bucket_name_is_stable() {
    assert_eq!(LEADER_BUCKET, "klieo-leaders");
}
