jwt-verify 0.1.0

use anyhow::Result;
use serde::de::DeserializeOwned;
use std::collections::HashMap;

use crate::claims::{OidcAccessTokenClaims, OidcIdTokenClaims, OidcJwtClaims};
use crate::common::error::{ErrorLogger, ErrorVerbosity, JwtError};
use crate::common::token::TokenParser;
use crate::jwk::provider::JwkProvider;
use crate::jwk::registry::{JwkProviderRegistry, RegistryError};
use crate::oidc::config::OidcProviderConfig;
use crate::oidc::discovery::OidcDiscovery;
use crate::verifier::{AccessTokenClaims, IdTokenClaims, JwtVerifier};

/// OIDC JWT verifier
///
/// This is the main entry point for verifying JWTs from OIDC-compatible identity providers.
/// It manages multiple OIDC providers and provides methods for verifying different types of tokens.
///
/// The verifier uses issuer-based provider selection to determine which provider
/// should verify a given token. This means that tokens must contain a valid issuer
/// claim that matches one of the registered providers.
///
/// # Examples
///
/// ```
/// use jwt_verify::oidc::{OidcJwtVerifier, OidcProviderConfig};
/// use std::time::Duration;
///
/// // Create configurations for multiple OIDC providers
/// let config1 = OidcProviderConfig::new(
///     "https://accounts.example1.com",
///     Some("https://accounts.example1.com/.well-known/jwks.json"),
///     &["client1".to_string()],
/// ).unwrap();
///
/// let config2 = OidcProviderConfig::new(
///     "https://accounts.example2.com",
///     Some("https://accounts.example2.com/.well-known/jwks.json"),
///     &["client2".to_string()],
/// ).unwrap();
///
/// // Create a verifier with multiple providers
/// let verifier = OidcJwtVerifier::new(vec![config1, config2]).unwrap();
///
/// // Verify a token (the issuer in the token will determine which provider to use)
/// // let token = "..."; // JWT token string
/// // let claims = verifier.verify_id_token(token).await.unwrap();
/// ```
#[derive(Debug)]
pub struct OidcJwtVerifier {
    /// JWK provider registry
    jwk_registry: JwkProviderRegistry,
    /// Configurations for different OIDC providers
    configs: HashMap<String, OidcProviderConfig>,
    /// OIDC discovery service
    discovery: OidcDiscovery,
    /// Error logger
    error_logger: ErrorLogger,
}

impl OidcJwtVerifier {
    /// Create a new verifier with multiple OIDC provider configurations
    ///
    /// This constructor takes a vector of `OidcProviderConfig` objects, each representing
    /// a different OIDC provider. The verifier will register all these providers
    /// and use them for token verification based on the issuer claim in the tokens.
    ///
    /// # Parameters
    ///
    /// * `configs` - Vector of configurations for different OIDC providers
    ///
    /// # Returns
    ///
    /// Returns a `Result` containing the new `OidcJwtVerifier` if successful, or a `JwtError`
    /// if any of the configurations are invalid or if there's an error registering the providers.
    ///
    /// # Examples
    ///
    /// ```
    /// use jwt_verify::oidc::{OidcJwtVerifier, OidcProviderConfig};
    ///
    /// // Create configurations for multiple OIDC providers
    /// let config1 = OidcProviderConfig::new(
    ///     "https://accounts.example1.com",
    ///     Some("https://accounts.example1.com/.well-known/jwks.json"),
    ///     &["client1".to_string()],
    /// ).unwrap();
    ///
    /// let config2 = OidcProviderConfig::new(
    ///     "https://accounts.example2.com",
    ///     Some("https://accounts.example2.com/.well-known/jwks.json"),
    ///     &["client2".to_string()],
    /// ).unwrap();
    ///
    /// // Create a verifier with multiple providers
    /// let verifier = OidcJwtVerifier::new(vec![config1, config2]).unwrap();
    /// ```
    pub fn new(configs: Vec<OidcProviderConfig>) -> Result<Self, JwtError> {
        // Create a verifier with an empty registry
        let mut verifier = Self {
            jwk_registry: JwkProviderRegistry::new(),
            configs: HashMap::new(),
            discovery: OidcDiscovery::new(std::time::Duration::from_secs(3600 * 24)), // Default 24 hours
            error_logger: ErrorLogger::new(ErrorVerbosity::Standard),
        };

        // Register each provider
        for config in configs {
            let issuer = config.issuer.clone();
            verifier.add_provider(&issuer, config)?;
        }

        Ok(verifier)
    }

