argentor-gateway 1.4.7

//! SSO authentication module for enterprise single sign-on.
//!
//! Provides session management, domain-based access control, and axum-compatible
//! middleware and route handlers. The [`SsoProvider::ApiKey`] flow is fully
//! functional end-to-end (see `POST /auth/api-key`).
//!
//! # Supported providers
//!
//! - [`SsoProvider::Oidc`] — OpenID Connect (Google, Auth0, Okta, Azure AD).
//!   **Fully functional** — performs OIDC discovery, token exchange, JWT payload
//!   decoding, issuer validation, and domain-based access control. Uses `reqwest`
//!   for HTTP and manual base64url JWT decoding (no `openidconnect` crate needed).
//! - [`SsoProvider::Saml`] — SAML 2.0 (enterprise identity providers).
//!   **Functional** — base64-decodes the SAML response, extracts NameID and
//!   attributes via regex, validates the StatusCode, and performs domain-based
//!   access control. Does **not** verify XML signatures (assumes an IdP proxy
//!   or WAF has already validated the assertion).
//! - [`SsoProvider::ApiKey`] — API key bridge (creates SSO sessions from existing
//!   API keys). **Fully functional** — no external dependencies required.
//!
//! # Flow
//!
//! 1. Client visits `GET /auth/login` — redirected to the SSO provider.
//! 2. Provider authenticates the user and redirects back to `GET /auth/callback`.
//! 3. Gateway validates the response, creates a [`UserIdentity`] and session token.
//! 4. Subsequent requests include the token via `Authorization: Bearer <token>` or
//!    the `argentor_session` cookie.
//! 5. The [`sso_auth_middleware`] validates the session and injects `UserIdentity`
//!    into request extensions.

use axum::{
    extract::{Query, Request, State},
    http::{header, HeaderMap, StatusCode},
    middleware::Next,
    response::{IntoResponse, Response},
    routing::get,
    Json, Router,
};
use base64::{
    engine::general_purpose::{STANDARD, URL_SAFE_NO_PAD},
    Engine as _,
};
use chrono::{DateTime, Duration, Utc};
use regex::Regex;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use tracing::{info, warn};
use uuid::Uuid;

use argentor_core::{ArgentorError, ArgentorResult};

// ─── Configuration types ────────────────────────────────────────────────────

/// SSO provider configuration.
///
/// Contains all the information needed to initiate and validate an SSO flow
/// with an external identity provider.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SsoConfig {
    /// The SSO provider type.
    pub provider: SsoProvider,
    /// OAuth2 / OIDC client ID.
    pub client_id: String,
    /// OAuth2 / OIDC client secret.
    pub client_secret: String,
    /// The redirect URI registered with the provider (e.g., `https://app.example.com/auth/callback`).
    pub redirect_uri: String,
    /// Issuer URL for OIDC discovery (e.g., `https://accounts.google.com`).
    pub issuer_url: String,
    /// Email domains that are allowed to authenticate (e.g., `["company.com"]`).
    /// Empty list means all domains are allowed.
    pub allowed_domains: Vec<String>,
    /// OAuth2 scopes to request.
    pub scopes: Vec<String>,
    /// Session time-to-live in hours (default: 24).
    pub session_ttl_hours: u32,
}

impl Default for SsoConfig {
    fn default() -> Self {
        Self {
            provider: SsoProvider::Oidc,
            client_id: String::new(),
            client_secret: String::new(),
            redirect_uri: String::new(),
            issuer_url: String::new(),
            allowed_domains: Vec::new(),
            scopes: vec!["openid".into(), "profile".into(), "email".into()],
            session_ttl_hours: 24,
        }
    }
}

/// Supported SSO provider types.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum SsoProvider {
    /// OpenID Connect (Google, Auth0, Okta, Azure AD).
    Oidc,
    /// SAML 2.0 (enterprise identity providers).
    Saml,
    /// API Key bridge — creates SSO sessions from existing API keys.
    ApiKey,
}

// ─── OIDC helpers ───────────────────────────────────────────────────────────

/// Endpoints discovered from an OIDC provider's `.well-known/openid-configuration`.
#[derive(Debug, Clone)]
#[allow(dead_code)]
pub(crate) struct OidcEndpoints {
    /// The authorization endpoint for login redirects.
    /// Used by `login_url_async()` or external callers that need the discovered endpoint.
    pub authorization_endpoint: String,
    /// The token endpoint for exchanging authorization codes.
    pub token_endpoint: String,
    /// The userinfo endpoint (optional, not all providers expose it).
    pub userinfo_endpoint: Option<String>,
    /// The issuer identifier (used to validate `id_token` claims).
    pub issuer: String,
}

/// Discover OIDC endpoints from the issuer's `.well-known/openid-configuration`.
///
/// Performs a `GET` request to `{issuer_url}/.well-known/openid-configuration`,
/// parses the JSON response, and extracts the required endpoints.
async fn discover_oidc_endpoints(issuer_url: &str) -> ArgentorResult<OidcEndpoints> {
    let url = format!(
        "{}/.well-known/openid-configuration",
        issuer_url.trim_end_matches('/')
    );
    let resp = reqwest::get(&url)
        .await
        .map_err(|e| ArgentorError::Http(format!("OIDC discovery request failed: {e}")))?;

    if !resp.status().is_success() {
        return Err(ArgentorError::Http(format!(
            "OIDC discovery returned HTTP {}",
            resp.status()
        )));
    }

    let config: serde_json::Value = resp
        .json()
        .await
        .map_err(|e| ArgentorError::Http(format!("OIDC discovery response not valid JSON: {e}")))?;

    let authorization_endpoint = config["authorization_endpoint"]
        .as_str()
        .ok_or_else(|| {
            ArgentorError::Http("OIDC discovery: missing authorization_endpoint".into())
        })?
        .to_string();

    let token_endpoint = config["token_endpoint"]
        .as_str()
        .ok_or_else(|| ArgentorError::Http("OIDC discovery: missing token_endpoint".into()))?
        .to_string();

    let userinfo_endpoint = config["userinfo_endpoint"].as_str().map(String::from);

    let issuer = config["issuer"]
        .as_str()
        .ok_or_else(|| ArgentorError::Http("OIDC discovery: missing issuer".into()))?
        .to_string();

    Ok(OidcEndpoints {
        authorization_endpoint,
        token_endpoint,
        userinfo_endpoint,
        issuer,
    })
}

/// Decode the payload (claims) from a JWT token without signature verification.
///
/// A JWT has three base64url-encoded parts separated by `.`: header, payload,
/// signature. This function decodes only the payload (middle part) and parses
/// it as a JSON object.
///
/// # Security note
///
/// This does **not** verify the JWT signature. Signature verification is
/// intentionally omitted because the `id_token` is received directly from the
/// token endpoint over TLS — the transport itself authenticates the issuer.
/// For tokens received from untrusted sources (e.g., client-supplied), you
/// must verify the signature against the provider's JWKS.
fn decode_jwt_payload(token: &str) -> ArgentorResult<serde_json::Value> {
    let parts: Vec<&str> = token.split('.').collect();
    if parts.len() != 3 {
        return Err(ArgentorError::Security(format!(
            "Invalid JWT format: expected 3 parts separated by '.', got {}",
            parts.len()
        )));
    }

    let payload_bytes = URL_SAFE_NO_PAD.decode(parts[1]).map_err(|e| {
        ArgentorError::Security(format!("JWT payload base64url decode failed: {e}"))
    })?;

    let payload: serde_json::Value = serde_json::from_slice(&payload_bytes)
        .map_err(|e| ArgentorError::Security(format!("JWT payload is not valid JSON: {e}")))?;

    Ok(payload)
}

// ─── SAML helpers ──────────────────────────────────────────────────────────

/// Claims extracted from a SAML response.
#[derive(Debug, Clone)]
struct SamlClaims {
    /// The NameID value (typically an email or opaque identifier).
    name_id: String,
    /// The user's email address (from NameID or attribute).
    email: String,
    /// The user's display name, if present.
    name: Option<String>,
    /// Roles assigned to the user.
    roles: Vec<String>,
}

/// Extract the text content of an XML element by its local tag name.
///
/// Matches `<tag>value</tag>` or `<ns:tag>value</ns:tag>` and returns the
/// inner text. Only the first match is returned.
fn extract_xml_element(xml: &str, tag: &str) -> Option<String> {
    // Match <tag>...</tag> or <prefix:tag>...</prefix:tag> or <prefix:tag ...>...</prefix:tag>
    let pattern = format!(
        r"<(?:[a-zA-Z0-9_]+:)?{tag}(?:\s[^>]*)?>([^<]+)</(?:[a-zA-Z0-9_]+:)?{tag}>",
        tag = regex::escape(tag)
    );
    let re = Regex::new(&pattern).ok()?;
    re.captures(xml)
        .and_then(|caps| caps.get(1))
        .map(|m| m.as_str().trim().to_string())
}

