// ═══════════════════════════════════════════════════════════════════════
//  Auth Module — Pluggable authentication for PLP connections
//
//  Supports three auth modes:
//    • None    — No authentication (dev/local mode, default)
//    • JWT     — Verify JWTs using a JWKS endpoint (Clerk, Auth0, Google, etc.)
//    • Webhook — POST the token to an external URL for verification
// ═══════════════════════════════════════════════════════════════════════

use pulse_core::control::{AuthContext, AuthenticatedUser};
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::{debug, info, warn};

// ─── Auth result ─────────────────────────────────────────────────────

/// Outcome of verifying an auth token.
#[derive(Debug, Clone)]
pub enum AuthResult {
    /// Authentication succeeded; the user is identified.
    Authenticated(AuthenticatedUser),
    /// Authentication was not required (auth mode = none).
    Skipped,
    /// Authentication failed with a reason.
    Rejected(String),
}

// ─── Auth configuration ──────────────────────────────────────────────

/// Which auth mode the relay is operating in.
#[derive(Debug, Clone, PartialEq, Eq, clap::ValueEnum, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum AuthMode {
    /// No authentication — all connections accepted (default, dev mode).
    None,
    /// Verify JWTs locally using a JWKS endpoint.
    Jwt,
    /// Delegate verification to an external webhook URL.
    Webhook,
}

impl std::fmt::Display for AuthMode {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::None => write!(f, "none"),
            Self::Jwt => write!(f, "jwt"),
            Self::Webhook => write!(f, "webhook"),
        }
    }
}

/// Full auth configuration for the relay.
#[derive(Debug, Clone)]
pub struct AuthConfig {
    pub mode: AuthMode,
    /// JWKS URL for JWT mode (e.g. https://clerk.your-app.com/.well-known/jwks.json).
    pub jwt_issuer: Option<String>,
    /// Expected audience for JWT verification (optional).
    pub jwt_audience: Option<String>,
    /// Webhook URL for webhook mode.
    pub webhook_url: Option<String>,
}

impl Default for AuthConfig {
    fn default() -> Self {
        Self {
            mode: AuthMode::None,
            jwt_issuer: None,
            jwt_audience: None,
            webhook_url: None,
        }
    }
}

// ─── Auth Verifier ───────────────────────────────────────────────────

/// The verifier holds config and cached JWKS keys (for JWT mode).
pub struct AuthVerifier {
    config: AuthConfig,
    /// Cached JWKS keys (refreshed periodically).
    jwks_cache: RwLock<Option<CachedJwks>>,
    http_client: reqwest::Client,
}

struct CachedJwks {
    keys: jsonwebtoken::jwk::JwkSet,
    fetched_at: std::time::Instant,
}

const JWKS_CACHE_TTL_SECS: u64 = 300; // 5 minutes

impl AuthVerifier {
    pub fn new(config: AuthConfig) -> Arc<Self> {
        info!("[Auth] Mode: {}", config.mode);
        if let Some(ref issuer) = config.jwt_issuer {
            info!("[Auth] JWT issuer: {issuer}");
        }
        if let Some(ref aud) = config.jwt_audience {
            info!("[Auth] JWT audience: {aud}");
        }
        if let Some(ref url) = config.webhook_url {
            info!("[Auth] Webhook URL: {url}");
        }
        Arc::new(Self {
            config,
            jwks_cache: RwLock::new(None),
            http_client: reqwest::Client::new(),
        })
    }

    /// Verify an auth context from a CONNECT payload.
    /// Returns `AuthResult::Skipped` when auth mode is `None`.
    pub async fn verify(&self, auth: Option<&AuthContext>) -> AuthResult {
        match self.config.mode {
            AuthMode::None => AuthResult::Skipped,
            AuthMode::Jwt => self.verify_jwt(auth).await,
            AuthMode::Webhook => self.verify_webhook(auth).await,
        }
    }

    // ── JWT Verification ────────────────────────────────────────────

