//! JWT/OIDC token verification
//!
//! Provides JWT signature verification with JWKS (JSON Web Key Set) support.

use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};

use base64::Engine;
use jsonwebtoken::{decode, decode_header, Algorithm, DecodingKey, Validation};
use serde::Deserialize;
use tokio::sync::RwLock;

use iam_types::{Error, IamError, JwtClaims, Result, TokenValidationError};

/// JWKS (JSON Web Key Set) structure
#[derive(Debug, Clone, Deserialize)]
pub struct Jwks {
    pub keys: Vec<Jwk>,
}

/// Individual JWK (JSON Web Key)
#[derive(Debug, Clone, Deserialize)]
pub struct Jwk {
    pub kty: String,
    pub kid: Option<String>,
    #[serde(rename = "use")]
    pub key_use: Option<String>,
    pub alg: Option<String>,
    pub n: Option<String>,   // RSA modulus
    pub e: Option<String>,   // RSA exponent
    pub x: Option<String>,   // EC x coordinate
    pub y: Option<String>,   // EC y coordinate
    pub crv: Option<String>, // EC curve
}

impl Jwk {
    /// Convert to jsonwebtoken DecodingKey
    pub fn to_decoding_key(&self) -> std::result::Result<DecodingKey, TokenValidationError> {
        match self.kty.as_str() {
            "RSA" => {
                let n = self
                    .n
                    .as_ref()
                    .ok_or_else(|| TokenValidationError::Malformed("RSA key missing n".into()))?;
                let e = self
                    .e
                    .as_ref()
                    .ok_or_else(|| TokenValidationError::Malformed("RSA key missing e".into()))?;
                DecodingKey::from_rsa_components(n, e)
                    .map_err(|e| TokenValidationError::Malformed(e.to_string()))
            }
            "EC" => {
                let x = self
                    .x
                    .as_ref()
                    .ok_or_else(|| TokenValidationError::Malformed("EC key missing x".into()))?;
                let y = self
                    .y
                    .as_ref()
                    .ok_or_else(|| TokenValidationError::Malformed("EC key missing y".into()))?;
                DecodingKey::from_ec_components(x, y)
                    .map_err(|e| TokenValidationError::Malformed(e.to_string()))
            }
            other => Err(TokenValidationError::Malformed(format!(
                "Unsupported key type: {}",
                other
            ))),
        }
    }

    /// Get the algorithm for this key
    pub fn algorithm(&self) -> std::result::Result<Algorithm, TokenValidationError> {
        match self.alg.as_deref() {
            Some("RS256") => Ok(Algorithm::RS256),
            Some("RS384") => Ok(Algorithm::RS384),
            Some("RS512") => Ok(Algorithm::RS512),
            Some("ES256") => Ok(Algorithm::ES256),
            Some("ES384") => Ok(Algorithm::ES384),
            Some(other) => Err(TokenValidationError::Malformed(format!(
                "Unsupported algorithm: {}",
                other
            ))),
            None => {
                // Infer from key type
                match self.kty.as_str() {
                    "RSA" => Ok(Algorithm::RS256),
                    "EC" => match self.crv.as_deref() {
                        Some("P-256") => Ok(Algorithm::ES256),
                        Some("P-384") => Ok(Algorithm::ES384),
                        _ => Ok(Algorithm::ES256),
                    },
                    _ => Err(TokenValidationError::Malformed(
                        "Cannot infer algorithm".into(),
                    )),
                }
            }
        }
    }
}

/// Cached JWKS with expiration
struct CachedJwks {
    jwks: Jwks,
    fetched_at: Instant,
}

/// Configuration for JWT verifier
#[derive(Debug, Clone)]
pub struct JwtVerifierConfig {
    /// JWKS endpoint URL
    pub jwks_url: String,
    /// Expected issuer
    pub issuer: String,
    /// Expected audience
    pub audience: String,
    /// JWKS cache TTL
    pub jwks_cache_ttl: Duration,
    /// Allow clock skew (seconds)
    pub clock_skew_seconds: u64,
}

impl JwtVerifierConfig {
    /// Create a new config
    pub fn new(
        jwks_url: impl Into<String>,
        issuer: impl Into<String>,
        audience: impl Into<String>,
    ) -> Self {
        Self {
            jwks_url: jwks_url.into(),
            issuer: issuer.into(),
            audience: audience.into(),
            jwks_cache_ttl: Duration::from_secs(3600), // 1 hour default
            clock_skew_seconds: 60,
        }
    }

