use anyhow::{anyhow, Result};
use chrono::{DateTime, Utc};
use jsonwebtoken::{decode, DecodingKey, TokenData, Validation, Algorithm};
use reqwest::Client;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use tokio::sync::Mutex;

/// JWT claims from Authentik
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JwtClaims {
    pub sub: String,
    pub iss: String,
    pub aud: String,
    pub exp: i64,
    pub iat: i64,
    pub nbf: Option<i64>,
    pub permissions: Option<Vec<String>>,
    pub groups: Option<Vec<String>>,
}

/// JWKS (JSON Web Key Set) response from Authentik
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JwksResponse {
    pub keys: Vec<JsonWebKey>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JsonWebKey {
    pub kty: String,
    pub use_: Option<String>,
    #[serde(rename = "kid")]
    pub key_id: Option<String>,
    pub n: Option<String>,
    pub e: Option<String>,
    pub alg: Option<String>,
}

/// JWT validator with JWKS caching
pub struct JwtValidator {
    pub issuer: String,
    pub audience: String,
    client: Client,
    jwks_cache: Arc<Mutex<(Option<JwksResponse>, DateTime<Utc>)>>,
    jwks_cache_ttl_secs: i64,
}

impl JwtValidator {
    pub fn new(issuer: String, audience: String, jwks_cache_ttl_secs: i64) -> Self {
        Self {
            issuer,
            audience,
            client: Client::new(),
            jwks_cache: Arc::new(Mutex::new((None, Utc::now()))),
            jwks_cache_ttl_secs,
        }
    }

    /// Fetch JWKS from issuer discovery endpoint
    async fn fetch_jwks(&self) -> Result<JwksResponse> {
        let discovery_url = format!("{}/.well-known/openid-configuration", self.issuer);
        tracing::debug!("Fetching OIDC discovery from {}", discovery_url);

        let discovery: serde_json::Value = self
            .client
            .get(&discovery_url)
            .send()
            .await?
            .json()
            .await?;

        let jwks_uri = discovery
            .get("jwks_uri")
            .and_then(|v| v.as_str())
            .ok_or_else(|| anyhow!("No jwks_uri in discovery doc"))?;

        tracing::debug!("Fetching JWKS from {}", jwks_uri);
        let jwks: JwksResponse = self.client.get(jwks_uri).send().await?.json().await?;

        if jwks.keys.is_empty() {
            return Err(anyhow!("No keys in JWKS response"));
        }

        Ok(jwks)
    }

    /// Get JWKS from cache or fetch fresh
    async fn get_jwks(&self) -> Result<JwksResponse> {
        let cache = self.jwks_cache.lock().await;
        let (cached_jwks, cached_at) = cache.clone();

        // Check if cache is still valid
        if let Some(jwks) = cached_jwks {
            let age = (Utc::now() - cached_at).num_seconds();
            if age < self.jwks_cache_ttl_secs {
                drop(cache);
                tracing::debug!("JWKS from cache (age: {}s)", age);
                return Ok(jwks);
            }
        }

        drop(cache);

        // Fetch fresh JWKS
        let jwks = self.fetch_jwks().await?;
        let mut cache = self.jwks_cache.lock().await;
        *cache = (Some(jwks.clone()), Utc::now());
        Ok(jwks)
    }

    /// Convert JWKS key to DecodingKey for RS256 validation
    fn jwks_to_decoding_key(key: &JsonWebKey) -> Result<DecodingKey> {
        // Only support RSA keys
        if key.kty != "RSA" {
            return Err(anyhow!("Unsupported key type: {}", key.kty));
        }

        let n = key.n.as_ref().ok_or_else(|| anyhow!("Missing RSA modulus"))?;
        let e = key.e.as_ref().ok_or_else(|| anyhow!("Missing RSA exponent"))?;

        DecodingKey::from_rsa_components(n, e).map_err(|e| anyhow!("Invalid RSA key: {}", e))
    }

    /// Validate JWT token and extract claims
    pub async fn validate_token(&self, token: &str) -> Result<JwtClaims> {
        // Decode header to check algorithm
        let header = jsonwebtoken::decode_header(token)
            .map_err(|e| anyhow!("Invalid token header: {}", e))?;

        // Pin to RS256 only (defense against algorithm confusion)
        if header.alg != Algorithm::RS256 {
            return Err(anyhow!(
                "Invalid algorithm: {:?}, expected RS256",
                header.alg
            ));
        }

        let kid = header
            .kid
            .as_ref()
            .ok_or_else(|| anyhow!("Token missing 'kid' header"))?;

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

        // Find key by kid
        let key = jwks
            .keys
            .iter()
            .find(|k| k.key_id.as_ref() == Some(kid))
            .ok_or_else(|| anyhow!("Key not found in JWKS: {}", kid))?;

        // Convert to DecodingKey
        let decoding_key = Self::jwks_to_decoding_key(key)?;

        // Validate token signature + claims
        let mut validation = Validation::new(Algorithm::RS256);
        validation.set_issuer(&[self.issuer.clone()]);
        validation.set_audience(&[self.audience.clone()]);
        validation.leeway = 60; // 60s clock skew tolerance

        let token_data: TokenData<JwtClaims> =
            decode::<JwtClaims>(token, &decoding_key, &validation)
                .map_err(|e| anyhow!("Token validation failed: {}", e))?;

        Ok(token_data.claims)
    }

    /// Extract bearer token from Authorization header
    pub fn extract_bearer_token(auth_header: &str) -> Result<String> {
        let parts: Vec<&str> = auth_header.split_whitespace().collect();
        if parts.len() != 2 || parts[0].to_lowercase() != "bearer" {
            return Err(anyhow!("Invalid Authorization header format"));
        }
        Ok(parts[1].to_string())
    }
}

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

    #[test]
    fn test_extract_bearer_token_valid() {
        let header = "Bearer eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0";
        let token = JwtValidator::extract_bearer_token(header).unwrap();
        assert_eq!(
            token,
            "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0"
        );
    }

    #[test]
    fn test_extract_bearer_token_invalid_format() {
        let header = "Basic dXNlcjpwYXNz";
        let result = JwtValidator::extract_bearer_token(header);
        assert!(result.is_err());
    }

    #[test]
    fn test_extract_bearer_token_missing() {
        let header = "Bearer";
        let result = JwtValidator::extract_bearer_token(header);
        assert!(result.is_err());
    }
}