    /// Create a new verifier with a single OIDC provider configuration
    ///
    /// This is a convenience constructor that creates a verifier with a single OIDC provider.
    ///
    /// # Parameters
    ///
    /// * `issuer` - The issuer URL of the OIDC provider
    /// * `jwks_url` - Optional JWKS URL (if None, will be discovered from issuer)
    /// * `client_ids` - List of allowed client IDs for this provider
    ///
    /// # Returns
    ///
    /// Returns a `Result` containing the new `OidcJwtVerifier` if successful, or a `JwtError`
    /// if the configuration is invalid or if there's an error registering the provider.
    ///
    /// # Examples
    ///
    /// ```
    /// use jwt_verify::oidc::OidcJwtVerifier;
    ///
    /// // Create a verifier with a single OIDC provider
    /// let verifier = OidcJwtVerifier::new_single_provider(
    ///     "https://accounts.example.com",
    ///     Some("https://accounts.example.com/.well-known/jwks.json"),
    ///     &["client1".to_string()],
    /// ).unwrap();
    /// ```
    pub fn new_single_provider(
        issuer: &str,
        jwks_url: Option<&str>,
        client_ids: &[String],
    ) -> Result<Self, JwtError> {
        let config = OidcProviderConfig::new(issuer, jwks_url, client_ids, None)?;
        Self::new(vec![config])
    }

    /// Add an OIDC provider with configuration
    ///
    /// This method adds a new OIDC provider to the verifier. The provider is identified by
    /// the provided ID, which should be unique among all registered providers.
    ///
    /// # Parameters
    ///
    /// * `id` - Unique identifier for the provider
    /// * `config` - Configuration for the provider
    ///
    /// # Returns
    ///
    /// Returns `Ok(())` if the provider was successfully added, or a `JwtError` if there
    /// was an error registering the provider.
    ///
    /// # Examples
    ///
    /// ```
    /// use jwt_verify::oidc::{OidcJwtVerifier, OidcProviderConfig};
    ///
    /// // Create a verifier
    /// let mut verifier = OidcJwtVerifier::new(vec![]).unwrap();
    ///
    /// // Create a configuration for an OIDC provider
    /// let config = OidcProviderConfig::new(
    ///     "https://accounts.example.com",
    ///     Some("https://accounts.example.com/.well-known/jwks.json"),
    ///     &["client1".to_string()],
    /// ).unwrap();
    ///
    /// // Add the provider to the verifier
    /// verifier.add_provider("example", config).unwrap();
    /// ```
    pub fn add_provider(&mut self, id: &str, config: OidcProviderConfig) -> Result<(), JwtError> {
        // For OIDC providers, we need to create a JwkProvider that works with the registry

        // Extract the domain from the issuer URL to use as the "region"
        let issuer = config.issuer.clone();

        // Create JWK provider from JWKS URL
        let jwks_url = config
            .jwks_url
            .as_ref()
            .ok_or_else(|| JwtError::ConfigurationError {
                parameter: Some("jwks_url".to_string()),
                error: "JWKS URL is required".to_string(),
            })?;

        let jwk_provider =
            JwkProvider::from_jwks_url(jwks_url, &issuer, config.jwk_cache_duration)?;

        // Add to registry and configs
        self.jwk_registry
            .register(id, jwk_provider)
            .map_err(|e| match e {
                RegistryError::JwtError(err) => err,
                other => JwtError::ConfigurationError {
                    parameter: Some("provider_registration".to_string()),
                    error: other.to_string(),
                },
            })?;

        self.configs.insert(id.to_string(), config);

        Ok(())
    }

