//! Pluggable authentication for the A2A HTTP listener.
//!
//! The default listener (`serve` / `build_router`) ships open — fine
//! for local development and for deployments behind an authenticating
//! reverse proxy, **not** safe for direct internet exposure.
//! Embedders that want enforcement use [`serve_with_auth`] /
//! [`build_router_with_auth`] in the parent module and supply an
//! [`AuthValidator`].
//!
//! ## What's enforced
//!
//! Auth is applied to:
//!
//! - `POST /` and `POST /a2a` (JSON-RPC + streaming RPC)
//! - `GET /a2a/stream/:task_id` (SSE)
//!
//! `GET /.well-known/agent-card.json` and the legacy alias stay open
//! per the A2A v1.0 spec — Agent Cards are discovery metadata.
//!
//! ## Bring your own
//!
//! Built-in implementations cover the common shapes:
//!
//! - [`NoAuth`] — the default; allows everything.
//! - [`BearerKeyAuth`] — accepts `Authorization: Bearer <key>` for any
//!   key in a configured allow-list.
//! - [`ApiKeyHeaderAuth`] — accepts `X-API-Key: <key>` (or any
//!   custom header name).
//!
//! Anything more elaborate (OAuth2 introspection, mTLS, signed
//! requests) implements the [`AuthValidator`] trait directly.

use axum::http::HeaderMap;
use std::collections::{HashMap, HashSet};

/// Why authentication failed. Always renders to a `401 Unauthorized`.
#[derive(Debug, Clone)]
pub enum AuthError {
    /// No credentials present on the request.
    Missing,
    /// Credentials present but rejected.
    Invalid,
}

/// Verified caller identity surfaced by an [`AuthValidator`].
///
/// Returned by [`AuthValidator::validate`] when authentication
/// succeeds AND the validator can map the credential to a stable
/// principal. The dispatcher copies this onto
/// `proposal.context["a2a_caller_verified"]` so tools and policies
/// can cross-check it against the peer-supplied `a2a_caller`
/// claims (Parslee-ai/car#187 phase 2).
///
/// Phase 1's `a2a_caller` (carried on `Message.metadata`) is what
/// the peer *claims*; `Identity.subject` here is what the server
/// has *verified* about the bearer / API-key / token. The two
/// distinct surfaces let policies enforce "claims must agree with
/// verified subject" without each policy re-running the auth
/// check.
#[derive(Debug, Clone, Default)]
pub struct Identity {
    /// Stable principal identifier. Format is validator-defined —
    /// could be an OIDC subject, an email address, an account id,
    /// or any other identifier the embedder considers a stable
    /// caller key. Empty string means "no subject known."
    pub subject: String,
    /// Additional verified claims (scopes, tenant id, org id, etc.)
    /// the validator extracted from the credential. Surfaced
    /// alongside the subject under `a2a_caller_verified.claims`.
    pub claims: HashMap<String, serde_json::Value>,
}

#[async_trait::async_trait]
pub trait AuthValidator: Send + Sync {
    /// Validate the request and (optionally) extract the caller's
    /// verified identity. Returns:
    ///
    /// - `Ok(None)` — request passes but no identity is known
    ///   (e.g. `NoAuth`, or an allow-list validator that doesn't
    ///   map keys to subjects).
    /// - `Ok(Some(identity))` — request passes with a verified
    ///   principal the dispatcher can surface to policies.
    /// - `Err(_)` — request is rejected.
    ///
    /// Custom validators (OAuth2 introspection, OIDC, mTLS,
    /// signed-request schemes) return `Ok(Some(_))` to enable the
    /// downstream multi-tenant scoping work (Parslee-ai/car#187).
    async fn validate(&self, headers: &HeaderMap) -> Result<Option<Identity>, AuthError>;

    /// Optional `WWW-Authenticate` challenge value sent on 401 so
    /// peers know what credential to supply. Default `Bearer`.
    fn challenge(&self) -> String {
        "Bearer".to_string()
    }
}

/// Pass-through validator — every request is allowed. Default.
pub struct NoAuth;

#[async_trait::async_trait]
impl AuthValidator for NoAuth {
    async fn validate(&self, _: &HeaderMap) -> Result<Option<Identity>, AuthError> {
        Ok(None)
    }
}

/// Validates the standard `Authorization: Bearer <key>` header
/// against an allow-list.
pub struct BearerKeyAuth {
    allowed: HashSet<String>,
}

impl BearerKeyAuth {
    pub fn new<I, S>(keys: I) -> Self
    where
        I: IntoIterator<Item = S>,
        S: Into<String>,
    {
        Self {
            allowed: keys.into_iter().map(Into::into).collect(),
        }
    }

    pub fn single(key: impl Into<String>) -> Self {
        Self::new(std::iter::once(key.into()))
    }
}

