//! A2A cluster-0.24 follower-side failover re-invoke tests.
//!
//! Drives `A2aDispatcher::handle_dead_leader_orphan` directly with a
//! pre-seeded `LeaderEntry` in the KV bucket. Production wiring calls
//! the gate via `dispatch_subscribe_to_task` after a `LeaderProbe::Dead`
//! probe, but the probe and `lookup_entry_with_revision` both go
//! through `kv.get` — either the entry exists for both or for
//! neither. The test-only `pub` on `handle_dead_leader_orphan` (gated
//! by the `test-fixtures` feature) lets these tests verify the gate
//! semantics without contorting the KV layer.

#![cfg(all(feature = "http", feature = "test-fixtures"))]

use async_trait::async_trait;
use bytes::Bytes;
use klieo_auth_common::{AllowAnonymous, Authenticator, Identity};
use klieo_a2a::auth::RequestContext;
use klieo_a2a::envelope::A2aHeaders;
use klieo_a2a::error::A2aError;
use klieo_a2a::handler::{A2aHandler, EchoHandler};
use klieo_a2a::server::{A2aDispatcher, TaskEventStream};
use klieo_a2a::task_store::{A2aTaskStore, DEFAULT_BUCKET};
use klieo_a2a::types::{Message, Role, SendMessageParams, SendMessageResult, Task, TaskStatus};
use klieo_bus_memory::MemoryBus;
use klieo_core::resume::{KvResumeBuffer, ResumeBuffer};
use klieo_core::{KvStore, LeaderEntry, Pubsub};
use std::sync::Arc;
use std::time::Duration;
use tokio_stream::StreamExt;
use uuid::Uuid;

const TASK_ID: &str = "t-failover";
const LEADER_BUCKET: &str = "klieo-leaders";
const PRINCIPAL: &str = "alice@example.com";