    /// Add an OIDC provider with discovery
    ///
    /// This method adds a new OIDC provider to the verifier using auto-discovery
    /// to find the JWKS URL and other provider configuration.
    ///
    /// # Parameters
    ///
    /// * `id` - Unique identifier for the provider
    /// * `issuer` - The issuer URL of the OIDC provider
    /// * `client_ids` - List of allowed client IDs for this provider
    ///
    /// # Returns
    ///
    /// Returns `Ok(())` if the provider was successfully added, or a `JwtError` if there
    /// was an error discovering or registering the provider.
    ///
    /// # Examples
    ///
    /// ```
    /// use jwt_verify::oidc::OidcJwtVerifier;
    ///
    /// // Create a verifier
    /// let mut verifier = OidcJwtVerifier::new(vec![]).unwrap();
    ///
    /// // Add a provider with discovery
    /// verifier.add_provider_with_discovery(
    ///     "example",
    ///     "https://accounts.example.com",
    ///     &["client1".to_string()],
    /// ).unwrap();
    /// ```
    pub fn add_provider_with_discovery(
        &mut self,
        id: &str,
        issuer: &str,
        client_ids: &[String],
    ) -> Result<(), JwtError> {
        let config = OidcProviderConfig::with_discovery(issuer, client_ids)?;
        self.add_provider(id, config)
    }

    /// Get the list of registered provider IDs
    ///
    /// This method returns a list of all provider IDs that have been registered
    /// with the verifier.
    ///
    /// # Returns
    ///
    /// Returns a vector of provider IDs.
    ///
    /// # Examples
    ///
    /// ```
    /// use jwt_verify::oidc::{OidcJwtVerifier, OidcProviderConfig};
    ///
    /// // Create a verifier with multiple providers
    /// let config1 = OidcProviderConfig::new(
    ///     "https://accounts.example1.com",
    ///     Some("https://accounts.example1.com/.well-known/jwks.json"),
    ///     &["client1".to_string()],
    /// ).unwrap();
    ///
    /// let config2 = OidcProviderConfig::new(
    ///     "https://accounts.example2.com",
    ///     Some("https://accounts.example2.com/.well-known/jwks.json"),
    ///     &["client2".to_string()],
    /// ).unwrap();
    ///
    /// let verifier = OidcJwtVerifier::new(vec![config1, config2]).unwrap();
    ///
    /// // Get the list of provider IDs
    /// let provider_ids = verifier.get_provider_ids();
    /// assert_eq!(provider_ids.len(), 2);
    /// ```
    pub fn get_provider_ids(&self) -> Vec<String> {
        self.jwk_registry.list_ids()
    }

    /// Remove a provider
    ///
    /// This method removes a provider from the verifier. The provider is identified
    /// by the provided ID.
    ///
    /// # Parameters
    ///
    /// * `id` - Unique identifier for the provider to remove
    ///
    /// # Returns
    ///
    /// Returns `Ok(())` if the provider was successfully removed, or a `JwtError` if
    /// the provider was not found.
    ///
    /// # Examples
    ///
    /// ```
    /// use jwt_verify::oidc::{OidcJwtVerifier, OidcProviderConfig};
    ///
    /// // Create a verifier with a provider
    /// let config = OidcProviderConfig::new(
    ///     "https://accounts.example.com",
    ///     Some("https://accounts.example.com/.well-known/jwks.json"),
    ///     &["client1".to_string()],
    /// ).unwrap();
    ///
    /// let mut verifier = OidcJwtVerifier::new(vec![config]).unwrap();
    ///
    /// // Remove the provider
    /// verifier.remove_provider("https://accounts.example.com").unwrap();
    /// ```
    pub fn remove_provider(&mut self, id: &str) -> Result<(), JwtError> {
        // Remove from registry
        self.jwk_registry.remove(id).map_err(|e| match e {
            RegistryError::ProviderNotFound(_) => JwtError::ConfigurationError {
                parameter: Some("provider_id".to_string()),
                error: format!("Provider '{}' not found", id),
            },
            RegistryError::JwtError(err) => err,
            other => JwtError::ConfigurationError {
                parameter: Some("provider_id".to_string()),
                error: other.to_string(),
            },
        })?;

        // Remove from configs
        self.configs.remove(id);

        Ok(())
    }

