jwt-verify 0.1.0

use anyhow::Result;
use jsonwebtoken::DecodingKey;
use reqwest::Client;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fmt::Debug;
use std::sync::RwLock;
use std::time::{Duration, Instant};
use tokio::sync::Mutex;

use crate::common::error::JwtError;

/// RSA key from JWK set
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct RSAKey {
    /// Key ID
    pub kid: String,
    /// Algorithm
    pub alg: String,
    /// Modulus
    pub n: String,
    /// Exponent
    pub e: String,
    /// Key use
    #[serde(rename = "use")]
    pub use_for: String,
}

/// JWK set response
#[derive(Debug, Deserialize)]
pub struct JwkSet {
    /// Keys in the set
    pub keys: Vec<RSAKey>,
}

/// Cached JWK
struct CachedJwk {
    /// Decoding key
    key: DecodingKey,
    /// Time when the key was inserted into the cache
    inserted_at: Instant,
}

impl std::fmt::Debug for CachedJwk {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CachedJwk")
            .field("inserted_at", &self.inserted_at)
            .finish()
    }
}

/// JWK provider for fetching and caching JWKs
#[derive(Debug)]
pub struct JwkProvider {
    /// JWK URL
    jwk_url: String,
    /// Issuer URL
    issuer: String,
    /// Keys cache
    keys_cache: RwLock<HashMap<String, CachedJwk>>,
    /// Last refresh time
    last_refresh: Mutex<Option<Instant>>,
    /// Cache duration
    cache_duration: Duration,
    /// Minimum refresh interval
    min_refresh_interval: Duration,
    /// HTTP client
    client: Client,
}

impl JwkProvider {
    /// Create a new JWK provider
    ///
    /// This constructor creates a JwkProvider without prefetching keys.
    /// Keys will be fetched on the first request.
    pub fn new(
        region: &str,
        user_pool_id: &str,
        cache_duration: Duration,
    ) -> Result<Self, JwtError> {
        // Default minimum refresh interval is 60 seconds for production use
        // Validate region
        if region.is_empty() {
            return Err(JwtError::ConfigurationError {
                parameter: Some("region".to_string()),
                error: "Region cannot be empty".to_string(),
            });
        }

        // Validate user pool ID
        if user_pool_id.is_empty() {
            return Err(JwtError::ConfigurationError {
                parameter: Some("user_pool_id".to_string()),
                error: "User pool ID cannot be empty".to_string(),
            });
        }

        // Validate user pool ID format (region_code-region-number)
        // AWS Cognito user pool IDs follow the format: region_code_region-number
        // For example: us-east-1_abcdefghi
        if !Self::is_valid_user_pool_id(user_pool_id) {
            return Err(JwtError::ConfigurationError {
                parameter: Some("user_pool_id".to_string()),
                error: format!(
                    "Invalid user pool ID format: {}. Expected format: region_number",
                    user_pool_id
                ),
            });
        }

        // Create JWK URL
        let issuer = format!(
            "https://cognito-idp.{}.amazonaws.com/{}",
            region, user_pool_id
        );
        let jwk_url = format!("{}/.well-known/jwks.json", issuer);

        // Create HTTP client
        let client = Client::builder().use_rustls_tls().build().map_err(|e| {
            JwtError::ConfigurationError {
                parameter: Some("http_client".to_string()),
                error: format!("Failed to create HTTP client: {}", e),
            }
        })?;

        let provider = Self {
            jwk_url,
            issuer,
            keys_cache: RwLock::new(HashMap::new()),
            last_refresh: Mutex::new(None),
            cache_duration,
            min_refresh_interval: Duration::from_secs(60), // Default to 60 seconds for production
            client,
        };

        Ok(provider)
    }