    /// Set JWKS cache TTL
    pub fn with_cache_ttl(mut self, ttl: Duration) -> Self {
        self.jwks_cache_ttl = ttl;
        self
    }

    /// Set allowed clock skew
    pub fn with_clock_skew(mut self, seconds: u64) -> Self {
        self.clock_skew_seconds = seconds;
        self
    }
}

/// JWT token verifier with JWKS support
pub struct JwtVerifier {
    config: JwtVerifierConfig,
    jwks_cache: RwLock<Option<CachedJwks>>,
    http_client: reqwest::Client,
}

impl JwtVerifier {
    /// Create a new JWT verifier
    pub fn new(config: JwtVerifierConfig) -> Self {
        Self {
            config,
            jwks_cache: RwLock::new(None),
            http_client: reqwest::Client::new(),
        }
    }

    /// Verify a JWT token and return its claims
    pub async fn verify(&self, token: &str) -> Result<JwtClaims> {
        // 1. Decode header to get kid
        let header =
            decode_header(token).map_err(|e| Error::Iam(IamError::InvalidToken(e.to_string())))?;

        // 2. Get the signing key
        let jwks = self.get_jwks().await?;
        let key = self.find_key(&jwks, header.kid.as_deref())?;

        // 3. Build validation
        let mut validation = Validation::new(
            key.algorithm()
                .map_err(|e| Error::Iam(IamError::InvalidToken(e.to_string())))?,
        );
        validation.set_issuer(&[&self.config.issuer]);
        validation.set_audience(&[&self.config.audience]);
        validation.leeway = self.config.clock_skew_seconds;

        // 4. Decode and verify
        let decoding_key = key
            .to_decoding_key()
            .map_err(|e| Error::Iam(IamError::InvalidToken(e.to_string())))?;

        let token_data = decode::<JwtClaims>(token, &decoding_key, &validation)
            .map_err(|e| Error::Iam(IamError::InvalidToken(e.to_string())))?;

        Ok(token_data.claims)
    }

    /// Verify a token and validate specific claims
    pub async fn verify_with_validation<F>(
        &self,
        token: &str,
        additional_validation: F,
    ) -> Result<JwtClaims>
    where
        F: FnOnce(&JwtClaims) -> Result<()>,
    {
        let claims = self.verify(token).await?;
        additional_validation(&claims)?;
        Ok(claims)
    }

    /// Get JWKS (from cache or fetch)
    async fn get_jwks(&self) -> Result<Jwks> {
        // Check cache
        {
            let cache = self.jwks_cache.read().await;
            if let Some(cached) = cache.as_ref() {
                if cached.fetched_at.elapsed() < self.config.jwks_cache_ttl {
                    return Ok(cached.jwks.clone());
                }
            }
        }

        // Fetch new JWKS
        let jwks = self.fetch_jwks().await?;

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

        Ok(jwks)
    }

    /// Fetch JWKS from the configured URL
    async fn fetch_jwks(&self) -> Result<Jwks> {
        let response = self
            .http_client
            .get(&self.config.jwks_url)
            .send()
            .await
            .map_err(|e| Error::Internal(format!("Failed to fetch JWKS: {}", e)))?;

        if !response.status().is_success() {
            return Err(Error::Internal(format!(
                "JWKS fetch failed with status: {}",
                response.status()
            )));
        }

        let jwks: Jwks = response
            .json()
            .await
            .map_err(|e| Error::Internal(format!("Failed to parse JWKS: {}", e)))?;

        Ok(jwks)
    }