    /// Set the error verbosity level
    ///
    /// This method sets the verbosity level for error logging and reporting.
    ///
    /// # Parameters
    ///
    /// * `verbosity` - The error verbosity level
    ///
    /// # Examples
    ///
    /// ```
    /// use jwt_verify::oidc::{OidcJwtVerifier, OidcProviderConfig};
    /// use jwt_verify::ErrorVerbosity;
    ///
    /// // Create a verifier
    /// let mut verifier = OidcJwtVerifier::new(vec![]).unwrap();
    ///
    /// // Set the error verbosity level
    /// verifier.set_error_verbosity(ErrorVerbosity::Detailed);
    /// ```
    pub fn set_error_verbosity(&mut self, verbosity: ErrorVerbosity) {
        self.error_logger = ErrorLogger::new(verbosity);
    }

    /// Set the discovery cache duration
    ///
    /// This method sets the duration for which OIDC discovery documents are cached.
    ///
    /// # Parameters
    ///
    /// * `duration` - The cache duration
    ///
    /// # Examples
    ///
    /// ```
    /// use jwt_verify::oidc::OidcJwtVerifier;
    /// use std::time::Duration;
    ///
    /// // Create a verifier
    /// let mut verifier = OidcJwtVerifier::new(vec![]).unwrap();
    ///
    /// // Set the discovery cache duration
    /// verifier.set_discovery_cache_duration(Duration::from_secs(3600 * 12)); // 12 hours
    /// ```
    pub fn set_discovery_cache_duration(&mut self, duration: std::time::Duration) {
        self.discovery = OidcDiscovery::new(duration);
    }

    /// Prefetch JWKs for all providers
    ///
    /// This method prefetches JWKs for all registered providers. This can be useful
    /// to warm up the cache before handling requests.
    ///
    /// # Returns
    ///
    /// Returns a vector of tuples containing the provider ID and the result of the
    /// prefetch operation.
    ///
    /// # Examples
    ///
    /// ```
    /// use jwt_verify::oidc::{OidcJwtVerifier, OidcProviderConfig};
    ///
    /// // Create a verifier with multiple providers
    /// let config1 = OidcProviderConfig::new(
    ///     "https://accounts.example1.com",
    ///     Some("https://accounts.example1.com/.well-known/jwks.json"),
    ///     &["client1".to_string()],
    /// ).unwrap();
    ///
    /// let config2 = OidcProviderConfig::new(
    ///     "https://accounts.example2.com",
    ///     Some("https://accounts.example2.com/.well-known/jwks.json"),
    ///     &["client2".to_string()],
    /// ).unwrap();
    ///
    /// let verifier = OidcJwtVerifier::new(vec![config1, config2]).unwrap();
    ///
    /// // Prefetch JWKs for all providers
    /// // let results = verifier.hydrate().await;
    /// ```
    pub async fn hydrate(&self) -> Vec<(String, Result<(), JwtError>)> {
        self.jwk_registry.hydrate().await
    }

    /// Verify a token with generic type support
    ///
    /// This method verifies a JWT token and returns the claims as the specified type.
    /// It automatically selects the appropriate provider based on the issuer claim in the token.
    ///
    /// # Parameters
    ///
    /// * `token` - The JWT token to verify
    ///
    /// # Returns
    ///
    /// Returns a `Result` containing the verified claims if successful, or a `JwtError`
    /// if verification fails.
    ///
    /// # Examples
    ///
    /// ```
    /// use jwt_verify::oidc::{OidcJwtVerifier, OidcIdTokenClaims};
    ///
    /// // Create a verifier
    /// let verifier = OidcJwtVerifier::new_single_provider(
    ///     "https://accounts.example.com",
    ///     Some("https://accounts.example.com/.well-known/jwks.json"),
    ///     &["client1".to_string()],
    /// ).unwrap();
    ///
    /// // Verify a token
    /// let token = "..."; // JWT token string
    /// let claims = verifier.verify::<OidcIdTokenClaims>(token).await.unwrap();
    /// ```
    pub async fn verify<T>(&self, token: &str) -> Result<T, JwtError>
    where
        T: DeserializeOwned + TryFrom<OidcJwtClaims, Error = JwtError>,
    {
        // Extract issuer from token without full validation yet
        let issuer = TokenParser::extract_issuer(token)?;

        // Find the provider ID based on the issuer
        let provider_id = self.find_provider_id_by_issuer(&issuer)?;

        // Verify with the specific provider
        self.verify_generic_with_provider::<T>(token, &provider_id)
            .await
    }