/// Extract the value of an XML attribute from an element's opening tag.
///
/// Given XML like `<StatusCode Value="urn:oasis:names:tc:SAML:2.0:status:Success"/>`,
/// calling `extract_xml_attr(xml, "StatusCode", "Value")` returns `Some("urn:...Success")`.
fn extract_xml_attr(xml: &str, element: &str, attr: &str) -> Option<String> {
    let pattern = format!(
        r#"<(?:[a-zA-Z0-9_]+:)?{element}[^>]*?\s{attr}\s*=\s*"([^"]+)""#,
        element = regex::escape(element),
        attr = regex::escape(attr),
    );
    let re = Regex::new(&pattern).ok()?;
    re.captures(xml)
        .and_then(|caps| caps.get(1))
        .map(|m| m.as_str().to_string())
}

/// Extract a single SAML attribute value by its `Name`.
///
/// Matches SAML `<Attribute Name="..."><AttributeValue>...</AttributeValue></Attribute>`
/// blocks and returns the first `AttributeValue` for the matching attribute name.
fn extract_saml_attribute(xml: &str, attr_name: &str) -> Option<String> {
    let pattern = format!(
        r#"<(?:[a-zA-Z0-9_]+:)?Attribute\s[^>]*Name\s*=\s*"{name}"[^>]*>.*?<(?:[a-zA-Z0-9_]+:)?AttributeValue[^>]*>([^<]+)</(?:[a-zA-Z0-9_]+:)?AttributeValue>"#,
        name = regex::escape(attr_name),
    );
    let re = Regex::new(&pattern).ok()?;
    re.captures(xml)
        .and_then(|caps| caps.get(1))
        .map(|m| m.as_str().trim().to_string())
}

/// Extract all SAML attribute values for a given attribute `Name`.
///
/// Unlike [`extract_saml_attribute`], this returns **all** `<AttributeValue>`
/// elements within the matching `<Attribute>` block (e.g., multiple roles).
/// Returns `None` if the attribute is not found.
fn extract_saml_attribute_values(xml: &str, attr_name: &str) -> Option<Vec<String>> {
    // First, find the full <Attribute Name="...">...</Attribute> block
    let block_pattern = format!(
        r#"<(?:[a-zA-Z0-9_]+:)?Attribute\s[^>]*Name\s*=\s*"{name}"[^>]*>([\s\S]*?)</(?:[a-zA-Z0-9_]+:)?Attribute>"#,
        name = regex::escape(attr_name),
    );
    let block_re = Regex::new(&block_pattern).ok()?;
    let block = block_re.captures(xml)?.get(1)?.as_str();

    // Then extract all AttributeValue elements within that block
    let value_re = Regex::new(
        r"<(?:[a-zA-Z0-9_]+:)?AttributeValue[^>]*>([^<]+)</(?:[a-zA-Z0-9_]+:)?AttributeValue>",
    )
    .ok()?;
    let values: Vec<String> = value_re
        .captures_iter(block)
        .filter_map(|caps| caps.get(1).map(|m| m.as_str().trim().to_string()))
        .collect();

    if values.is_empty() {
        None
    } else {
        Some(values)
    }
}

/// Parse a base64-encoded SAML response and extract identity claims.
///
/// # Process
///
/// 1. Base64-decode the `SAMLResponse` value.
/// 2. Extract `NameID` (required — typically an email).
/// 3. Extract standard attributes: `name` / `displayName`, `email`, `role`.
///    Also checks full-URI claim names used by Azure AD / ADFS.
/// 4. Validate the `StatusCode` indicates success.
///
/// # Security note
///
/// This does **not** verify the XML signature. In enterprise deployments the
/// SAML response has already been validated by an IdP proxy (e.g., Shibboleth,
/// mod_auth_mellon, or a cloud WAF) before reaching the application. If you
/// need signature verification, use a dedicated SAML library.
fn parse_saml_response(saml_response: &str) -> ArgentorResult<SamlClaims> {
    // 1. Base64-decode the SAMLResponse
    let decoded_bytes = STANDARD
        .decode(saml_response)
        .map_err(|e| ArgentorError::Security(format!("SAML response base64 decode failed: {e}")))?;

    let xml = String::from_utf8(decoded_bytes)
        .map_err(|e| ArgentorError::Security(format!("SAML response is not valid UTF-8: {e}")))?;

    // 2. Extract NameID (required)
    let name_id = extract_xml_element(&xml, "NameID")
        .ok_or_else(|| ArgentorError::Security("SAML response missing NameID element".into()))?;

    // 3. Extract attributes
    let name = extract_saml_attribute(&xml, "name")
        .or_else(|| extract_saml_attribute(&xml, "displayName"))
        .or_else(|| {
            extract_saml_attribute(
                &xml,
                "http://schemas.xmlsoap.org/ws/2005/05/identity/claims/name",
            )
        });

    let email = extract_saml_attribute(&xml, "email")
        .or_else(|| {
            extract_saml_attribute(
                &xml,
                "http://schemas.xmlsoap.org/ws/2005/05/identity/claims/emailaddress",
            )
        })
        .unwrap_or_else(|| name_id.clone());

    let roles: Vec<String> = extract_saml_attribute_values(&xml, "role")
        .or_else(|| {
            extract_saml_attribute_values(
                &xml,
                "http://schemas.microsoft.com/ws/2008/06/identity/claims/role",
            )
        })
        .unwrap_or_default();

    // 4. Validate StatusCode
    if let Some(status_value) = extract_xml_attr(&xml, "StatusCode", "Value") {
        if !status_value.contains("Success") {
            return Err(ArgentorError::Security(format!(
                "SAML authentication failed with status: {status_value}"
            )));
        }
    }

    Ok(SamlClaims {
        name_id,
        email,
        name,
        roles,
    })
}

// ─── User identity ──────────────────────────────────────────────────────────

/// Authenticated user identity from an SSO provider.
///
/// Created after a successful SSO callback and stored in the session map.
/// Injected into request extensions by the SSO middleware for downstream
/// handlers to consume.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UserIdentity {
    /// Unique user identifier (from the IdP `sub` claim or equivalent).
    pub id: String,
    /// User's email address.
    pub email: String,
    /// User's display name, if available.
    pub name: Option<String>,
    /// Email domain (e.g., `company.com`).
    pub domain: String,
    /// Roles assigned to this user.
    pub roles: Vec<String>,
    /// Tenant identifier for multi-tenant deployments.
    pub tenant_id: Option<String>,
    /// When the user was authenticated.
    pub authenticated_at: DateTime<Utc>,
    /// When the session expires.
    pub expires_at: DateTime<Utc>,
}

impl UserIdentity {
    /// Returns `true` if this identity's session has expired.
    pub fn is_expired(&self) -> bool {
        Utc::now() >= self.expires_at
    }
}

// ─── SSO Manager ────────────────────────────────────────────────────────────

/// SSO session manager.
///
/// Handles login URL generation, callback processing, session lifecycle
/// (create/validate/revoke), and domain-based access control.
///
/// # Provider Support
///
/// | Feature | ApiKey | OIDC | SAML |
/// |---|---|---|---|
/// | Login URL generation | Yes | Yes | Yes |
/// | Session management | Yes | Yes | Yes |
/// | Domain-based access control | Yes | Yes | Yes |
/// | Middleware auth | Yes | Yes | Yes |
/// | Token exchange (callback) | N/A | Yes | Yes (regex-based) |
///
/// The **ApiKey** provider is fully functional: call `POST /auth/api-key` to
/// create a session, then use the returned token for all subsequent requests.
///
/// The **OIDC** provider is fully functional: performs OpenID Connect discovery,
/// exchanges the authorization code for tokens, decodes the JWT `id_token`,
/// validates the issuer and `email_verified` claim, checks domain restrictions,
/// and creates a session. No external OIDC library is required.
///
/// The **SAML** provider is functional: base64-decodes the SAMLResponse,
/// extracts NameID and attributes (name, email, roles) via regex, validates
/// the SAML StatusCode, checks domain restrictions, and creates a session.
/// XML signature verification is intentionally omitted — it is expected that
/// an IdP proxy (e.g., Shibboleth, mod\_auth\_mellon) validates signatures
/// before the response reaches this handler.
pub struct SsoManager {
    config: SsoConfig,
    /// Active sessions: session_token -> UserIdentity.
    sessions: RwLock<HashMap<String, UserIdentity>>,
    /// Pending OIDC/SAML states: state_param -> (created_at, nonce).
    /// Used to prevent CSRF in the callback flow.
    pending_states: RwLock<HashMap<String, DateTime<Utc>>>,
}

impl SsoManager {
    /// Create a new `SsoManager` with the given configuration.
    pub fn new(config: SsoConfig) -> Self {
        Self {
            config,
            sessions: RwLock::new(HashMap::new()),
            pending_states: RwLock::new(HashMap::new()),
        }
    }

    /// Return a reference to the SSO configuration.
    pub fn config(&self) -> &SsoConfig {
        &self.config
    }

    // ── Login flow ─────────────────────────────────────────────────────