    /// Find a signing key in JWKS by kid
    fn find_key<'a>(&self, jwks: &'a Jwks, kid: Option<&str>) -> Result<&'a Jwk> {
        let is_signing_key = |key: &&Jwk| match key.key_use.as_deref() {
            None => true,
            Some(use_type) => use_type.eq_ignore_ascii_case("sig"),
        };

        match kid {
            Some(kid) => {
                let key = jwks
                    .keys
                    .iter()
                    .find(|k| k.kid.as_deref() == Some(kid))
                    .ok_or_else(|| {
                        Error::Iam(IamError::InvalidToken(format!("Key not found: {}", kid)))
                    })?;

                if !is_signing_key(&key) {
                    return Err(Error::Iam(IamError::InvalidToken(format!(
                        "Key {} is not a signing key",
                        kid
                    ))));
                }

                Ok(key)
            }
            None => {
                let mut keys = jwks.keys.iter().filter(is_signing_key);
                let key = keys.next().ok_or_else(|| {
                    Error::Iam(IamError::InvalidToken(
                        "JWKS has no signing keys".into(),
                    ))
                })?;
                if keys.next().is_some() {
                    return Err(Error::Iam(IamError::InvalidToken(
                        "JWKS has multiple signing keys but kid is missing".into(),
                    )));
                }
                Ok(key)
            }
        }
    }

    /// Force refresh the JWKS cache
    pub async fn refresh_jwks(&self) -> Result<()> {
        let jwks = self.fetch_jwks().await?;
        let mut cache = self.jwks_cache.write().await;
        *cache = Some(CachedJwks {
            jwks,
            fetched_at: Instant::now(),
        });
        Ok(())
    }
}

/// Multi-issuer JWT verifier
pub struct MultiIssuerJwtVerifier {
    verifiers: HashMap<String, Arc<JwtVerifier>>,
}

impl MultiIssuerJwtVerifier {
    /// Create a new multi-issuer verifier
    pub fn new() -> Self {
        Self {
            verifiers: HashMap::new(),
        }
    }

    /// Add a verifier for an issuer
    pub fn add_issuer(&mut self, issuer: impl Into<String>, verifier: JwtVerifier) {
        self.verifiers.insert(issuer.into(), Arc::new(verifier));
    }

    /// Verify a token (automatically selects issuer from token)
    pub async fn verify(&self, token: &str) -> Result<JwtClaims> {
        // Extract issuer from token without verification
        let issuer = extract_issuer_from_token(token)?;

        let verifier = self.verifiers.get(&issuer).ok_or_else(|| {
            Error::Iam(IamError::InvalidToken(format!(
                "Unknown issuer: {}",
                issuer
            )))
        })?;

        verifier.verify(token).await
    }
}

/// Extract issuer from JWT without verification (for routing purposes)
fn extract_issuer_from_token(token: &str) -> Result<String> {
    let parts: Vec<&str> = token.split('.').collect();
    if parts.len() != 3 {
        return Err(Error::Iam(IamError::InvalidToken(
            "Invalid JWT format".into(),
        )));
    }

    // Decode payload (middle part)
    let payload_json = base64::engine::general_purpose::URL_SAFE_NO_PAD
        .decode(parts[1])
        .map_err(|e| Error::Iam(IamError::InvalidToken(format!("Invalid base64: {}", e))))?;

    let claims: UnverifiedClaims = serde_json::from_slice(&payload_json)
        .map_err(|e| Error::Iam(IamError::InvalidToken(format!("Invalid JSON: {}", e))))?;

    Ok(claims.iss)
}

impl Default for MultiIssuerJwtVerifier {
    fn default() -> Self {
        Self::new()
    }
}

/// Minimal claims for issuer detection
#[derive(Debug, Deserialize)]
struct UnverifiedClaims {
    iss: String,
}

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

    #[test]
    fn test_jwk_algorithm_inference() {
        let rsa_key = Jwk {
            kty: "RSA".into(),
            kid: Some("key1".into()),
            key_use: Some("sig".into()),
            alg: None,
            n: Some("abc".into()),
            e: Some("AQAB".into()),
            x: None,
            y: None,
            crv: None,
        };

        assert!(matches!(rsa_key.algorithm(), Ok(Algorithm::RS256)));

        let ec_key = Jwk {
            kty: "EC".into(),
            kid: Some("key2".into()),
            key_use: Some("sig".into()),
            alg: Some("ES256".into()),
            n: None,
            e: None,
            x: Some("x".into()),
            y: Some("y".into()),
            crv: Some("P-256".into()),
        };

        assert!(matches!(ec_key.algorithm(), Ok(Algorithm::ES256)));
    }

    #[test]
    fn test_verifier_config() {
        let config = JwtVerifierConfig::new(
            "https://example.com/.well-known/jwks.json",
            "https://example.com",
            "my-api",
        )
        .with_cache_ttl(Duration::from_secs(7200))
        .with_clock_skew(120);

        assert_eq!(config.jwks_cache_ttl, Duration::from_secs(7200));
        assert_eq!(config.clock_skew_seconds, 120);
    }
}