fn build_dispatcher(
    handler: Arc<dyn A2aHandler>,
    pubsub: Arc<dyn Pubsub>,
    kv: Arc<dyn KvStore>,
) -> Arc<A2aDispatcher> {
    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 anon_ctx() -> RequestContext {
    RequestContext::new(
        A2aHeaders::decode_from(&klieo_core::Headers::new()),
        Some(Identity::anonymous()),
    )
}

fn build_payload(text: &str) -> Bytes {
    let params = SendMessageParams {
        message: Message {
            messageId: format!("m-{}", Uuid::new_v4()),
            contextId: Some("ctx-failover".into()),
            taskId: Some(TASK_ID.into()),
            role: Role::User,
            parts: vec![klieo_a2a::types::Part::Text {
                content: text.into(),
                metadata: None,
                mediaType: None,
                filename: None,
            }],
            metadata: None,
            extensions: vec![],
            referenceTaskIds: vec![],
        },
        configuration: None,
    };
    Bytes::from(serde_json::to_vec(&params).unwrap())
}

async fn seed_leader_entry(kv: &Arc<dyn KvStore>, entry: LeaderEntry) {
    let key = format!("a2a.{TASK_ID}");
    kv.put(LEADER_BUCKET, &key, entry.to_bytes().unwrap())
        .await
        .unwrap();
}

async fn seed_resume_event(buf: &Arc<KvResumeBuffer>, id: u64) {
    buf.record(TASK_ID, id, Bytes::from_static(b"prior-leader-event"))
        .await
        .unwrap();
}

/// Idempotent handler: returns a single completed Task for any input,
/// reports `is_idempotent() == true`. Echoes the request so tests can
/// see the re-invoke ran.
#[derive(Default)]
struct IdempotentEchoHandler;

#[async_trait]
impl A2aHandler for IdempotentEchoHandler {
    async fn send_message(
        &self,
        _ctx: &RequestContext,
        params: SendMessageParams,
    ) -> Result<SendMessageResult, A2aError> {
        let task = Task {
            id: TASK_ID.into(),
            contextId: "ctx-failover".into(),
            status: TaskStatus::Completed,
            artifacts: vec![],
            history: vec![params.message],
            metadata: None,
        };
        Ok(SendMessageResult::Task(task))
    }

    fn is_idempotent(&self) -> bool {
        true
    }
}

async fn collect_stream(mut stream: TaskEventStream) -> Vec<klieo_a2a::server::TaskEvent> {
    let mut out = Vec::new();
    while let Some(ev) = tokio::time::timeout(Duration::from_millis(800), stream.next())
        .await
        .ok()
        .flatten()
    {
        let terminal = ev.final_event;
        out.push(ev);
        if terminal {
            break;
        }
    }
    out
}

#[tokio::test(flavor = "current_thread")]
async fn idempotent_handler_reinvokes_on_orphan() {
    let bus = MemoryBus::new();
    let pubsub: Arc<dyn Pubsub> = bus.pubsub.clone();
    let kv: Arc<dyn KvStore> = bus.kv.clone();
    let handler: Arc<dyn A2aHandler> = Arc::new(IdempotentEchoHandler);
    let dispatcher = build_dispatcher(handler, pubsub.clone(), kv.clone());

    let task_store = Arc::new(
        A2aTaskStore::new(kv.clone(), DEFAULT_BUCKET.into())
            .with_event_sink(dispatcher.event_sink()),
    );
    let buffer = Arc::new(KvResumeBuffer::new(kv.clone()));
    seed_resume_event(&buffer, 1).await;

    let payload = build_payload("hello-failover");
    let entry = LeaderEntry {
        replica_id: "replica-original".into(),
        payload: Some(payload),
        attempt: 0,
        principal: Some(PRINCIPAL.into()),
    };
    seed_leader_entry(&kv, entry).await;

    let result = match dispatcher
        .handle_dead_leader_orphan(&anon_ctx(), TASK_ID, &task_store, buffer.as_ref())
        .await
    {
        Ok(v) => v,
        Err(e) => panic!("re-invoke gate must not error: {e:?}"),
    };
    let stream = match result {
        Some(s) => s,
        None => panic!("idempotent handler must re-invoke (Some stream)"),
    };
    let events = collect_stream(stream).await;
    assert!(
        !events.is_empty(),
        "re-invoke must yield at least the synthetic completed event",
    );
    assert!(
        events.iter().any(|e| e.task_id == TASK_ID),
        "re-invoke event task_id must match seeded id; got {events:?}",
    );

    // Verify the new leader entry carries attempt+1 and the same payload.
    let updated = dispatcher
        .leader_registry()
        .unwrap()
        .lookup_entry(&format!("a2a.{TASK_ID}"))
        .await
        .unwrap()
        .expect("re-invoke must leave a fresh leader entry");
    assert_eq!(updated.attempt, 1);
    assert_eq!(updated.principal.as_deref(), Some(PRINCIPAL));

    // Verify the marker frame landed in the resume buffer at max+1 = 2.
    let mut replay = buffer.replay(TASK_ID, 1).await.unwrap();
    let mut saw_marker = false;
    while let Some((id, bytes)) = replay.next().await {
        if id == 2 {
            let text = String::from_utf8_lossy(&bytes);
            assert!(
                text.contains("failover-reinvoke"),
                "marker frame must carry the failover-reinvoke event tag; got: {text}",
            );
            saw_marker = true;
        }
    }
    assert!(saw_marker, "expected failover-reinvoke marker at id=2");
}

#[tokio::test(flavor = "current_thread")]
async fn non_idempotent_handler_terminates() {
    let bus = MemoryBus::new();
    let pubsub: Arc<dyn Pubsub> = bus.pubsub.clone();
    let kv: Arc<dyn KvStore> = bus.kv.clone();
    let handler: Arc<dyn A2aHandler> = Arc::new(EchoHandler::default());
    let dispatcher = build_dispatcher(handler, pubsub.clone(), kv.clone());

    let task_store = Arc::new(
        A2aTaskStore::new(kv.clone(), DEFAULT_BUCKET.into())
            .with_event_sink(dispatcher.event_sink()),
    );
    let buffer = Arc::new(KvResumeBuffer::new(kv.clone()));
    seed_resume_event(&buffer, 1).await;

    let entry = LeaderEntry {
        replica_id: "replica-original".into(),
        payload: Some(build_payload("ignored")),
        attempt: 0,
        principal: Some(PRINCIPAL.into()),
    };
    seed_leader_entry(&kv, entry).await;

    let err = match dispatcher
        .handle_dead_leader_orphan(&anon_ctx(), TASK_ID, &task_store, buffer.as_ref())
        .await
    {
        Err(e) => e,
        Ok(_) => panic!("non-idempotent handler must terminate via LeaderDied"),
    };
    assert!(matches!(err, A2aError::LeaderDied { .. }));
    assert!(
        buffer.is_terminal(TASK_ID).await.unwrap(),
        "terminate path must mark the resume buffer terminal",
    );
}

#[tokio::test(flavor = "current_thread")]
async fn attempt_cap_terminates_after_three() {
    let bus = MemoryBus::new();
    let pubsub: Arc<dyn Pubsub> = bus.pubsub.clone();
    let kv: Arc<dyn KvStore> = bus.kv.clone();
    let handler: Arc<dyn A2aHandler> = Arc::new(IdempotentEchoHandler);
    let dispatcher = build_dispatcher(handler, pubsub.clone(), kv.clone());

    let task_store = Arc::new(
        A2aTaskStore::new(kv.clone(), DEFAULT_BUCKET.into())
            .with_event_sink(dispatcher.event_sink()),
    );
    let buffer = Arc::new(KvResumeBuffer::new(kv.clone()));
    seed_resume_event(&buffer, 1).await;

    let entry = LeaderEntry {
        replica_id: "replica-third".into(),
        payload: Some(build_payload("at-cap")),
        attempt: klieo_core::FAILOVER_ATTEMPT_CAP,
        principal: Some(PRINCIPAL.into()),
    };
    seed_leader_entry(&kv, entry).await;

    let err = match dispatcher
        .handle_dead_leader_orphan(&anon_ctx(), TASK_ID, &task_store, buffer.as_ref())
        .await
    {
        Err(e) => e,
        Ok(_) => panic!("attempt cap reached must terminate"),
    };
    assert!(matches!(err, A2aError::LeaderDied { .. }));
}

/// `KvStore` wrapper that forwards every op to an inner store EXCEPT
/// `cas`, which always returns `BusError::CasConflict`. Used to wedge
/// the CAS branch deterministically — a same-task racing claim can't
/// produce a conflict because the gate's own lookup observes the
/// racing claim's bumped revision and CAS'es against it.
struct CasConflictKv {
    inner: Arc<dyn KvStore>,
}

#[async_trait]
impl KvStore for CasConflictKv {
    async fn get(
        &self,
        bucket: &str,
        key: &str,
    ) -> Result<Option<klieo_core::KvEntry>, klieo_core::BusError> {
        self.inner.get(bucket, key).await
    }
    async fn put(
        &self,
        bucket: &str,
        key: &str,
        value: Bytes,
    ) -> Result<klieo_core::Revision, klieo_core::BusError> {
        self.inner.put(bucket, key, value).await
    }
    async fn cas(
        &self,
        _bucket: &str,
        _key: &str,
        _value: Bytes,
        expected: Option<klieo_core::Revision>,
    ) -> Result<klieo_core::Revision, klieo_core::BusError> {
        Err(klieo_core::BusError::CasConflict {
            expected: expected.unwrap_or(0),
            actual: 0,
        })
    }
    async fn delete(&self, bucket: &str, key: &str) -> Result<(), klieo_core::BusError> {
        self.inner.delete(bucket, key).await
    }
    async fn lease(
        &self,
        bucket: &str,
        key: &str,
        ttl: Duration,
    ) -> Result<klieo_core::Lease, klieo_core::BusError> {
        self.inner.lease(bucket, key, ttl).await
    }
}

#[tokio::test(flavor = "current_thread")]
async fn cas_conflict_falls_back_to_terminate() {
    let bus = MemoryBus::new();
    let pubsub: Arc<dyn Pubsub> = bus.pubsub.clone();
    // Real KV under the wrapper — lookups succeed, CAS deterministically
    // returns BusError::CasConflict to exercise the gate's race-loss
    // branch without time-of-check/time-of-use shenanigans.
    let wrapped_kv: Arc<dyn KvStore> = Arc::new(CasConflictKv {
        inner: bus.kv.clone(),
    });
    let handler: Arc<dyn A2aHandler> = Arc::new(IdempotentEchoHandler);
    let dispatcher = build_dispatcher(handler, pubsub.clone(), wrapped_kv.clone());

    let task_store = Arc::new(
        A2aTaskStore::new(bus.kv.clone(), DEFAULT_BUCKET.into())
            .with_event_sink(dispatcher.event_sink()),
    );
    let buffer = Arc::new(KvResumeBuffer::new(bus.kv.clone()));
    seed_resume_event(&buffer, 1).await;

    let entry = LeaderEntry {
        replica_id: "replica-original".into(),
        payload: Some(build_payload("racy")),
        attempt: 0,
        principal: Some(PRINCIPAL.into()),
    };
    // Seed via the wrapped store so the dispatcher's lookup
    // sees the entry; the wrapper's get/put pass through to the
    // inner MemoryKv.
    let key = format!("a2a.{TASK_ID}");
    wrapped_kv
        .put(LEADER_BUCKET, &key, entry.to_bytes().unwrap())
        .await
        .unwrap();

    let err = match dispatcher
        .handle_dead_leader_orphan(&anon_ctx(), TASK_ID, &task_store, buffer.as_ref())
        .await
    {
        Err(e) => e,
        Ok(_) => panic!("CAS conflict must terminate"),
    };
    assert!(matches!(err, A2aError::LeaderDied { .. }));
}

#[tokio::test(flavor = "current_thread")]
async fn no_cached_payload_terminates() {
    let bus = MemoryBus::new();
    let pubsub: Arc<dyn Pubsub> = bus.pubsub.clone();
    let kv: Arc<dyn KvStore> = bus.kv.clone();
    let handler: Arc<dyn A2aHandler> = Arc::new(IdempotentEchoHandler);
    let dispatcher = build_dispatcher(handler, pubsub.clone(), kv.clone());

    let task_store = Arc::new(
        A2aTaskStore::new(kv.clone(), DEFAULT_BUCKET.into())
            .with_event_sink(dispatcher.event_sink()),
    );
    let buffer = Arc::new(KvResumeBuffer::new(kv.clone()));
    seed_resume_event(&buffer, 1).await;

    let entry = LeaderEntry {
        replica_id: "replica-no-payload".into(),
        payload: None,
        attempt: 0,
        principal: Some(PRINCIPAL.into()),
    };
    seed_leader_entry(&kv, entry).await;

    let err = match dispatcher
        .handle_dead_leader_orphan(&anon_ctx(), TASK_ID, &task_store, buffer.as_ref())
        .await
    {
        Err(e) => e,
        Ok(_) => panic!("missing payload must terminate"),
    };
    assert!(matches!(err, A2aError::LeaderDied { .. }));
}