    /// Generate the SSO login URL.
    ///
    /// The `state` parameter is an opaque value that will be echoed back in
    /// the callback to prevent CSRF attacks. The caller should generate a
    /// random value and store it (e.g., in a cookie) for later verification.
    ///
    /// The returned URL is where the user's browser should be redirected.
    pub fn login_url(&self, state: &str) -> String {
        // Record the pending state for CSRF verification
        if let Ok(mut states) = self.pending_states.write() {
            states.insert(state.to_string(), Utc::now());
        }

        match self.config.provider {
            SsoProvider::Oidc => {
                let scopes = self.config.scopes.join("+");
                format!(
                    "{issuer}/authorize?client_id={client_id}&redirect_uri={redirect_uri}&response_type=code&scope={scopes}&state={state}",
                    issuer = self.config.issuer_url.trim_end_matches('/'),
                    client_id = self.config.client_id,
                    redirect_uri = self.config.redirect_uri,
                )
            }
            SsoProvider::Saml => {
                // SAML AuthnRequest URL — the issuer_url is the IdP SSO endpoint.
                format!(
                    "{issuer}?SAMLRequest=placeholder&RelayState={state}",
                    issuer = self.config.issuer_url.trim_end_matches('/'),
                )
            }
            SsoProvider::ApiKey => {
                // API key auth doesn't use browser redirect.
                format!(
                    "{redirect_uri}?mode=api_key&state={state}",
                    redirect_uri = self.config.redirect_uri,
                )
            }
        }
    }

    // ── Callback handling ──────────────────────────────────────────────

    /// Handle the callback from the SSO provider.
    ///
    /// Validates the authorization code/token, extracts user identity, checks
    /// domain restrictions, and creates a new session.
    ///
    /// Returns `(session_token, UserIdentity)` on success.
    ///
    /// # OIDC
    ///
    /// For OIDC, this method:
    /// 1. Discovers the token endpoint via `.well-known/openid-configuration`.
    /// 2. Exchanges the authorization code for an `id_token` and `access_token`.
    /// 3. Decodes the JWT `id_token` payload (base64url, no signature check
    ///    since the token comes directly from the provider over TLS).
    /// 4. Validates the issuer matches the configured `issuer_url`.
    /// 5. Checks `email_verified` is true.
    /// 6. Validates the email domain against `allowed_domains`.
    /// 7. Creates a session and returns the token + identity.
    ///
    /// # SAML
    ///
    /// For SAML, the `code` parameter contains the base64-encoded SAMLResponse.
    /// This method:
    /// 1. Base64-decodes the SAMLResponse to XML.
    /// 2. Extracts the `NameID` element (required).
    /// 3. Extracts attributes: `name`/`displayName`, `email`, `role` (including
    ///    full-URI claim names used by Azure AD / ADFS).
    /// 4. Validates the `StatusCode` indicates success.
    /// 5. Checks the email domain against `allowed_domains`.
    /// 6. Creates a session and returns the token + identity.
    pub async fn handle_callback(
        &self,
        code: &str,
        state: &str,
    ) -> ArgentorResult<(String, UserIdentity)> {
        // Verify the state parameter was issued by us (CSRF protection)
        let state_valid = self
            .pending_states
            .read()
            .map_err(|e| ArgentorError::Gateway(format!("lock poisoned: {e}")))?
            .contains_key(state);

        if !state_valid {
            return Err(ArgentorError::Security(
                "Invalid or expired SSO state parameter — possible CSRF attack".into(),
            ));
        }

        // Remove the used state
        if let Ok(mut states) = self.pending_states.write() {
            states.remove(state);
        }

        // Validate the authorization code
        if code.is_empty() {
            return Err(ArgentorError::Gateway(
                "Empty authorization code in SSO callback".into(),
            ));
        }

        // ── Provider-specific token exchange ─────────────────────────────
        match self.config.provider {
            SsoProvider::Oidc => {
                // 1. Discover OIDC endpoints
                let endpoints = discover_oidc_endpoints(&self.config.issuer_url).await?;

                // 2. Exchange authorization code for tokens
                let client = reqwest::Client::new();
                let token_response = client
                    .post(&endpoints.token_endpoint)
                    .form(&[
                        ("grant_type", "authorization_code"),
                        ("code", code),
                        ("redirect_uri", self.config.redirect_uri.as_str()),
                        ("client_id", self.config.client_id.as_str()),
                        ("client_secret", self.config.client_secret.as_str()),
                    ])
                    .send()
                    .await
                    .map_err(|e| {
                        ArgentorError::Http(format!("OIDC token exchange request failed: {e}"))
                    })?;

                if !token_response.status().is_success() {
                    let status = token_response.status();
                    let body = token_response
                        .text()
                        .await
                        .unwrap_or_else(|_| "<unreadable>".into());
                    return Err(ArgentorError::Http(format!(
                        "OIDC token endpoint returned HTTP {status}: {body}"
                    )));
                }

                // 3. Parse the token response JSON
                let token_json: serde_json::Value = token_response.json().await.map_err(|e| {
                    ArgentorError::Http(format!("OIDC token response not valid JSON: {e}"))
                })?;

                let id_token_str = token_json["id_token"].as_str().ok_or_else(|| {
                    ArgentorError::Security("OIDC token response missing id_token field".into())
                })?;

                // 4. Decode the JWT id_token payload
                let payload = decode_jwt_payload(id_token_str)?;

                // 5. Validate issuer — the `iss` claim must match the discovered
                //    issuer, which itself must match our configured issuer_url.
                let discovered_issuer = endpoints.issuer.trim_end_matches('/');
                let configured_issuer = self.config.issuer_url.trim_end_matches('/');
                if discovered_issuer != configured_issuer {
                    return Err(ArgentorError::Security(format!(
                        "OIDC discovered issuer mismatch: configured '{configured_issuer}', \
                         discovered '{discovered_issuer}'"
                    )));
                }

                let iss = payload["iss"].as_str().unwrap_or("");
                if iss.trim_end_matches('/') != discovered_issuer {
                    return Err(ArgentorError::Security(format!(
                        "OIDC id_token issuer mismatch: expected '{discovered_issuer}', got '{iss}'"
                    )));
                }

                // 6. Validate email_verified claim
                let email_verified = payload["email_verified"].as_bool().unwrap_or(false);
                if !email_verified {
                    return Err(ArgentorError::Security(
                        "OIDC email not verified by the identity provider".into(),
                    ));
                }

                // 7. Extract claims
                let email = payload["email"].as_str().ok_or_else(|| {
                    ArgentorError::Security("OIDC id_token missing email claim".into())
                })?;
                let name = payload["name"].as_str().map(String::from);
                let sub = payload["sub"].as_str().unwrap_or("unknown");

                // 8. Validate domain
                if !self.is_domain_allowed(email) {
                    return Err(ArgentorError::Security(format!(
                        "Email domain not in allowed list for '{email}'"
                    )));
                }

                // 9. Build identity and create session
                let identity = self.build_identity(
                    sub,
                    email,
                    name.as_deref(),
                    vec!["oidc-user".into()],
                    None,
                );
                let session_token = self.create_session(identity.clone());

                info!(
                    email = email,
                    sub = sub,
                    "OIDC authentication successful, session created"
                );

                Ok((session_token, identity))
            }
            SsoProvider::Saml => {
                // Parse the base64-encoded SAML response (passed as the `code` param)
                let claims = parse_saml_response(code)?;

                // Validate domain restriction
                if !self.is_domain_allowed(&claims.email) {
                    return Err(ArgentorError::Security(format!(
                        "Email domain not in allowed list for '{}'",
                        claims.email
                    )));
                }

                // Build identity and create session
                let mut roles = claims.roles.clone();
                if roles.is_empty() {
                    roles.push("saml-user".into());
                }
                let identity = self.build_identity(
                    &claims.name_id,
                    &claims.email,
                    claims.name.as_deref(),
                    roles,
                    None,
                );
                let session_token = self.create_session(identity.clone());

                info!(
                    email = %claims.email,
                    name_id = %claims.name_id,
                    "SAML authentication successful, session created"
                );

                Ok((session_token, identity))
            }
            SsoProvider::ApiKey => {
                // API key mode doesn't go through the callback flow.
                Err(ArgentorError::Agent(
                    "API key authentication does not use the SSO callback flow \u{2014} \
                     use POST /auth/api-key instead"
                        .into(),
                ))
            }
        }
    }

    // ── Session management ─────────────────────────────────────────────

    /// Create a new session for the given identity.
    ///
    /// Returns the session token (a UUID v4 string). This method is used
    /// both internally after successful SSO callbacks and externally for
    /// API key migration flows.
    pub fn create_session(&self, identity: UserIdentity) -> String {
        let token = Uuid::new_v4().to_string();
        if let Ok(mut sessions) = self.sessions.write() {
            sessions.insert(token.clone(), identity);
        }
        token
    }

    /// Validate a session token and return the associated identity.
    ///
    /// Returns `None` if the token is unknown or the session has expired.
    /// Expired sessions are automatically removed.
    pub fn validate_session(&self, token: &str) -> Option<UserIdentity> {
        // First check if the session exists and is valid.
        let identity = {
            let sessions = self.sessions.read().ok()?;
            sessions.get(token).cloned()
        };

        match identity {
            Some(id) if id.is_expired() => {
                // Auto-revoke expired sessions.
                self.revoke_session(token);
                None
            }
            other => other,
        }
    }

    /// Revoke (logout) a session by its token.
    ///
    /// Returns `true` if the session existed and was removed.
    pub fn revoke_session(&self, token: &str) -> bool {
        if let Ok(mut sessions) = self.sessions.write() {
            sessions.remove(token).is_some()
        } else {
            false
        }
    }