#[async_trait::async_trait]
impl AuthValidator for BearerKeyAuth {
    async fn validate(&self, headers: &HeaderMap) -> Result<Option<Identity>, AuthError> {
        let header = headers
            .get(axum::http::header::AUTHORIZATION)
            .and_then(|v| v.to_str().ok())
            .ok_or(AuthError::Missing)?;
        let key = header
            .strip_prefix("Bearer ")
            .or_else(|| header.strip_prefix("bearer "))
            .ok_or(AuthError::Invalid)?;
        if self.allowed.contains(key) {
            // Allow-list validators have no peer→subject mapping —
            // the key is shared and tells us nothing about *who*
            // sent it. Embedders that need a verified subject ship
            // a custom validator (OIDC, OAuth2 introspection,
            // mTLS) that returns Ok(Some(Identity{...})). Returning
            // Ok(None) here is the honest "passed but anonymous"
            // answer.
            Ok(None)
        } else {
            Err(AuthError::Invalid)
        }
    }
}

/// Validates a custom header (typically `X-API-Key`) against an
/// allow-list. Useful when peers can't or won't set
/// `Authorization`.
pub struct ApiKeyHeaderAuth {
    header_name: String,
    allowed: HashSet<String>,
}

impl ApiKeyHeaderAuth {
    pub fn new<I, S>(header_name: impl Into<String>, keys: I) -> Self
    where
        I: IntoIterator<Item = S>,
        S: Into<String>,
    {
        Self {
            header_name: header_name.into(),
            allowed: keys.into_iter().map(Into::into).collect(),
        }
    }
}

#[async_trait::async_trait]
impl AuthValidator for ApiKeyHeaderAuth {
    async fn validate(&self, headers: &HeaderMap) -> Result<Option<Identity>, AuthError> {
        let key = headers
            .get(self.header_name.as_str())
            .and_then(|v| v.to_str().ok())
            .ok_or(AuthError::Missing)?;
        if self.allowed.contains(key) {
            Ok(None)
        } else {
            Err(AuthError::Invalid)
        }
    }

    fn challenge(&self) -> String {
        format!("ApiKey realm=\"{}\"", self.header_name)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use axum::http::HeaderValue;

    fn hdrs(pairs: &[(&str, &str)]) -> HeaderMap {
        use axum::http::HeaderName;
        let mut h = HeaderMap::new();
        for (k, v) in pairs {
            let name = HeaderName::from_bytes(k.as_bytes()).unwrap();
            h.insert(name, HeaderValue::from_str(v).unwrap());
        }
        h
    }

    #[tokio::test]
    async fn no_auth_passes_everything() {
        let v = NoAuth;
        // NoAuth allows everything and returns no identity.
        assert!(matches!(v.validate(&hdrs(&[])).await, Ok(None)));
    }

    #[tokio::test]
    async fn bearer_accepts_known_key() {
        let v = BearerKeyAuth::single("alpha");
        // Allow-list validators authorize without surfacing a
        // subject — the key is shared, so it identifies the
        // *credential* but not the *caller*. Ok(None) is the
        // honest answer (Parslee-ai/car#187 phase 2 rationale).
        assert!(matches!(
            v.validate(&hdrs(&[("authorization", "Bearer alpha")]))
                .await,
            Ok(None)
        ));
    }

    #[tokio::test]
    async fn bearer_rejects_unknown_and_missing() {
        let v = BearerKeyAuth::single("alpha");
        assert!(matches!(
            v.validate(&hdrs(&[("authorization", "Bearer beta")])).await,
            Err(AuthError::Invalid)
        ));
        assert!(matches!(
            v.validate(&hdrs(&[])).await,
            Err(AuthError::Missing)
        ));
    }

    #[tokio::test]
    async fn api_key_header_accepts_known_key() {
        let v = ApiKeyHeaderAuth::new("X-API-Key", ["secret"]);
        assert!(matches!(
            v.validate(&hdrs(&[("x-api-key", "secret")])).await,
            Ok(None)
        ));
    }

    #[tokio::test]
    async fn custom_validator_can_surface_identity() {
        // Sanity check that the trait shape supports the multi-tenant
        // use case (Parslee-ai/car#187 phase 2). A custom validator
        // mapping an in-memory token table to subjects should
        // produce Ok(Some(Identity{...})) without surface friction.
        struct TokenTableAuth(HashMap<&'static str, &'static str>);
        #[async_trait::async_trait]
        impl AuthValidator for TokenTableAuth {
            async fn validate(&self, headers: &HeaderMap) -> Result<Option<Identity>, AuthError> {
                let token = headers
                    .get(axum::http::header::AUTHORIZATION)
                    .and_then(|v| v.to_str().ok())
                    .and_then(|h| {
                        h.strip_prefix("Bearer ")
                            .or_else(|| h.strip_prefix("bearer "))
                    })
                    .ok_or(AuthError::Missing)?;
                let subject = self.0.get(token).copied().ok_or(AuthError::Invalid)?;
                Ok(Some(Identity {
                    subject: subject.into(),
                    claims: HashMap::new(),
                }))
            }
        }
        let mut table = HashMap::new();
        table.insert("tok-1", "alice@example.com");
        let v = TokenTableAuth(table);
        let id = v
            .validate(&hdrs(&[("authorization", "Bearer tok-1")]))
            .await
            .expect("accepted")
            .expect("identity returned");
        assert_eq!(id.subject, "alice@example.com");
    }
}