    async fn verify_jwt(&self, auth: Option<&AuthContext>) -> AuthResult {
        let auth = match auth {
            Some(a) => a,
            None => return AuthResult::Rejected("auth required but no token provided".into()),
        };

        let token = &auth.token;
        if token.is_empty() {
            return AuthResult::Rejected("empty auth token".into());
        }

        // Decode the JWT header to find the `kid` (key ID)
        let header = match jsonwebtoken::decode_header(token) {
            Ok(h) => h,
            Err(e) => return AuthResult::Rejected(format!("invalid JWT header: {e}")),
        };

        let kid = match &header.kid {
            Some(k) => k.clone(),
            None => return AuthResult::Rejected("JWT missing kid".into()),
        };

        // Fetch or use cached JWKS
        let jwks = match self.get_jwks().await {
            Ok(j) => j,
            Err(e) => return AuthResult::Rejected(format!("failed to fetch JWKS: {e}")),
        };

        // Find the matching key
        let jwk = match jwks.keys.keys.iter().find(|k| {
            k.common.key_id.as_deref() == Some(&kid)
        }) {
            Some(k) => k,
            None => return AuthResult::Rejected(format!("no JWKS key found for kid={kid}")),
        };

        // Build the decoding key from the JWK
        let decoding_key = match jsonwebtoken::DecodingKey::from_jwk(jwk) {
            Ok(dk) => dk,
            Err(e) => return AuthResult::Rejected(format!("invalid JWK: {e}")),
        };

        // Configure validation
        let mut validation = jsonwebtoken::Validation::new(header.alg);
        if let Some(ref aud) = self.config.jwt_audience {
            validation.set_audience(&[aud]);
        } else {
            validation.validate_aud = false;
        }

        // Decode and validate
        let token_data = match jsonwebtoken::decode::<serde_json::Value>(
            token,
            &decoding_key,
            &validation,
        ) {
            Ok(td) => td,
            Err(e) => return AuthResult::Rejected(format!("JWT verification failed: {e}")),
        };

        // Extract user identity from claims
        let claims = &token_data.claims;
        let sub = claims.get("sub").and_then(|v| v.as_str()).unwrap_or("unknown");

        let mut user_claims = BTreeMap::new();
        // Copy useful claims
        for key in ["iss", "aud", "email", "name", "role", "roles", "org_id", "org_slug"] {
            if let Some(val) = claims.get(key) {
                user_claims.insert(key.to_string(), val.to_string().trim_matches('"').to_string());
            }
        }

        debug!("[Auth] JWT verified: sub={sub}");

        AuthResult::Authenticated(AuthenticatedUser {
            id: sub.to_string(),
            claims: user_claims,
        })
    }

    async fn get_jwks(&self) -> Result<CachedJwks, String> {
        // Check cache
        {
            let cache = self.jwks_cache.read().await;
            if let Some(ref cached) = *cache {
                if cached.fetched_at.elapsed().as_secs() < JWKS_CACHE_TTL_SECS {
                    return Ok(CachedJwks {
                        keys: cached.keys.clone(),
                        fetched_at: cached.fetched_at,
                    });
                }
            }
        }

        // Fetch JWKS
        let issuer = self.config.jwt_issuer.as_deref()
            .ok_or_else(|| "JWT mode requires --auth-jwt-issuer".to_string())?;

        debug!("[Auth] Fetching JWKS from {issuer}");

        let resp = self.http_client.get(issuer)
            .send()
            .await
            .map_err(|e| format!("JWKS fetch failed: {e}"))?;

        let jwks: jsonwebtoken::jwk::JwkSet = resp.json()
            .await
            .map_err(|e| format!("JWKS parse failed: {e}"))?;

        let cached = CachedJwks {
            keys: jwks.clone(),
            fetched_at: std::time::Instant::now(),
        };

        // Update cache
        {
            let mut cache = self.jwks_cache.write().await;
            *cache = Some(CachedJwks {
                keys: jwks,
                fetched_at: std::time::Instant::now(),
            });
        }

        Ok(cached)
    }