    /// Check if an email domain is in the allowed list.
    ///
    /// If `allowed_domains` is empty, all domains are permitted.
    /// Domain matching is case-insensitive.
    pub fn is_domain_allowed(&self, email: &str) -> bool {
        if self.config.allowed_domains.is_empty() {
            return true;
        }

        let domain = match email.rsplit_once('@') {
            Some((_, d)) => d.to_lowercase(),
            None => return false, // Not a valid email.
        };

        self.config
            .allowed_domains
            .iter()
            .any(|d| d.to_lowercase() == domain)
    }

    /// List all active (non-expired) sessions.
    ///
    /// Returns `(session_token, identity)` pairs. Intended for admin use.
    pub fn active_sessions(&self) -> Vec<(String, UserIdentity)> {
        let sessions = match self.sessions.read() {
            Ok(s) => s,
            Err(_) => return Vec::new(),
        };

        sessions
            .iter()
            .filter(|(_, id)| !id.is_expired())
            .map(|(token, id)| (token.clone(), id.clone()))
            .collect()
    }

    /// Remove all expired sessions and return the count of sessions removed.
    pub fn cleanup_expired(&self) -> usize {
        let mut sessions = match self.sessions.write() {
            Ok(s) => s,
            Err(_) => return 0,
        };

        let before = sessions.len();
        sessions.retain(|_, id| !id.is_expired());
        before - sessions.len()
    }

    /// Create a `UserIdentity` from email and roles, using the configured TTL.
    ///
    /// Helper for building identities after successful provider verification.
    pub fn build_identity(
        &self,
        id: &str,
        email: &str,
        name: Option<&str>,
        roles: Vec<String>,
        tenant_id: Option<String>,
    ) -> UserIdentity {
        let domain = email
            .rsplit_once('@')
            .map(|(_, d)| d.to_string())
            .unwrap_or_default();

        let now = Utc::now();
        let ttl = Duration::hours(i64::from(self.config.session_ttl_hours));

        UserIdentity {
            id: id.to_string(),
            email: email.to_string(),
            name: name.map(ToString::to_string),
            domain,
            roles,
            tenant_id,
            authenticated_at: now,
            expires_at: now + ttl,
        }
    }
}

// ─── SSO route handlers ─────────────────────────────────────────────────────

/// Shared state for SSO route handlers.
#[derive(Clone)]
pub struct SsoState {
    /// The SSO manager instance.
    pub manager: Arc<SsoManager>,
}

/// Query parameters for the login endpoint.
#[derive(Deserialize)]
pub struct LoginQuery {
    /// Optional redirect URL after login completes.
    pub redirect: Option<String>,
}

/// Query parameters for the SSO callback.
#[derive(Deserialize)]
pub struct CallbackQuery {
    /// Authorization code from the provider.
    pub code: Option<String>,
    /// State parameter for CSRF verification.
    pub state: Option<String>,
}

/// Request body for API key authentication.
#[derive(Deserialize)]
pub struct ApiKeyAuthRequest {
    /// The API key to authenticate with.
    pub api_key: String,
    /// Email to associate with the session.
    pub email: Option<String>,
}

/// Response body for successful authentication.
#[derive(Serialize)]
pub struct AuthResponse {
    /// The session token to use for subsequent requests.
    pub token: String,
    /// The authenticated user's identity.
    pub identity: UserIdentity,
    /// When the session expires.
    pub expires_at: DateTime<Utc>,
}

/// `GET /auth/login` — Redirects the user to the SSO provider login page.
async fn sso_login_handler(
    State(state): State<Arc<SsoState>>,
    Query(query): Query<LoginQuery>,
) -> impl IntoResponse {
    // Generate a random state for CSRF protection.
    let csrf_state = Uuid::new_v4().to_string();

    let login_url = state.manager.login_url(&csrf_state);

    // In production, set the state in a secure cookie for callback verification.
    // For now, redirect directly.
    let _redirect = query.redirect.unwrap_or_else(|| "/".to_string());

    (
        StatusCode::TEMPORARY_REDIRECT,
        [(header::LOCATION, login_url)],
        "",
    )
        .into_response()
}

/// `GET /auth/callback` — Handles the SSO provider callback after authentication.
async fn sso_callback_handler(
    State(state): State<Arc<SsoState>>,
    Query(query): Query<CallbackQuery>,
) -> impl IntoResponse {
    let code = match query.code {
        Some(c) => c,
        None => {
            return (
                StatusCode::BAD_REQUEST,
                Json(serde_json::json!({
                    "error": "missing_code",
                    "message": "Authorization code is required in the callback"
                })),
            )
                .into_response();
        }
    };

    let callback_state = query.state.unwrap_or_default();

    match state.manager.handle_callback(&code, &callback_state).await {
        Ok((token, identity)) => {
            let response = AuthResponse {
                expires_at: identity.expires_at,
                token,
                identity,
            };
            (StatusCode::OK, Json(response)).into_response()
        }
        Err(e) => (
            StatusCode::BAD_REQUEST,
            Json(serde_json::json!({
                "error": "sso_callback_failed",
                "message": e.to_string()
            })),
        )
            .into_response(),
    }
}

/// `GET /auth/logout` — Revokes the current session.
async fn sso_logout_handler(
    State(state): State<Arc<SsoState>>,
    headers: HeaderMap,
) -> impl IntoResponse {
    let token = extract_session_token(&headers);

    match token {
        Some(t) => {
            let revoked = state.manager.revoke_session(&t);
            if revoked {
                (
                    StatusCode::OK,
                    Json(serde_json::json!({
                        "status": "logged_out",
                        "message": "Session revoked successfully"
                    })),
                )
                    .into_response()
            } else {
                (
                    StatusCode::OK,
                    Json(serde_json::json!({
                        "status": "no_session",
                        "message": "No active session found for the given token"
                    })),
                )
                    .into_response()
            }
        }
        None => (
            StatusCode::BAD_REQUEST,
            Json(serde_json::json!({
                "error": "missing_token",
                "message": "No session token found in Authorization header or cookie"
            })),
        )
            .into_response(),
    }
}

/// `GET /auth/me` — Returns the current user's identity.
async fn sso_me_handler(
    State(state): State<Arc<SsoState>>,
    headers: HeaderMap,
) -> impl IntoResponse {
    let token = extract_session_token(&headers);

    match token.and_then(|t| state.manager.validate_session(&t)) {
        Some(identity) => (StatusCode::OK, Json(identity)).into_response(),
        None => (
            StatusCode::UNAUTHORIZED,
            Json(serde_json::json!({
                "error": "not_authenticated",
                "message": "No valid session — please log in via /auth/login"
            })),
        )
            .into_response(),
    }
}

/// `POST /auth/api-key` — Creates an SSO session from an API key.
///
/// This endpoint bridges the existing API key auth with the SSO session system,
/// allowing API key holders to get a session token for cookie-based auth.
async fn sso_api_key_handler(
    State(state): State<Arc<SsoState>>,
    Json(body): Json<ApiKeyAuthRequest>,
) -> impl IntoResponse {
    if body.api_key.is_empty() {
        return (
            StatusCode::BAD_REQUEST,
            Json(serde_json::json!({
                "error": "missing_api_key",
                "message": "API key is required"
            })),
        )
            .into_response();
    }

    // Create identity from the API key.
    // In production, the API key would be validated against the auth service.
    let email = body.email.unwrap_or_else(|| "api-key-user@local".into());

    if !state.manager.is_domain_allowed(&email) {
        return (
            StatusCode::FORBIDDEN,
            Json(serde_json::json!({
                "error": "domain_not_allowed",
                "message": format!(
                    "Email domain is not in the allowed list: {:?}",
                    state.manager.config().allowed_domains
                )
            })),
        )
            .into_response();
    }

    let identity = state.manager.build_identity(
        &format!("apikey-{}", &body.api_key[..8.min(body.api_key.len())]),
        &email,
        Some("API Key User"),
        vec!["api-user".into()],
        None,
    );

    let token = state.manager.create_session(identity.clone());

    let response = AuthResponse {
        expires_at: identity.expires_at,
        token,
        identity,
    };

    (StatusCode::OK, Json(response)).into_response()
}

// ─── Helper: extract session token ──────────────────────────────────────────

/// Extract the session token from the request.
///
/// Checks (in order):
/// 1. `Authorization: Bearer <token>` header
/// 2. `argentor_session` cookie
fn extract_session_token(headers: &HeaderMap) -> Option<String> {
    // Check Authorization header
    if let Some(auth) = headers
        .get("authorization")
        .and_then(|v| v.to_str().ok())
        .and_then(|v| v.strip_prefix("Bearer "))
    {
        return Some(auth.to_string());
    }

    // Check cookie
    if let Some(cookie_header) = headers.get("cookie").and_then(|v| v.to_str().ok()) {
        for cookie in cookie_header.split(';') {
            let cookie = cookie.trim();
            if let Some(value) = cookie.strip_prefix("argentor_session=") {
                return Some(value.to_string());
            }
        }
    }

    None
}

// ─── SSO middleware ─────────────────────────────────────────────────────────

/// Shared state for the SSO auth middleware.
#[derive(Clone)]
pub struct SsoMiddlewareState {
    /// The SSO manager for session validation.
    pub manager: Arc<SsoManager>,
}