    /// Create a new JWK provider from a JWKS URL (for OIDC providers)
    ///
    /// This constructor creates a JwkProvider for generic OIDC providers.
    /// Keys will be fetched on the first request.
    ///
    /// # Arguments
    /// * `jwks_url` - The full URL to the JWKS endpoint (e.g., "https://example.com/.well-known/jwks.json")
    /// * `issuer` - The issuer URL (e.g., "https://example.com")
    /// * `cache_duration` - How long to cache keys before refreshing
    pub fn from_jwks_url(
        jwks_url: &str,
        issuer: &str,
        cache_duration: Duration,
    ) -> Result<Self, JwtError> {
        // Validate JWKS URL
        if jwks_url.is_empty() {
            return Err(JwtError::ConfigurationError {
                parameter: Some("jwks_url".to_string()),
                error: "JWKS URL cannot be empty".to_string(),
            });
        }

        // Validate issuer
        if issuer.is_empty() {
            return Err(JwtError::ConfigurationError {
                parameter: Some("issuer".to_string()),
                error: "Issuer cannot be empty".to_string(),
            });
        }

        // Validate URL format
        if !jwks_url.starts_with("http://") && !jwks_url.starts_with("https://") {
            return Err(JwtError::ConfigurationError {
                parameter: Some("jwks_url".to_string()),
                error: "JWKS URL must start with http:// or https://".to_string(),
            });
        }

        // Create HTTP client
        let client = Client::builder().use_rustls_tls().build().map_err(|e| {
            JwtError::ConfigurationError {
                parameter: Some("http_client".to_string()),
                error: format!("Failed to create HTTP client: {}", e),
            }
        })?;

        let provider = Self {
            jwk_url: jwks_url.to_string(),
            issuer: issuer.to_string(),
            keys_cache: RwLock::new(HashMap::new()),
            last_refresh: Mutex::new(None),
            cache_duration,
            min_refresh_interval: Duration::from_secs(60), // Default to 60 seconds for production
            client,
        };

        Ok(provider)
    }

    /// Validate user pool ID format
    fn is_valid_user_pool_id(user_pool_id: &str) -> bool {
        // AWS Cognito user pool IDs follow the format: region_code_region-number
        // For example: us-east-1_abcdefghi
        let parts: Vec<&str> = user_pool_id.split('_').collect();

        if parts.len() != 2 {
            return false;
        }

        // The second part should be alphanumeric
        parts[1].chars().all(|c| c.is_alphanumeric())
    }

    /// Create a new JWK provider with a custom base URL (for testing)
    #[cfg(test)]
    pub fn new_with_base_url(
        base_url: &str,
        issuer: &str,
        cache_duration: Duration,
    ) -> Result<Self, JwtError> {
        // For tests, use a shorter minimum refresh interval (1 second)
        Self::new_with_base_url_and_refresh_interval(
            base_url,
            issuer,
            cache_duration,
            Duration::from_secs(1),
        )
    }

    #[cfg(test)]
    pub fn new_with_base_url_and_refresh_interval(
        base_url: &str,
        issuer: &str,
        cache_duration: Duration,
        min_refresh_interval: Duration,
    ) -> Result<Self, JwtError> {
        // Create JWK URL
        let jwk_url = format!("{}/.well-known/jwks.json", base_url);

        // Create HTTP client
        let client = Client::builder().use_rustls_tls().build().map_err(|e| {
            JwtError::ConfigurationError {
                parameter: Some("http_client".to_string()),
                error: format!("Failed to create HTTP client: {}", e),
            }
        })?;

        let provider = Self {
            jwk_url,
            issuer: issuer.to_string(),
            keys_cache: RwLock::new(HashMap::new()),
            last_refresh: Mutex::new(None),
            cache_duration,
            min_refresh_interval,
            client,
        };

        Ok(provider)
    }

    /// Get the issuer URL
    pub fn get_issuer(&self) -> &str {
        &self.issuer
    }

    /// Prefetch JWKs from the Cognito user pool
    /// This method should be called when the system starts to ensure JWKs are available
    pub async fn prefetch_keys(&self) -> Result<(), JwtError> {
        self.refresh_keys().await
    }

    /// Get a JWK by key ID
    pub async fn get_key(&self, kid: &str) -> Result<DecodingKey, JwtError> {
        // Check if key is in cache
        {
            let cache = self.keys_cache.read().unwrap();
            if let Some(cached_jwk) = cache.get(kid) {
                let now = Instant::now();
                if now.duration_since(cached_jwk.inserted_at) < self.cache_duration {
                    // Cache hit
                    return Ok(cached_jwk.key.clone());
                }
                // Cache entry expired
            } else {
                // Cache miss
            }
        }

        // Key not found or expired, refresh keys
        self.refresh_keys().await?;

        // Try to get key from cache again
        {
            let cache = self.keys_cache.read().unwrap();
            if let Some(cached_jwk) = cache.get(kid) {
                return Ok(cached_jwk.key.clone());
            }
        }

        // Key still not found
        Err(JwtError::KeyNotFound(kid.to_string()))
    }