    /// Find a provider ID by issuer
    ///
    /// This method finds the provider ID that matches the given issuer.
    ///
    /// # Parameters
    ///
    /// * `issuer` - The issuer to match
    ///
    /// # Returns
    ///
    /// Returns a `Result` containing the provider ID if found, or a `JwtError`
    /// if no matching provider is found.
    fn find_provider_id_by_issuer(&self, issuer: &str) -> Result<String, JwtError> {
        for (id, config) in &self.configs {
            if config.issuer == issuer {
                return Ok(id.clone());
            }
        }

        Err(JwtError::InvalidIssuer {
            expected: "a registered OIDC provider".to_string(),
            actual: issuer.to_string(),
        })
    }

    /// Verify a token with a specific provider with generic type support
    async fn verify_generic_with_provider<T>(
        &self,
        token: &str,
        provider_id: &str,
    ) -> Result<T, JwtError>
    where
        T: DeserializeOwned + TryFrom<OidcJwtClaims, Error = JwtError>,
    {
        // Get provider and config
        let jwk_provider = self.jwk_registry.get(provider_id).map_err(|e| match e {
            RegistryError::ProviderNotFound(_) => JwtError::ConfigurationError {
                parameter: Some("provider_id".to_string()),
                error: format!("Provider '{}' not found", provider_id),
            },
            RegistryError::JwtError(err) => err,
            other => JwtError::ConfigurationError {
                parameter: Some("provider_id".to_string()),
                error: other.to_string(),
            },
        })?;

        let config = self
            .configs
            .get(provider_id)
            .ok_or_else(|| JwtError::ConfigurationError {
                parameter: Some("provider_id".to_string()),
                error: format!("Configuration for provider '{}' not found", provider_id),
            })?;

        // Parse header
        let header = TokenParser::parse_token_header(token)?;

        // Get key
        let key = jwk_provider.get_key(&header.kid).await?;

        // Create validation using the registry
        let validation = self
            .jwk_registry
            .create_validation_for_issuer(
                jwk_provider.get_issuer(),
                config.clock_skew,
                &config.client_ids,
            )
            .map_err(|e| match e {
                RegistryError::JwtError(err) => err,
                other => JwtError::ConfigurationError {
                    parameter: Some("validation".to_string()),
                    error: other.to_string(),
                },
            })?;

        // Parse claims as base type first
        let claims: OidcJwtClaims = TokenParser::parse_token_claims(token, &key, &validation)?;

        // Validate base claims
        self.validate_claims(&claims, config)?;

        // Convert to the requested type
        T::try_from(claims)
    }