/// Axum middleware that authenticates requests via SSO sessions.
///
/// Extracts the session token from the `Authorization: Bearer` header or
/// `argentor_session` cookie, validates it against the SSO manager, and
/// injects the [`UserIdentity`] into request extensions.
///
/// Returns `401 Unauthorized` if no valid session is found.
pub async fn sso_auth_middleware(
    State(state): State<Arc<SsoMiddlewareState>>,
    headers: HeaderMap,
    mut request: Request,
    next: Next,
) -> Response {
    let token = extract_session_token(&headers);

    match token {
        Some(t) => match state.manager.validate_session(&t) {
            Some(identity) => {
                request.extensions_mut().insert(identity);
                next.run(request).await
            }
            None => {
                warn!("SSO session token invalid or expired");
                (
                    StatusCode::UNAUTHORIZED,
                    Json(serde_json::json!({
                        "error": "session_expired",
                        "message": "Session token is invalid or expired — please log in again"
                    })),
                )
                    .into_response()
            }
        },
        None => {
            warn!("No SSO session token in request");
            (
                StatusCode::UNAUTHORIZED,
                Json(serde_json::json!({
                    "error": "not_authenticated",
                    "message": "Authentication required — include a Bearer token or argentor_session cookie"
                })),
            )
                .into_response()
        }
    }
}

// ─── Router builder ─────────────────────────────────────────────────────────

/// Build the SSO authentication router.
///
/// Mounts the following routes:
/// - `GET  /auth/login`    — Redirect to SSO provider
/// - `GET  /auth/callback` — Handle SSO callback
/// - `GET  /auth/logout`   — Revoke session
/// - `GET  /auth/me`       — Current user identity
/// - `POST /auth/api-key`  — API key to session exchange
pub fn sso_router(manager: Arc<SsoManager>) -> Router {
    let state = Arc::new(SsoState { manager });

    Router::new()
        .route("/auth/login", get(sso_login_handler))
        .route("/auth/callback", get(sso_callback_handler))
        .route("/auth/logout", get(sso_logout_handler))
        .route("/auth/me", get(sso_me_handler))
        .route("/auth/api-key", axum::routing::post(sso_api_key_handler))
        .with_state(state)
}

