//! Auth smoke tests for `A2aServer`.
//!
//! Spec-level guarantees we exercise:
//!
//! 1. A missing/invalid `Authorization` header causes the server to reply
//!    with JSON-RPC `-32001 Unauthenticated`. The request body never
//!    reaches the handler.
//! 2. The verified [`Identity`] propagates into every handler call via
//!    [`RequestContext::caller`] — handlers can use it for per-method
//!    authorisation.
//! 3. `AllowAnonymous` produces `Identity::anonymous()` for every
//!    request; the request body DOES reach the handler.
//!
//! These tests run against the in-process `MemoryBus` (no docker
//! dependency) so they execute on every `cargo test`.

use bytes::Bytes;
use klieo_a2a::auth::RequestContext;
use klieo_auth_common::{AllowAnonymous, AuthError, BearerTokenAuthenticator, Identity};
use klieo_a2a::envelope::{JsonRpcRequest, JsonRpcResponse};
use klieo_a2a::error::A2aError;
use klieo_a2a::handler::A2aHandler;
use klieo_a2a::server::A2aServer;
use klieo_a2a::types::{
    AgentCapabilities, AgentCard, GetExtendedAgentCardParams, SendMessageParams, SendMessageResult,
};
use klieo_bus_memory::MemoryBus;
use klieo_core::{DurableName, Headers};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio_stream::StreamExt;

const APP: &str = "test";
const AGENT: &str = "agent-x";

fn rpc_subject() -> String {
    format!("{APP}.a2a.{AGENT}.rpc")
}

/// Handler that records the verified identity it saw on the most-recent
/// `send_message`. Lets the test assert identity propagation.
#[derive(Default)]
struct CapturingHandler {
    last_caller: Mutex<Option<Option<Identity>>>,
}

#[async_trait::async_trait]
impl A2aHandler for CapturingHandler {
    async fn send_message(
        &self,
        ctx: &RequestContext,
        _params: SendMessageParams,
    ) -> Result<SendMessageResult, A2aError> {
        *self.last_caller.lock().unwrap() = Some(ctx.caller.clone());
        // Reject anonymous callers — exercises the "handler decides what
        // to do with anonymous" contract. Production code should follow
        // this pattern for mutating methods.
        if ctx
            .caller
            .as_ref()
            .map(Identity::is_anonymous)
            .unwrap_or(true)
        {
            return Err(A2aError::Server("anonymous not permitted".into()));
        }
        Err(A2aError::MethodNotFound("ok-but-anon-rejected".into()))
    }

    async fn get_extended_agent_card(
        &self,
        _ctx: &RequestContext,
        _params: GetExtendedAgentCardParams,
    ) -> Result<AgentCard, A2aError> {
        Ok(AgentCard {
            id: "cap".into(),
            name: "Capturing".into(),
            description: None,
            provider: None,
            capabilities: AgentCapabilities {
                streaming: false,
                pushNotifications: false,
                stateTransitionHistory: false,
            },
            skills: vec![],
            interfaces: vec![],
            securitySchemes: serde_json::Map::new(),
            security: vec![],
            extensions: vec![],
            signature: None,
        })
    }
}

async fn round_trip_with_headers(
    bus: &Arc<MemoryBus>,
    reply_to: &str,
    headers: Headers,
    payload: Bytes,
) -> JsonRpcResponse {
    let mut reply_stream = bus
        .pubsub
        .subscribe(reply_to, DurableName::new("test-reply"))
        .await
        .unwrap();
    bus.pubsub
        .publish(&rpc_subject(), payload, headers)
        .await
        .unwrap();
    let msg = tokio::time::timeout(Duration::from_secs(2), reply_stream.next())
        .await
        .expect("reply timeout")
        .expect("stream ended")
        .expect("publish err");
    serde_json::from_slice(&msg.payload).expect("decode reply")
}

fn run_server_with(
    handler: Arc<dyn A2aHandler>,
    authenticator: Arc<dyn klieo_auth_common::Authenticator>,
    bus: Arc<MemoryBus>,
) {
    let server = A2aServer::new(
        APP.into(),
        AGENT.into(),
        handler,
        authenticator,
        bus.pubsub.clone(),
    );
    tokio::spawn(async move {
        server.run().await.expect("server.run");
    });
}

fn send_message_payload() -> Bytes {
    let req = JsonRpcRequest {
        jsonrpc: "2.0".into(),
        method: "SendMessage".into(),
        params: serde_json::json!({
            "message": {
                "messageId": "m-1",
                "contextId": null,
                "taskId": null,
                "role": "user",
                "parts": [],
                "metadata": null,
                "extensions": [],
                "referenceTaskIds": [],
            },
            "configuration": null,
        }),
        id: serde_json::json!(1),
    };
    Bytes::from(serde_json::to_vec(&req).unwrap())
}