    /// Validate claims
    ///
    /// This method validates the claims of a token against the provider's configuration.
    ///
    /// # Parameters
    ///
    /// * `claims` - The claims to validate
    /// * `config` - The provider configuration
    ///
    /// # Returns
    ///
    /// Returns `Ok(())` if the claims are valid, or a `JwtError` if validation fails.
    fn validate_claims(
        &self,
        claims: &OidcJwtClaims,
        config: &OidcProviderConfig,
    ) -> Result<(), JwtError> {
        // Validate issuer
        if !claims.validate_issuer(&config.issuer) {
            return Err(JwtError::InvalidClaim {
                claim: "iss".to_string(),
                reason: format!("Expected issuer {}, got {}", config.issuer, claims.iss),
                value: Some(claims.iss.clone()),
            });
        }

        // Validate client ID (using audience and/or azp)
        if !claims.validate_client_id(&config.client_ids) {
            return Err(JwtError::InvalidClaim {
                claim: "aud/azp".to_string(),
                reason: format!("Client ID not allowed: {}", claims.aud),
                value: Some(claims.aud.clone()),
            });
        }

        // Check required claims
        for claim_name in &config.required_claims {
            match claim_name.as_str() {
                "sub" => {
                    if claims.sub.is_empty() {
                        return Err(JwtError::InvalidClaim {
                            claim: "sub".to_string(),
                            reason: "Subject claim is empty".to_string(),
                            value: None,
                        });
                    }
                }
                "iss" => {
                    if claims.iss.is_empty() {
                        return Err(JwtError::InvalidClaim {
                            claim: "iss".to_string(),
                            reason: "Issuer claim is empty".to_string(),
                            value: None,
                        });
                    }
                }
                "aud" => {
                    // Already validated by jsonwebtoken
                    // if claims.aud.is_empty() {
                    //     return Err(JwtError::InvalidClaim {
                    //         claim: "aud".to_string(),
                    //         reason: "Audience claim is empty".to_string(),
                    //         value: None,
                    //     });
                    // }
                }
                "exp" => {
                    // Already validated by jsonwebtoken
                }
                "iat" => {
                    // Already validated by jsonwebtoken
                }
                "auth_time" => {
                    if claims.auth_time.is_none() {
                        return Err(JwtError::InvalidClaim {
                            claim: "auth_time".to_string(),
                            reason: "Authentication time claim is missing".to_string(),
                            value: None,
                        });
                    }
                }
                "nonce" => {
                    if claims.nonce.is_none() {
                        return Err(JwtError::InvalidClaim {
                            claim: "nonce".to_string(),
                            reason: "Nonce claim is missing".to_string(),
                            value: None,
                        });
                    }
                }
                _ => {
                    // Check if the claim exists in custom_claims
                    if !claims.custom_claims.contains_key(claim_name) {
                        return Err(JwtError::InvalidClaim {
                            claim: claim_name.clone(),
                            reason: format!("Required claim {} is missing", claim_name),
                            value: None,
                        });
                    }
                }
            }
        }

        // We can't run the custom validators directly because they expect CognitoJwtClaims
        // For now, we'll skip custom validation for OIDC tokens
        // In a future implementation, we could create a separate trait for OIDC validation

        Ok(())
    }
}

// Implement the JwtVerifier trait for OidcJwtVerifier
impl JwtVerifier for OidcJwtVerifier {
    // We've removed the generic verify method to make the trait object-safe
    // Instead, we'll use the specific methods for ID and Access tokens

    async fn verify_id_token(&self, token: &str) -> Result<Box<dyn IdTokenClaims>, JwtError> {
        // Use the generic verify method with OidcIdTokenClaims as the target type
        let claims = self.verify::<OidcIdTokenClaims>(token).await?;
        Ok(Box::new(claims))
    }

    async fn verify_access_token(
        &self,
        token: &str,
    ) -> Result<Box<dyn AccessTokenClaims>, JwtError> {
        // Use the generic verify method with OidcAccessTokenClaims as the target type
        let claims = self.verify::<OidcAccessTokenClaims>(token).await?;
        Ok(Box::new(claims))
    }
}

// Implement IdTokenClaims for OidcIdTokenClaims
impl IdTokenClaims for OidcIdTokenClaims {
    fn get_sub(&self) -> &str {
        &self.base.sub
    }

    fn get_iss(&self) -> &str {
        &self.base.iss
    }

    fn get_aud(&self) -> &str {
        &self.base.aud
    }

    fn get_exp(&self) -> u64 {
        self.base.exp
    }

    fn get_iat(&self) -> u64 {
        self.base.iat
    }

    fn get_email(&self) -> Option<&str> {
        self.email.as_deref()
    }

    fn is_email_verified(&self) -> bool {
        self.email_verified.unwrap_or(false)
    }

    fn get_name(&self) -> Option<&str> {
        self.name.as_deref()
    }
}

// Implement AccessTokenClaims for OidcAccessTokenClaims
impl AccessTokenClaims for OidcAccessTokenClaims {
    fn get_sub(&self) -> &str {
        &self.base.sub
    }

    fn get_iss(&self) -> &str {
        &self.base.iss
    }

    fn get_aud(&self) -> &str {
        &self.base.aud
    }

    fn get_exp(&self) -> u64 {
        self.base.exp
    }

    fn get_iat(&self) -> u64 {
        self.base.iat
    }

    fn get_scopes(&self) -> Vec<String> {
        match &self.scope {
            Some(scope) => scope.split_whitespace().map(|s| s.to_string()).collect(),
            None => Vec::new(),
        }
    }

    fn has_scope(&self, scope: &str) -> bool {
        self.get_scopes().contains(&scope.to_string())
    }

    fn get_client_id(&self) -> Option<&str> {
        self.client_id.as_deref()
    }
}