// ─── Tests ──────────────────────────────────────────────────────────────────

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
    use super::*;
    use axum::body::Body;
    use axum::http::Request as HttpRequest;
    use tower::ServiceExt;

    /// Helper: create a test SSO config.
    fn test_config() -> SsoConfig {
        SsoConfig {
            provider: SsoProvider::Oidc,
            client_id: "test-client-id".into(),
            client_secret: "test-client-secret".into(),
            redirect_uri: "https://app.example.com/auth/callback".into(),
            issuer_url: "https://accounts.example.com".into(),
            allowed_domains: vec!["example.com".into(), "corp.example.com".into()],
            scopes: vec!["openid".into(), "profile".into(), "email".into()],
            session_ttl_hours: 24,
        }
    }

    /// Helper: create a test SsoManager.
    fn test_manager() -> SsoManager {
        SsoManager::new(test_config())
    }

    /// Helper: create a test UserIdentity that is valid (not expired).
    fn test_identity(email: &str) -> UserIdentity {
        let domain = email
            .rsplit_once('@')
            .map(|(_, d)| d.to_string())
            .unwrap_or_default();
        UserIdentity {
            id: Uuid::new_v4().to_string(),
            email: email.to_string(),
            name: Some("Test User".into()),
            domain,
            roles: vec!["user".into()],
            tenant_id: Some("tenant-1".into()),
            authenticated_at: Utc::now(),
            expires_at: Utc::now() + Duration::hours(24),
        }
    }

    /// Helper: create an expired UserIdentity.
    fn expired_identity(email: &str) -> UserIdentity {
        let domain = email
            .rsplit_once('@')
            .map(|(_, d)| d.to_string())
            .unwrap_or_default();
        UserIdentity {
            id: Uuid::new_v4().to_string(),
            email: email.to_string(),
            name: Some("Expired User".into()),
            domain,
            roles: vec!["user".into()],
            tenant_id: None,
            authenticated_at: Utc::now() - Duration::hours(48),
            expires_at: Utc::now() - Duration::hours(1),
        }
    }

    // ── Config tests ────────────────────────────────────────────────────

    #[test]
    fn sso_config_default_values() {
        let config = SsoConfig::default();
        assert_eq!(config.provider, SsoProvider::Oidc);
        assert_eq!(config.session_ttl_hours, 24);
        assert!(config.scopes.contains(&"openid".to_string()));
        assert!(config.allowed_domains.is_empty());
    }

    #[test]
    fn sso_config_serialization_roundtrip() {
        let config = test_config();
        let json = serde_json::to_string(&config).unwrap();
        let deserialized: SsoConfig = serde_json::from_str(&json).unwrap();
        assert_eq!(deserialized.client_id, "test-client-id");
        assert_eq!(deserialized.allowed_domains.len(), 2);
        assert_eq!(deserialized.provider, SsoProvider::Oidc);
    }

    #[test]
    fn sso_provider_variants_serialize() {
        let oidc = serde_json::to_string(&SsoProvider::Oidc).unwrap();
        let saml = serde_json::to_string(&SsoProvider::Saml).unwrap();
        let api = serde_json::to_string(&SsoProvider::ApiKey).unwrap();
        assert_eq!(oidc, "\"Oidc\"");
        assert_eq!(saml, "\"Saml\"");
        assert_eq!(api, "\"ApiKey\"");
    }

    // ── Login URL tests ─────────────────────────────────────────────────

    #[test]
    fn login_url_oidc_contains_required_params() {
        let manager = test_manager();
        let url = manager.login_url("csrf-state-123");
        assert!(url.contains("accounts.example.com/authorize"));
        assert!(url.contains("client_id=test-client-id"));
        assert!(url.contains("redirect_uri=https://app.example.com/auth/callback"));
        assert!(url.contains("response_type=code"));
        assert!(url.contains("state=csrf-state-123"));
        assert!(url.contains("scope=openid+profile+email"));
    }

    #[test]
    fn login_url_saml_contains_relay_state() {
        let mut config = test_config();
        config.provider = SsoProvider::Saml;
        let manager = SsoManager::new(config);
        let url = manager.login_url("saml-state-456");
        assert!(url.contains("SAMLRequest="));
        assert!(url.contains("RelayState=saml-state-456"));
    }

    #[test]
    fn login_url_api_key_uses_redirect_uri() {
        let mut config = test_config();
        config.provider = SsoProvider::ApiKey;
        let manager = SsoManager::new(config);
        let url = manager.login_url("api-state");
        assert!(url.contains("app.example.com/auth/callback"));
        assert!(url.contains("mode=api_key"));
    }

    #[test]
    fn login_url_records_pending_state() {
        let manager = test_manager();
        let _url = manager.login_url("track-me");
        let states = manager.pending_states.read().unwrap();
        assert!(states.contains_key("track-me"));
    }

    // ── Domain validation tests ─────────────────────────────────────────

    #[test]
    fn domain_allowed_exact_match() {
        let manager = test_manager();
        assert!(manager.is_domain_allowed("alice@example.com"));
        assert!(manager.is_domain_allowed("bob@corp.example.com"));
    }

    #[test]
    fn domain_blocked_when_not_in_list() {
        let manager = test_manager();
        assert!(!manager.is_domain_allowed("eve@evil.com"));
        assert!(!manager.is_domain_allowed("mallory@other.org"));
    }

    #[test]
    fn domain_case_insensitive() {
        let manager = test_manager();
        assert!(manager.is_domain_allowed("Alice@EXAMPLE.COM"));
        assert!(manager.is_domain_allowed("bob@Corp.Example.Com"));
    }

    #[test]
    fn domain_all_allowed_when_list_empty() {
        let mut config = test_config();
        config.allowed_domains = vec![];
        let manager = SsoManager::new(config);
        assert!(manager.is_domain_allowed("anyone@anywhere.com"));
    }

    #[test]
    fn domain_invalid_email_rejected() {
        let manager = test_manager();
        assert!(!manager.is_domain_allowed("not-an-email"));
        assert!(!manager.is_domain_allowed(""));
    }

    // ── Session lifecycle tests ─────────────────────────────────────────

    #[test]
    fn session_create_and_validate() {
        let manager = test_manager();
        let identity = test_identity("alice@example.com");
        let token = manager.create_session(identity.clone());

        let retrieved = manager.validate_session(&token);
        assert!(retrieved.is_some());
        let retrieved = retrieved.unwrap();
        assert_eq!(retrieved.email, "alice@example.com");
        assert_eq!(retrieved.name, Some("Test User".into()));
    }

    #[test]
    fn session_validate_unknown_token_returns_none() {
        let manager = test_manager();
        assert!(manager.validate_session("nonexistent-token").is_none());
    }

    #[test]
    fn session_revoke_removes_session() {
        let manager = test_manager();
        let identity = test_identity("bob@example.com");
        let token = manager.create_session(identity);

        assert!(manager.validate_session(&token).is_some());
        assert!(manager.revoke_session(&token));
        assert!(manager.validate_session(&token).is_none());
    }

    #[test]
    fn session_revoke_unknown_returns_false() {
        let manager = test_manager();
        assert!(!manager.revoke_session("no-such-token"));
    }

    #[test]
    fn session_expired_auto_revoked_on_validate() {
        let manager = test_manager();
        let identity = expired_identity("expired@example.com");
        let token = manager.create_session(identity);

        // Should return None because the session is expired.
        assert!(manager.validate_session(&token).is_none());
        // Session should have been auto-revoked.
        let sessions = manager.sessions.read().unwrap();
        assert!(!sessions.contains_key(&token));
    }

    // ── Cleanup tests ───────────────────────────────────────────────────

    #[test]
    fn cleanup_removes_expired_sessions() {
        let manager = test_manager();

        // Add one valid and two expired sessions.
        let _valid_token = manager.create_session(test_identity("valid@example.com"));
        let _expired1 = manager.create_session(expired_identity("old1@example.com"));
        let _expired2 = manager.create_session(expired_identity("old2@example.com"));

        let removed = manager.cleanup_expired();
        assert_eq!(removed, 2);

        let active = manager.active_sessions();
        assert_eq!(active.len(), 1);
        assert_eq!(active[0].1.email, "valid@example.com");
    }

    #[test]
    fn cleanup_returns_zero_when_none_expired() {
        let manager = test_manager();
        let _t = manager.create_session(test_identity("fresh@example.com"));
        assert_eq!(manager.cleanup_expired(), 0);
    }

    // ── Active sessions listing ─────────────────────────────────────────

    #[test]
    fn active_sessions_lists_only_valid() {
        let manager = test_manager();
        let _t1 = manager.create_session(test_identity("a@example.com"));
        let _t2 = manager.create_session(test_identity("b@example.com"));
        let _t3 = manager.create_session(expired_identity("c@example.com"));

        let active = manager.active_sessions();
        assert_eq!(active.len(), 2);

        let emails: Vec<&str> = active.iter().map(|(_, id)| id.email.as_str()).collect();
        assert!(emails.contains(&"a@example.com"));
        assert!(emails.contains(&"b@example.com"));
    }

    // ── Build identity helper test ──────────────────────────────────────

    #[test]
    fn build_identity_sets_domain_and_ttl() {
        let manager = test_manager();
        let identity = manager.build_identity(
            "user-1",
            "alice@example.com",
            Some("Alice"),
            vec!["admin".into()],
            Some("tenant-x".into()),
        );

        assert_eq!(identity.id, "user-1");
        assert_eq!(identity.email, "alice@example.com");
        assert_eq!(identity.domain, "example.com");
        assert_eq!(identity.name, Some("Alice".into()));
        assert_eq!(identity.roles, vec!["admin"]);
        assert_eq!(identity.tenant_id, Some("tenant-x".into()));
        assert!(identity.expires_at > Utc::now());
        // TTL should be ~24 hours
        let ttl = identity.expires_at - identity.authenticated_at;
        assert_eq!(ttl.num_hours(), 24);
    }

    // ── UserIdentity expiration test ────────────────────────────────────

    #[test]
    fn user_identity_is_expired() {
        let valid = test_identity("valid@example.com");
        assert!(!valid.is_expired());

        let expired = expired_identity("old@example.com");
        assert!(expired.is_expired());
    }

    // ── Middleware tests ─────────────────────────────────────────────────

    #[tokio::test]
    async fn middleware_valid_session_passes_through() {
        let manager = Arc::new(test_manager());
        let identity = test_identity("auth@example.com");
        let token = manager.create_session(identity);

        let mw_state = Arc::new(SsoMiddlewareState {
            manager: manager.clone(),
        });

        // Build a minimal app with the middleware.
        let app = Router::new()
            .route(
                "/protected",
                get(|req: HttpRequest<Body>| async move {
                    let id = req.extensions().get::<UserIdentity>();
                    match id {
                        Some(u) => (StatusCode::OK, u.email.clone()).into_response(),
                        None => (StatusCode::INTERNAL_SERVER_ERROR, "no identity").into_response(),
                    }
                }),
            )
            .layer(axum::middleware::from_fn_with_state(
                mw_state,
                sso_auth_middleware,
            ));

        let response = app
            .oneshot(
                HttpRequest::builder()
                    .uri("/protected")
                    .header("authorization", format!("Bearer {token}"))
                    .body(Body::empty())
                    .unwrap(),
            )
            .await
            .unwrap();

        assert_eq!(response.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn middleware_missing_session_returns_401() {
        let manager = Arc::new(test_manager());
        let mw_state = Arc::new(SsoMiddlewareState {
            manager: manager.clone(),
        });

        let app = Router::new()
            .route("/protected", get(|| async { "ok" }))
            .layer(axum::middleware::from_fn_with_state(
                mw_state,
                sso_auth_middleware,
            ));

        let response = app
            .oneshot(
                HttpRequest::builder()
                    .uri("/protected")
                    .body(Body::empty())
                    .unwrap(),
            )
            .await
            .unwrap();

        assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
    }

    #[tokio::test]
    async fn middleware_expired_session_returns_401() {
        let manager = Arc::new(test_manager());
        let identity = expired_identity("old@example.com");

        // Manually insert the expired session (bypassing create_session validation).
        let token = "expired-token-123".to_string();
        manager
            .sessions
            .write()
            .unwrap()
            .insert(token.clone(), identity);

        let mw_state = Arc::new(SsoMiddlewareState {
            manager: manager.clone(),
        });

        let app = Router::new()
            .route("/protected", get(|| async { "ok" }))
            .layer(axum::middleware::from_fn_with_state(
                mw_state,
                sso_auth_middleware,
            ));

        let response = app
            .oneshot(
                HttpRequest::builder()
                    .uri("/protected")
                    .header("authorization", format!("Bearer {token}"))
                    .body(Body::empty())
                    .unwrap(),
            )
            .await
            .unwrap();

        assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
    }

    #[tokio::test]
    async fn middleware_invalid_token_returns_401() {
        let manager = Arc::new(test_manager());
        let mw_state = Arc::new(SsoMiddlewareState {
            manager: manager.clone(),
        });

        let app = Router::new()
            .route("/protected", get(|| async { "ok" }))
            .layer(axum::middleware::from_fn_with_state(
                mw_state,
                sso_auth_middleware,
            ));

        let response = app
            .oneshot(
                HttpRequest::builder()
                    .uri("/protected")
                    .header("authorization", "Bearer bogus-token")
                    .body(Body::empty())
                    .unwrap(),
            )
            .await
            .unwrap();

        assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
    }

    // ── Token extraction tests ──────────────────────────────────────────

    #[test]
    fn extract_token_from_bearer_header() {
        let mut headers = HeaderMap::new();
        headers.insert("authorization", "Bearer my-token-123".parse().unwrap());
        assert_eq!(extract_session_token(&headers), Some("my-token-123".into()));
    }

    #[test]
    fn extract_token_from_cookie() {
        let mut headers = HeaderMap::new();
        headers.insert(
            "cookie",
            "other=val; argentor_session=cookie-token-456; another=x"
                .parse()
                .unwrap(),
        );
        assert_eq!(
            extract_session_token(&headers),
            Some("cookie-token-456".into())
        );
    }

    #[test]
    fn extract_token_bearer_takes_precedence_over_cookie() {
        let mut headers = HeaderMap::new();
        headers.insert("authorization", "Bearer header-token".parse().unwrap());
        headers.insert("cookie", "argentor_session=cookie-token".parse().unwrap());
        assert_eq!(extract_session_token(&headers), Some("header-token".into()));
    }

    #[test]
    fn extract_token_returns_none_when_absent() {
        let headers = HeaderMap::new();
        assert_eq!(extract_session_token(&headers), None);
    }

    // ── SSO callback tests ────────────────────────────────────────────

    #[tokio::test]
    async fn callback_rejects_invalid_state() {
        let manager = test_manager();
        let result = manager.handle_callback("some-code", "unknown-state").await;
        assert!(result.is_err());
        let err = result.unwrap_err().to_string();
        assert!(err.contains("CSRF"), "Error should mention CSRF: {err}");
    }

    #[tokio::test]
    async fn callback_rejects_empty_code() {
        let manager = test_manager();
        // First create a valid state
        let _url = manager.login_url("valid-state");
        let result = manager.handle_callback("", "valid-state").await;
        assert!(result.is_err());
        let err = result.unwrap_err().to_string();
        assert!(
            err.contains("Empty authorization code"),
            "Error should mention empty code: {err}"
        );
    }

    #[tokio::test]
    async fn callback_oidc_discovery_fails_on_unreachable_issuer() {
        // With a real OIDC implementation, pointing at a non-existent issuer
        // should fail during discovery (HTTP error), not with "not configured".
        let manager = test_manager();
        let _url = manager.login_url("oidc-state");
        let result = manager.handle_callback("auth-code-123", "oidc-state").await;
        assert!(result.is_err());
        let err = result.unwrap_err().to_string();
        // The error should come from the HTTP layer (discovery failed),
        // not a placeholder "not configured" message.
        assert!(
            err.contains("OIDC discovery") || err.contains("HTTP"),
            "Should indicate OIDC discovery/HTTP failure: {err}"
        );
    }

    #[tokio::test]
    async fn callback_saml_invalid_base64_rejected() {
        let mut config = test_config();
        config.provider = SsoProvider::Saml;
        let manager = SsoManager::new(config);
        let _url = manager.login_url("saml-state");
        // "saml-response" is not valid base64
        let result = manager.handle_callback("saml-response", "saml-state").await;
        assert!(result.is_err());
        let err = result.unwrap_err().to_string();
        assert!(
            err.contains("base64 decode failed"),
            "Should indicate base64 decode failure: {err}"
        );
    }

    // ── SAML response parsing tests ─────────────────────────────────────

    /// Helper: build a realistic SAML XML response.
    fn make_saml_xml(name_id: &str, status: &str, attrs: &[(&str, &[&str])]) -> String {
        let mut attr_xml = String::new();
        for (name, values) in attrs {
            attr_xml.push_str(&format!(r#"<saml:Attribute Name="{name}">"#));
            for val in *values {
                attr_xml.push_str(&format!("<saml:AttributeValue>{val}</saml:AttributeValue>"));
            }
            attr_xml.push_str("</saml:Attribute>\n");
        }

        format!(
            r#"<samlp:Response xmlns:samlp="urn:oasis:names:tc:SAML:2.0:protocol" xmlns:saml="urn:oasis:names:tc:SAML:2.0:assertion" ID="_response123" Version="2.0" IssueInstant="2025-01-15T10:00:00Z" Destination="https://app.example.com/auth/callback">
  <saml:Issuer>https://idp.example.com</saml:Issuer>
  <samlp:Status>
    <samlp:StatusCode Value="{status}"/>
  </samlp:Status>
  <saml:Assertion ID="_assertion456" Version="2.0" IssueInstant="2025-01-15T10:00:00Z">
    <saml:Issuer>https://idp.example.com</saml:Issuer>
    <saml:Subject>
      <saml:NameID Format="urn:oasis:names:tc:SAML:1.1:nameid-format:emailAddress">{name_id}</saml:NameID>
    </saml:Subject>
    <saml:Conditions NotBefore="2025-01-15T09:55:00Z" NotOnOrAfter="2025-01-15T10:05:00Z">
      <saml:AudienceRestriction>
        <saml:Audience>https://app.example.com</saml:Audience>
      </saml:AudienceRestriction>
    </saml:Conditions>
    <saml:AttributeStatement>
      {attr_xml}
    </saml:AttributeStatement>
  </saml:Assertion>
</samlp:Response>"#
        )
    }

    /// Helper: base64-encode a SAML XML string.
    fn encode_saml(xml: &str) -> String {
        STANDARD.encode(xml.as_bytes())
    }

    #[test]
    fn saml_parse_valid_response() {
        let xml = make_saml_xml(
            "alice@example.com",
            "urn:oasis:names:tc:SAML:2.0:status:Success",
            &[
                ("name", &["Alice Smith"]),
                ("email", &["alice@example.com"]),
                ("role", &["admin", "viewer"]),
            ],
        );
        let encoded = encode_saml(&xml);
        let claims = parse_saml_response(&encoded).unwrap();

        assert_eq!(claims.name_id, "alice@example.com");
        assert_eq!(claims.email, "alice@example.com");
        assert_eq!(claims.name, Some("Alice Smith".into()));
        assert_eq!(claims.roles, vec!["admin", "viewer"]);
    }

    #[test]
    fn saml_parse_extracts_name_id() {
        let xml = make_saml_xml(
            "bob@corp.example.com",
            "urn:oasis:names:tc:SAML:2.0:status:Success",
            &[],
        );
        let encoded = encode_saml(&xml);
        let claims = parse_saml_response(&encoded).unwrap();

        assert_eq!(claims.name_id, "bob@corp.example.com");
        // With no email attribute, email falls back to name_id
        assert_eq!(claims.email, "bob@corp.example.com");
    }

    #[test]
    fn saml_parse_extracts_attributes() {
        let xml = make_saml_xml(
            "user123",
            "urn:oasis:names:tc:SAML:2.0:status:Success",
            &[
                ("name", &["Charlie Brown"]),
                ("email", &["charlie@example.com"]),
            ],
        );
        let encoded = encode_saml(&xml);
        let claims = parse_saml_response(&encoded).unwrap();

        assert_eq!(claims.name, Some("Charlie Brown".into()));
        assert_eq!(claims.email, "charlie@example.com");
    }

    #[test]
    fn saml_parse_missing_name_id_rejected() {
        // XML with no NameID element
        let xml = r#"<samlp:Response xmlns:samlp="urn:oasis:names:tc:SAML:2.0:protocol" xmlns:saml="urn:oasis:names:tc:SAML:2.0:assertion">
  <samlp:Status>
    <samlp:StatusCode Value="urn:oasis:names:tc:SAML:2.0:status:Success"/>
  </samlp:Status>
  <saml:Assertion>
    <saml:Subject></saml:Subject>
  </saml:Assertion>
</samlp:Response>"#;
        let encoded = encode_saml(xml);
        let result = parse_saml_response(&encoded);
        assert!(result.is_err());
        let err = result.unwrap_err().to_string();
        assert!(
            err.contains("missing NameID"),
            "Should mention missing NameID: {err}"
        );
    }

    #[test]
    fn saml_parse_failed_status_code_rejected() {
        let xml = make_saml_xml(
            "alice@example.com",
            "urn:oasis:names:tc:SAML:2.0:status:Requester",
            &[],
        );
        let encoded = encode_saml(&xml);
        let result = parse_saml_response(&encoded);
        assert!(result.is_err());
        let err = result.unwrap_err().to_string();
        assert!(
            err.contains("SAML authentication failed"),
            "Should indicate SAML auth failure: {err}"
        );
        assert!(
            err.contains("Requester"),
            "Should include the status code: {err}"
        );
    }

    #[tokio::test]
    async fn saml_callback_domain_validation() {
        let mut config = test_config();
        config.provider = SsoProvider::Saml;
        // allowed_domains: ["example.com", "corp.example.com"]
        let manager = SsoManager::new(config);
        let _url = manager.login_url("saml-state");

        // Build a SAML response with an email from an unauthorized domain
        let xml = make_saml_xml(
            "evil@attacker.com",
            "urn:oasis:names:tc:SAML:2.0:status:Success",
            &[("email", &["evil@attacker.com"])],
        );
        let encoded = encode_saml(&xml);

        let result = manager.handle_callback(&encoded, "saml-state").await;
        assert!(result.is_err());
        let err = result.unwrap_err().to_string();
        assert!(
            err.contains("not in allowed list"),
            "Should reject unauthorized domain: {err}"
        );
    }

    #[test]
    fn saml_parse_multiple_roles() {
        let xml = make_saml_xml(
            "admin@example.com",
            "urn:oasis:names:tc:SAML:2.0:status:Success",
            &[("role", &["admin", "editor", "viewer", "auditor"])],
        );
        let encoded = encode_saml(&xml);
        let claims = parse_saml_response(&encoded).unwrap();

        assert_eq!(claims.roles.len(), 4);
        assert_eq!(claims.roles, vec!["admin", "editor", "viewer", "auditor"]);
    }

    #[test]
    fn saml_parse_azure_ad_uri_attributes() {
        // Azure AD / ADFS uses full URI claim names
        let xml = make_saml_xml(
            "azure-user-id",
            "urn:oasis:names:tc:SAML:2.0:status:Success",
            &[
                (
                    "http://schemas.xmlsoap.org/ws/2005/05/identity/claims/name",
                    &["Azure User"],
                ),
                (
                    "http://schemas.xmlsoap.org/ws/2005/05/identity/claims/emailaddress",
                    &["azure@example.com"],
                ),
                (
                    "http://schemas.microsoft.com/ws/2008/06/identity/claims/role",
                    &["GlobalAdmin", "Reader"],
                ),
            ],
        );
        let encoded = encode_saml(&xml);
        let claims = parse_saml_response(&encoded).unwrap();

        assert_eq!(claims.name_id, "azure-user-id");
        assert_eq!(claims.name, Some("Azure User".into()));
        assert_eq!(claims.email, "azure@example.com");
        assert_eq!(claims.roles, vec!["GlobalAdmin", "Reader"]);
    }

    #[tokio::test]
    async fn saml_callback_full_success_flow() {
        let mut config = test_config();
        config.provider = SsoProvider::Saml;
        let manager = SsoManager::new(config);
        let _url = manager.login_url("saml-ok-state");

        let xml = make_saml_xml(
            "alice@example.com",
            "urn:oasis:names:tc:SAML:2.0:status:Success",
            &[
                ("name", &["Alice Smith"]),
                ("email", &["alice@example.com"]),
                ("role", &["engineer"]),
            ],
        );
        let encoded = encode_saml(&xml);

        let result = manager.handle_callback(&encoded, "saml-ok-state").await;
        assert!(result.is_ok(), "SAML callback should succeed");

        let (token, identity) = result.unwrap();
        assert!(!token.is_empty());
        assert_eq!(identity.email, "alice@example.com");
        assert_eq!(identity.name, Some("Alice Smith".into()));
        assert_eq!(identity.domain, "example.com");
        assert_eq!(identity.roles, vec!["engineer"]);

        // Session should be valid
        let session = manager.validate_session(&token);
        assert!(session.is_some());
    }

    #[test]
    fn saml_parse_displayname_fallback() {
        let xml = make_saml_xml(
            "user@example.com",
            "urn:oasis:names:tc:SAML:2.0:status:Success",
            &[("displayName", &["Display Name User"])],
        );
        let encoded = encode_saml(&xml);
        let claims = parse_saml_response(&encoded).unwrap();
        assert_eq!(claims.name, Some("Display Name User".into()));
    }

    #[test]
    fn saml_parse_invalid_base64() {
        let result = parse_saml_response("!!!not-base64!!!");
        assert!(result.is_err());
        let err = result.unwrap_err().to_string();
        assert!(
            err.contains("base64 decode failed"),
            "Should mention base64 failure: {err}"
        );
    }

    // ── SAML XML helper unit tests ──────────────────────────────────────

    #[test]
    fn extract_xml_element_simple() {
        let xml = "<NameID>alice@example.com</NameID>";
        assert_eq!(
            extract_xml_element(xml, "NameID"),
            Some("alice@example.com".into())
        );
    }

    #[test]
    fn extract_xml_element_with_namespace() {
        let xml = r#"<saml:NameID Format="email">alice@example.com</saml:NameID>"#;
        assert_eq!(
            extract_xml_element(xml, "NameID"),
            Some("alice@example.com".into())
        );
    }

    #[test]
    fn extract_xml_element_not_found() {
        let xml = "<Issuer>https://idp.example.com</Issuer>";
        assert_eq!(extract_xml_element(xml, "NameID"), None);
    }

    #[test]
    fn extract_xml_attr_status_code() {
        let xml = r#"<samlp:StatusCode Value="urn:oasis:names:tc:SAML:2.0:status:Success"/>"#;
        let val = extract_xml_attr(xml, "StatusCode", "Value");
        assert_eq!(
            val,
            Some("urn:oasis:names:tc:SAML:2.0:status:Success".into())
        );
    }

    #[test]
    fn extract_saml_attribute_simple() {
        let xml = r#"<saml:Attribute Name="email"><saml:AttributeValue>alice@example.com</saml:AttributeValue></saml:Attribute>"#;
        assert_eq!(
            extract_saml_attribute(xml, "email"),
            Some("alice@example.com".into())
        );
    }

    #[test]
    fn extract_saml_attribute_values_multiple() {
        let xml = r#"<saml:Attribute Name="role"><saml:AttributeValue>admin</saml:AttributeValue><saml:AttributeValue>user</saml:AttributeValue></saml:Attribute>"#;
        let values = extract_saml_attribute_values(xml, "role");
        assert_eq!(values, Some(vec!["admin".into(), "user".into()]));
    }

    // ── JWT decode tests ────────────────────────────────────────────────

    /// Helper: build a minimal JWT with the given payload (no real signature).
    fn make_test_jwt(payload: &serde_json::Value) -> String {
        use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};

        let header = serde_json::json!({"alg": "RS256", "typ": "JWT"});
        let header_b64 = URL_SAFE_NO_PAD.encode(header.to_string().as_bytes());
        let payload_b64 = URL_SAFE_NO_PAD.encode(payload.to_string().as_bytes());
        let sig_b64 = URL_SAFE_NO_PAD.encode(b"fake-signature");
        format!("{header_b64}.{payload_b64}.{sig_b64}")
    }

    #[test]
    fn decode_jwt_payload_valid_token() {
        let claims = serde_json::json!({
            "sub": "user-123",
            "email": "alice@example.com",
            "name": "Alice",
            "iss": "https://accounts.example.com",
            "email_verified": true
        });
        let token = make_test_jwt(&claims);
        let decoded = decode_jwt_payload(&token).unwrap();
        assert_eq!(decoded["sub"].as_str().unwrap(), "user-123");
        assert_eq!(decoded["email"].as_str().unwrap(), "alice@example.com");
        assert_eq!(decoded["name"].as_str().unwrap(), "Alice");
        assert!(decoded["email_verified"].as_bool().unwrap());
    }

    #[test]
    fn decode_jwt_payload_invalid_format_no_dots() {
        let result = decode_jwt_payload("not-a-jwt");
        assert!(result.is_err());
        let err = result.unwrap_err().to_string();
        assert!(
            err.contains("Invalid JWT format"),
            "Should mention invalid format: {err}"
        );
    }

    #[test]
    fn decode_jwt_payload_invalid_format_two_parts() {
        let result = decode_jwt_payload("part1.part2");
        assert!(result.is_err());
        let err = result.unwrap_err().to_string();
        assert!(
            err.contains("expected 3 parts"),
            "Should mention expected 3 parts: {err}"
        );
    }

    #[test]
    fn decode_jwt_payload_invalid_base64() {
        let result = decode_jwt_payload("header.!!!invalid!!!.signature");
        assert!(result.is_err());
        let err = result.unwrap_err().to_string();
        assert!(
            err.contains("base64url decode failed"),
            "Should mention base64 decode failure: {err}"
        );
    }

    #[test]
    fn decode_jwt_payload_invalid_json() {
        use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
        let header = URL_SAFE_NO_PAD.encode(b"{}");
        let payload = URL_SAFE_NO_PAD.encode(b"not json at all");
        let sig = URL_SAFE_NO_PAD.encode(b"sig");
        let token = format!("{header}.{payload}.{sig}");
        let result = decode_jwt_payload(&token);
        assert!(result.is_err());
        let err = result.unwrap_err().to_string();
        assert!(
            err.contains("not valid JSON"),
            "Should mention invalid JSON: {err}"
        );
    }

    #[test]
    fn decode_jwt_payload_extracts_nested_claims() {
        let claims = serde_json::json!({
            "sub": "u-456",
            "email": "bob@corp.example.com",
            "name": "Bob",
            "iss": "https://idp.example.com",
            "aud": "test-client-id",
            "email_verified": true,
            "custom_claim": {"org": "acme", "tier": "enterprise"}
        });
        let token = make_test_jwt(&claims);
        let decoded = decode_jwt_payload(&token).unwrap();
        assert_eq!(decoded["custom_claim"]["org"].as_str().unwrap(), "acme");
    }

    // ── OIDC domain validation after token decode ───────────────────────

    #[test]
    fn oidc_domain_validation_after_decode() {
        // Simulate what handle_callback does after decoding the JWT:
        // check that the email domain is in the allowed list.
        let manager = test_manager();
        // example.com is allowed
        assert!(manager.is_domain_allowed("alice@example.com"));
        // evil.com is not allowed
        assert!(!manager.is_domain_allowed("alice@evil.com"));
    }

    // ── OIDC endpoints struct test ──────────────────────────────────────

    #[test]
    fn oidc_endpoints_struct_fields() {
        let endpoints = OidcEndpoints {
            authorization_endpoint: "https://idp.example.com/authorize".into(),
            token_endpoint: "https://idp.example.com/token".into(),
            userinfo_endpoint: Some("https://idp.example.com/userinfo".into()),
            issuer: "https://idp.example.com".into(),
        };
        assert_eq!(endpoints.token_endpoint, "https://idp.example.com/token");
        assert_eq!(endpoints.issuer, "https://idp.example.com");
        assert!(endpoints.userinfo_endpoint.is_some());
    }

    #[test]
    fn oidc_endpoints_optional_userinfo() {
        let endpoints = OidcEndpoints {
            authorization_endpoint: "https://idp.example.com/authorize".into(),
            token_endpoint: "https://idp.example.com/token".into(),
            userinfo_endpoint: None,
            issuer: "https://idp.example.com".into(),
        };
        assert!(endpoints.userinfo_endpoint.is_none());
    }

    // ── OIDC issuer validation ──────────────────────────────────────────

    #[test]
    fn issuer_validation_trailing_slash_normalization() {
        // Both "https://accounts.example.com" and "https://accounts.example.com/"
        // should be considered the same issuer when trimmed.
        let iss_a = "https://accounts.example.com/";
        let iss_b = "https://accounts.example.com";
        assert_eq!(iss_a.trim_end_matches('/'), iss_b.trim_end_matches('/'));
    }

    // ── OIDC email_verified claim handling ──────────────────────────────

    #[test]
    fn email_verified_false_is_rejected() {
        // Simulate the check that handle_callback performs.
        let payload = serde_json::json!({
            "sub": "user-1",
            "email": "alice@example.com",
            "email_verified": false,
            "iss": "https://accounts.example.com"
        });
        let email_verified = payload["email_verified"].as_bool().unwrap_or(false);
        assert!(!email_verified, "email_verified=false should be rejected");
    }

    #[test]
    fn email_verified_missing_is_rejected() {
        // If the claim is missing entirely, we default to false.
        let payload = serde_json::json!({
            "sub": "user-1",
            "email": "alice@example.com",
            "iss": "https://accounts.example.com"
        });
        let email_verified = payload["email_verified"].as_bool().unwrap_or(false);
        assert!(
            !email_verified,
            "missing email_verified should default to false"
        );
    }

    // ── OIDC missing id_token handling ──────────────────────────────────

    #[test]
    fn missing_id_token_in_response_detected() {
        // Simulate the check: token_json["id_token"].as_str() returns None
        // when the field is missing from the token response.
        let token_json = serde_json::json!({
            "access_token": "abc123",
            "token_type": "Bearer"
        });
        assert!(
            token_json["id_token"].as_str().is_none(),
            "Missing id_token should be detected"
        );
    }
}