fn headers_with_reply(reply_to: &str, authz: Option<&str>) -> Headers {
    let mut h = Headers::new();
    h.insert("Reply-To".into(), reply_to.into());
    if let Some(val) = authz {
        h.insert("Authorization".into(), val.into());
    }
    h
}

#[tokio::test]
async fn missing_authorization_yields_minus_32001() {
    let bus = Arc::new(MemoryBus::new());
    let handler = Arc::new(CapturingHandler::default());
    let authn = Arc::new(BearerTokenAuthenticator::new(|_t| {
        Ok(Identity::new("never-reaches"))
    }));
    run_server_with(handler.clone(), authn, bus.clone());
    tokio::time::sleep(Duration::from_millis(50)).await;

    let reply_to = "test.reply.missing";
    let headers = headers_with_reply(reply_to, None);
    let resp = round_trip_with_headers(&bus, reply_to, headers, send_message_payload()).await;
    let err = resp.error.expect("expected error");
    assert_eq!(err.code, -32001, "expected -32001 Unauthenticated");
    assert!(resp.result.is_none(), "result must be absent on auth-fail");
    assert!(
        handler.last_caller.lock().unwrap().is_none(),
        "handler must not be invoked when auth fails"
    );
}

#[tokio::test]
async fn wrong_scheme_yields_minus_32001() {
    let bus = Arc::new(MemoryBus::new());
    let handler = Arc::new(CapturingHandler::default());
    let authn = Arc::new(BearerTokenAuthenticator::new(|_t| {
        Ok(Identity::new("never-reaches"))
    }));
    run_server_with(handler.clone(), authn, bus.clone());
    tokio::time::sleep(Duration::from_millis(50)).await;

    let reply_to = "test.reply.wrong-scheme";
    let headers = headers_with_reply(reply_to, Some("Basic xxx"));
    let resp = round_trip_with_headers(&bus, reply_to, headers, send_message_payload()).await;
    assert_eq!(resp.error.expect("err").code, -32001);
    assert!(handler.last_caller.lock().unwrap().is_none());
}

#[tokio::test]
async fn verifier_rejection_yields_minus_32001() {
    let bus = Arc::new(MemoryBus::new());
    let handler = Arc::new(CapturingHandler::default());
    let authn = Arc::new(BearerTokenAuthenticator::new(|t| {
        if t == "good" {
            Ok(Identity::new("alice"))
        } else {
            Err(AuthError::Rejected("bad".into()))
        }
    }));
    run_server_with(handler.clone(), authn, bus.clone());
    tokio::time::sleep(Duration::from_millis(50)).await;

    let reply_to = "test.reply.bad-token";
    let headers = headers_with_reply(reply_to, Some("Bearer bad"));
    let resp = round_trip_with_headers(&bus, reply_to, headers, send_message_payload()).await;
    assert_eq!(resp.error.expect("err").code, -32001);
    assert!(handler.last_caller.lock().unwrap().is_none());
}

#[tokio::test]
async fn verified_identity_propagates_into_handler_context() {
    let bus = Arc::new(MemoryBus::new());
    let handler = Arc::new(CapturingHandler::default());
    let authn = Arc::new(BearerTokenAuthenticator::new(|t| {
        if t == "good" {
            Ok(Identity::new("alice"))
        } else {
            Err(AuthError::Rejected("bad".into()))
        }
    }));
    run_server_with(handler.clone(), authn, bus.clone());
    tokio::time::sleep(Duration::from_millis(50)).await;

    let reply_to = "test.reply.good-token";
    let headers = headers_with_reply(reply_to, Some("Bearer good"));
    let _resp = round_trip_with_headers(&bus, reply_to, headers, send_message_payload()).await;
    let captured = handler
        .last_caller
        .lock()
        .unwrap()
        .clone()
        .expect("captured");
    let id = captured.expect("caller present");
    assert_eq!(id.as_str(), "alice");
    assert!(!id.is_anonymous());
}

#[tokio::test]
async fn allow_anonymous_propagates_anonymous_identity_into_handler() {
    let bus = Arc::new(MemoryBus::new());
    let handler = Arc::new(CapturingHandler::default());
    run_server_with(handler.clone(), Arc::new(AllowAnonymous), bus.clone());
    tokio::time::sleep(Duration::from_millis(50)).await;

    let reply_to = "test.reply.anon";
    let headers = headers_with_reply(reply_to, None);
    let _resp = round_trip_with_headers(&bus, reply_to, headers, send_message_payload()).await;
    let captured = handler
        .last_caller
        .lock()
        .unwrap()
        .clone()
        .expect("captured");
    let id = captured.expect("caller present");
    assert!(id.is_anonymous());
}