    /// Refresh JWKs
    async fn refresh_keys(&self) -> Result<(), JwtError> {
        // Check if we need to refresh
        {
            let mut last_refresh = self.last_refresh.lock().await;
            if let Some(time) = *last_refresh {
                let now = Instant::now();
                if now.duration_since(time) < self.min_refresh_interval {
                    // Don't refresh more than once per configured interval
                    tracing::debug!(
                        "Skipping JWK refresh, last refresh was less than {:?} ago",
                        self.min_refresh_interval
                    );
                    return Ok(());
                }
            }

            // Update last refresh time
            *last_refresh = Some(Instant::now());
        }

        tracing::debug!("Fetching JWKs from {}", self.jwk_url);

        // Fetch JWKs with retry for transient failures
        let mut retry_count = 0;
        let max_retries = 3;
        let mut last_error = None;

        while retry_count < max_retries {
            match self.fetch_and_parse_jwks().await {
                Ok(()) => {
                    tracing::debug!("Successfully refreshed JWKs");
                    return Ok(());
                }
                Err(e) => {
                    // Only retry on network errors, not on parsing errors
                    match &e {
                        JwtError::JwksFetchError { .. } => {
                            retry_count += 1;
                            if retry_count < max_retries {
                                tracing::warn!(
                                    "Failed to fetch JWKs (attempt {}/{}): {}. Retrying...",
                                    retry_count,
                                    max_retries,
                                    e
                                );
                                tokio::time::sleep(Duration::from_millis(500 * (1 << retry_count)))
                                    .await;
                            }
                            last_error = Some(e);
                        }
                        _ => {
                            // Don't retry on parsing errors
                            return Err(e);
                        }
                    }
                }
            }
        }

        // All retries failed
        Err(last_error.unwrap_or_else(|| JwtError::JwksFetchError {
            url: Some(self.jwk_url.clone()),
            error: "Failed to fetch JWKs after multiple attempts".to_string(),
        }))
    }

    /// Remove expired keys from the cache
    fn prune_expired_keys(&self, cache: &mut HashMap<String, CachedJwk>, now: Instant) {
        let expired_keys: Vec<String> = cache
            .iter()
            .filter(|(_, cached_jwk)| {
                now.duration_since(cached_jwk.inserted_at) >= self.cache_duration
            })
            .map(|(kid, _)| kid.clone())
            .collect();

        if !expired_keys.is_empty() {
            tracing::debug!("Pruning {} expired keys from cache", expired_keys.len());
            for kid in expired_keys {
                cache.remove(&kid);
            }
        }
    }

    /// Fetch and parse JWKs from the Cognito endpoint
    async fn fetch_and_parse_jwks(&self) -> Result<(), JwtError> {
        // Fetch JWKs
        let response = self
            .client
            .get(&self.jwk_url)
            .timeout(Duration::from_secs(5))
            .send()
            .await
            .map_err(|e| {
                let error_msg = if e.is_timeout() {
                    "Request timed out".to_string()
                } else if e.is_connect() {
                    "Connection error".to_string()
                } else {
                    format!("HTTP request failed: {}", e)
                };

                JwtError::JwksFetchError {
                    url: Some(self.jwk_url.clone()),
                    error: error_msg,
                }
            })?;

        // Check response status
        if !response.status().is_success() {
            return Err(JwtError::JwksFetchError {
                url: Some(self.jwk_url.clone()),
                error: format!("Failed to fetch JWKs: HTTP {}", response.status()),
            });
        }

        // Parse response
        let jwk_set: JwkSet = response.json().await.map_err(|e| JwtError::ParseError {
            part: Some("jwk_response".to_string()),
            error: format!("Failed to parse JWK response: {}", e),
        })?;

        // Check if we got any keys
        if jwk_set.keys.is_empty() {
            return Err(JwtError::JwksFetchError {
                url: Some(self.jwk_url.clone()),
                error: "JWK set is empty".to_string(),
            });
        }

        tracing::debug!("Fetched {} JWKs from Cognito", jwk_set.keys.len());

        // Update cache
        {
            let mut cache = self.keys_cache.write().unwrap();
            let now = Instant::now();

            // Remove expired keys before adding new ones
            self.prune_expired_keys(&mut cache, now);

            for key in jwk_set.keys {
                // Validate key fields
                if key.kid.is_empty() {
                    tracing::warn!("Skipping JWK with empty kid");
                    continue;
                }

                if key.n.is_empty() || key.e.is_empty() {
                    tracing::warn!("Skipping JWK with empty RSA components: kid={}", key.kid);
                    continue;
                }

                // Create decoding key
                let decoding_key =
                    DecodingKey::from_rsa_components(&key.n, &key.e).map_err(|e| {
                        JwtError::ParseError {
                            part: Some("jwk".to_string()),
                            error: format!("Failed to create decoding key: {}", e),
                        }
                    })?;

                // Add to cache
                cache.insert(
                    key.kid.clone(),
                    CachedJwk {
                        key: decoding_key,
                        inserted_at: now,
                    },
                );

                tracing::debug!("Cached JWK with kid={}", key.kid);
            }
        }

        Ok(())
    }
}