    // ── Webhook Verification ────────────────────────────────────────

    async fn verify_webhook(&self, auth: Option<&AuthContext>) -> AuthResult {
        let auth = match auth {
            Some(a) => a,
            None => return AuthResult::Rejected("auth required but no token provided".into()),
        };

        let webhook_url = match &self.config.webhook_url {
            Some(u) => u,
            None => return AuthResult::Rejected("webhook mode requires --auth-webhook".into()),
        };

        let body = serde_json::json!({
            "token": auth.token,
            "provider": auth.provider,
        });

        let resp = match self.http_client
            .post(webhook_url)
            .json(&body)
            .timeout(std::time::Duration::from_secs(5))
            .send()
            .await
        {
            Ok(r) => r,
            Err(e) => return AuthResult::Rejected(format!("webhook request failed: {e}")),
        };

        if !resp.status().is_success() {
            let status = resp.status();
            let body_text = resp.text().await.unwrap_or_default();
            warn!("[Auth] Webhook rejected: status={status}, body={body_text}");
            return AuthResult::Rejected(format!("webhook rejected: {status}"));
        }

        // Parse webhook response
        #[derive(Deserialize)]
        struct WebhookResponse {
            #[serde(default)]
            allow: bool,
            #[serde(default)]
            user_id: Option<String>,
            #[serde(default)]
            claims: Option<BTreeMap<String, String>>,
            #[serde(default)]
            reason: Option<String>,
        }

        let webhook_resp: WebhookResponse = match resp.json().await {
            Ok(r) => r,
            Err(e) => return AuthResult::Rejected(format!("webhook response parse failed: {e}")),
        };

        if !webhook_resp.allow {
            return AuthResult::Rejected(
                webhook_resp.reason.unwrap_or_else(|| "rejected by webhook".into())
            );
        }

        let user_id = webhook_resp.user_id.unwrap_or_else(|| "unknown".into());
        debug!("[Auth] Webhook verified: user_id={user_id}");

        AuthResult::Authenticated(AuthenticatedUser {
            id: user_id,
            claims: webhook_resp.claims.unwrap_or_default(),
        })
    }
}

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

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

    #[tokio::test]
    async fn none_mode_always_skips() {
        let v = AuthVerifier::new(AuthConfig::default());
        let result = v.verify(None).await;
        assert!(matches!(result, AuthResult::Skipped));

        let ctx = AuthContext { token: "anything".into(), provider: None };
        let result = v.verify(Some(&ctx)).await;
        assert!(matches!(result, AuthResult::Skipped));
    }

    #[tokio::test]
    async fn jwt_mode_rejects_missing_token() {
        let v = AuthVerifier::new(AuthConfig {
            mode: AuthMode::Jwt,
            jwt_issuer: Some("https://example.com/.well-known/jwks.json".into()),
            jwt_audience: None,
            webhook_url: None,
        });
        let result = v.verify(None).await;
        assert!(matches!(result, AuthResult::Rejected(_)));
    }

    #[tokio::test]
    async fn jwt_mode_rejects_empty_token() {
        let v = AuthVerifier::new(AuthConfig {
            mode: AuthMode::Jwt,
            jwt_issuer: Some("https://example.com/.well-known/jwks.json".into()),
            jwt_audience: None,
            webhook_url: None,
        });
        let ctx = AuthContext { token: "".into(), provider: None };
        let result = v.verify(Some(&ctx)).await;
        assert!(matches!(result, AuthResult::Rejected(_)));
    }

    #[tokio::test]
    async fn webhook_mode_rejects_missing_token() {
        let v = AuthVerifier::new(AuthConfig {
            mode: AuthMode::Webhook,
            jwt_issuer: None,
            jwt_audience: None,
            webhook_url: Some("http://localhost:9999/verify".into()),
        });
        let result = v.verify(None).await;
        assert!(matches!(result, AuthResult::Rejected(_)));
    }
}
