jwt-verify 0.1.3

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

use std::sync::Arc;

use crate::claims::{
    ClaimsValidator, CognitoAccessTokenClaims, CognitoIdTokenClaims, CognitoJwtClaims,
};
use crate::cognito::config::{self, TokenUse, VerifierConfig};
use crate::cognito::token::CognitoTokenParser;
use crate::common::error::{ErrorLogger, JwtError};
use crate::common::token::TokenParser;
use crate::jwk::provider::JwkProvider;
use crate::jwk::registry::{JwkProviderRegistry, RegistryError};
use crate::verifier::{AccessTokenClaims, IdTokenClaims, JwtVerifier};

/// Cognito JWT verifier
///
/// This is the main entry point for the JWT verification library. It manages multiple
/// user pools and provides methods for verifying different types of tokens.
///
/// The verifier uses issuer-based provider selection to determine which user pool
/// should verify a given token. This means that tokens must contain a valid issuer
/// claim that matches one of the registered user pools.
///
/// # Examples
///
/// ```
/// use jwt_verify::{CognitoJwtVerifier, VerifierConfig};
/// use std::time::Duration;
///
/// // Create configurations for multiple user pools
/// let config1 = VerifierConfig::new(
///     "us-east-1",
///     "us-east-1_example1",
///     &["client1".to_string()],
///     None,
/// ).unwrap();
///
/// let config2 = VerifierConfig::new(
///     "us-west-2",
///     "us-west-2_example2",
///     &["client2".to_string()],
///     None,
/// ).unwrap();
///
/// // Create a verifier with multiple user pools
/// let verifier = CognitoJwtVerifier::new(vec![config1, config2]).unwrap();
///
/// // Verify a token (the issuer in the token will determine which user pool to use)
/// let token = "..."; // JWT token string
/// let claims = verifier.verify_id_token(token).await.unwrap();
/// ```
#[derive(Debug)]
pub struct CognitoJwtVerifier {
    /// JWK provider registry
    jwk_registry: JwkProviderRegistry,
    /// Configurations for different user pools
    configs: HashMap<String, VerifierConfig>,
    /// Error logger
    error_logger: ErrorLogger,
}

impl CognitoJwtVerifier {
    /// Create a new verifier with multiple user pool configurations
    ///
    /// This constructor takes a vector of `VerifierConfig` objects, each representing
    /// a different Cognito user pool. The verifier will register all these user pools
    /// and use them for token verification based on the issuer claim in the tokens.
    ///
    /// # Parameters
    ///
    /// * `configs` - Vector of configurations for different user pools
    ///
    /// # Returns
    ///
    /// Returns a `Result` containing the new `CognitoJwtVerifier` if successful, or a `JwtError`
    /// if any of the configurations are invalid or if there's an error registering the user pools.
    ///
    /// # Examples
    ///
    /// ```
    /// use jwt_verify::{CognitoJwtVerifier, VerifierConfig};
    ///
    /// // Create configurations for multiple user pools
    /// let config1 = VerifierConfig::new(
    ///     "us-east-1",
    ///     "us-east-1_example1",
    ///     &["client1".to_string()],
    ///     None,
    /// ).unwrap();
    ///
    /// let config2 = VerifierConfig::new(
    ///     "us-west-2",
    ///     "us-west-2_example2",
    ///     &["client2".to_string()],
    ///     None,
    /// ).unwrap();
    ///
    /// // Create a verifier with multiple user pools
    /// let verifier = CognitoJwtVerifier::new(vec![config1, config2]).unwrap();
    /// ```
    pub fn new(configs: Vec<VerifierConfig>) -> Result<Self, JwtError> {
        let mut verifier = Self {
            jwk_registry: JwkProviderRegistry::new(),
            configs: HashMap::new(),
            error_logger: ErrorLogger::new(crate::common::error::ErrorVerbosity::Standard),
        };

        for config in configs {
            let id = format!("{}_{}", config.region, config.user_pool_id);
            verifier.add_user_pool(&id, config)?;
        }

        Ok(verifier)
    }

    /// Create a new verifier with a single user pool configuration
    ///
    /// This is a convenience constructor that creates a verifier with a single user pool.
    ///
    /// # Parameters
    ///
    /// * `region` - AWS region where the Cognito user pool is located (e.g., "us-east-1")
    /// * `user_pool_id` - Cognito user pool ID in the format "region_poolid"
    /// * `client_ids` - List of allowed client IDs for this user pool
    ///
    /// # Returns
    ///
    /// Returns a `Result` containing the new `CognitoJwtVerifier` if successful, or a `JwtError`
    /// if the configuration is invalid or if there's an error registering the user pool.
    ///
    /// # Examples
    ///
    /// ```
    /// use jwt_verify::CognitoJwtVerifier;
    ///
    /// // Create a verifier with a single user pool
    /// let verifier = CognitoJwtVerifier::new_single_pool(
    ///     "us-east-1",
    ///     "us-east-1_example",
    ///     &["client1".to_string()],
    /// ).unwrap();
    /// ```
    pub fn new_single_pool(
        region: &str,
        user_pool_id: &str,
        client_ids: &[String],
    ) -> Result<Self, JwtError> {
        let config = VerifierConfig::new(region, user_pool_id, client_ids, None)?;
        Self::new(vec![config])
    }

    /// Add a user pool with configuration
    ///
    /// This method adds a new user pool to the verifier. The user pool is identified by
    /// the provided ID, which should be unique among all registered user pools.
    ///
    /// # Parameters
    ///
    /// * `id` - Unique identifier for the user pool
    /// * `config` - Configuration for the user pool
    ///
    /// # Returns
    ///
    /// Returns `Ok(())` if the user pool was successfully added, or a `JwtError` if there
    /// was an error registering the user pool.
    ///
    /// # Examples
    ///
    /// ```
    /// use jwt_verify::{CognitoJwtVerifier, VerifierConfig};
    ///
    /// // Create a verifier
    /// let mut verifier = CognitoJwtVerifier::new(vec![]).unwrap();
    ///
    /// // Create a configuration for a user pool
    /// let config = VerifierConfig::new(
    ///     "us-east-1",
    ///     "us-east-1_example",
    ///     &["client1".to_string()],
    ///     None,
    /// ).unwrap();
    ///
    /// // Add the user pool to the verifier
    /// verifier.add_user_pool("my_pool", config).unwrap();
    /// ```
    pub fn add_user_pool(&mut self, id: &str, config: VerifierConfig) -> Result<(), JwtError> {
        // Create JWK provider
        let jwk_provider = JwkProvider::new(
            &config.region,
            &config.user_pool_id,
            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 a user pool with region, user pool ID, and client IDs
    ///
    /// This is a convenience method that creates a configuration for a user pool
    /// and adds it to the verifier.
    ///
    /// # Parameters
    ///
    /// * `id` - Unique identifier for the user pool
    /// * `region` - AWS region where the Cognito user pool is located (e.g., "us-east-1")
    /// * `user_pool_id` - Cognito user pool ID in the format "region_poolid"
    /// * `client_ids` - List of allowed client IDs for this user pool
    ///
    /// # Returns
    ///
    /// Returns `Ok(())` if the user pool was successfully added, or a `JwtError` if there
    /// was an error creating the configuration or registering the user pool.
    ///
    /// # Examples
    ///
    /// ```
    /// use jwt_verify::CognitoJwtVerifier;
    ///
    /// // Create a verifier
    /// let mut verifier = CognitoJwtVerifier::new(vec![]).unwrap();
    ///
    /// // Add a user pool
    /// verifier.add_user_pool_with_params(
    ///     "my_pool",
    ///     "us-east-1",
    ///     "us-east-1_example",
    ///     &["client1".to_string()],
    /// ).unwrap();
    /// ```
    pub fn add_user_pool_with_params(
        &mut self,
        id: &str,
        region: &str,
        user_pool_id: &str,
        client_ids: &[String],
    ) -> Result<(), JwtError> {
        let config = VerifierConfig::new(region, user_pool_id, client_ids, None)?;
        self.add_user_pool(id, config)
    }

    /// Get the list of registered user pool IDs
    ///
    /// This method returns a list of all user pool IDs that have been registered
    /// with the verifier.
    ///
    /// # Returns
    ///
    /// Returns a vector of user pool IDs.
    ///
    /// # Examples
    ///
    /// ```
    /// use jwt_verify::{CognitoJwtVerifier, VerifierConfig};
    ///
    /// // Create a verifier with multiple user pools
    /// let config1 = VerifierConfig::new(
    ///     "us-east-1",
    ///     "us-east-1_example1",
    ///     &["client1".to_string()],
    ///     None,
    /// ).unwrap();
    ///
    /// let config2 = VerifierConfig::new(
    ///     "us-west-2",
    ///     "us-west-2_example2",
    ///     &["client2".to_string()],
    ///     None,
    /// ).unwrap();
    ///
    /// let verifier = CognitoJwtVerifier::new(vec![config1, config2]).unwrap();
    ///
    /// // Get the list of user pool IDs
    /// let pool_ids = verifier.get_user_pool_ids();
    /// assert_eq!(pool_ids.len(), 2);
    /// ```
    pub fn get_user_pool_ids(&self) -> Vec<String> {
        self.jwk_registry.list_ids()
    }

    /// Remove a user pool
    ///
    /// This method removes a user pool from the verifier. The user pool is identified
    /// by the provided ID.
    ///
    /// # Parameters
    ///
    /// * `id` - Unique identifier for the user pool to remove
    ///
    /// # Returns
    ///
    /// Returns `Ok(())` if the user pool was successfully removed, or a `JwtError` if
    /// the user pool was not found.
    ///
    /// # Examples
    ///
    /// ```
    /// use jwt_verify::{CognitoJwtVerifier, VerifierConfig};
    ///
    /// // Create a verifier with a user pool
    /// let config = VerifierConfig::new(
    ///     "us-east-1",
    ///     "us-east-1_example",
    ///     &["client1".to_string()],
    ///     None,
    /// ).unwrap();
    ///
    /// let mut verifier = CognitoJwtVerifier::new(vec![config]).unwrap();
    ///
    /// // Remove the user pool
    /// verifier.remove_user_pool("us-east-1_us-east-1_example").unwrap();
    /// ```
    pub fn remove_user_pool(&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("pool_id".to_string()),
                error: format!("User pool '{}' not found", id),
            },
            RegistryError::JwtError(err) => err,
            other => JwtError::ConfigurationError {
                parameter: Some("pool_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::{CognitoJwtVerifier, VerifierConfig, ErrorVerbosity};
    ///
    /// // Create a verifier
    /// let mut verifier = CognitoJwtVerifier::new(vec![]).unwrap();
    ///
    /// // Set the error verbosity level
    /// verifier.set_error_verbosity(ErrorVerbosity::Detailed);
    /// ```
    pub fn set_error_verbosity(&mut self, verbosity: crate::common::error::ErrorVerbosity) {
        self.error_logger = ErrorLogger::new(verbosity);
    }

    /// Prefetch JWKs for all user pools
    ///
    /// This method prefetches JWKs for all registered user pools. This can be useful
    /// to warm up the cache before handling requests.
    ///
    /// # Returns
    ///
    /// Returns a vector of tuples containing the user pool ID and the result of the
    /// prefetch operation.
    ///
    /// # Examples
    ///
    /// ```
    /// use jwt_verify::{CognitoJwtVerifier, VerifierConfig};
    ///
    /// // Create a verifier with multiple user pools
    /// let config1 = VerifierConfig::new(
    ///     "us-east-1",
    ///     "us-east-1_example1",
    ///     &["client1".to_string()],
    ///     None,
    /// ).unwrap();
    ///
    /// let config2 = VerifierConfig::new(
    ///     "us-west-2",
    ///     "us-west-2_example2",
    ///     &["client2".to_string()],
    ///     None,
    /// ).unwrap();
    ///
    /// let verifier = CognitoJwtVerifier::new(vec![config1, config2]).unwrap();
    ///
    /// // Prefetch JWKs for all user pools
    /// 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 user pool 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::{CognitoJwtVerifier, CognitoIdTokenClaims};
    ///
    /// // Create a verifier
    /// let verifier = CognitoJwtVerifier::new_single_pool(
    ///     "us-east-1",
    ///     "us-east-1_example",
    ///     &["client1".to_string()],
    /// ).unwrap();
    ///
    /// // Verify a token
    /// let token = "..."; // JWT token string
    /// let claims = verifier.verify::<CognitoIdTokenClaims>(token).await.unwrap();
    /// ```
    pub async fn verify<T>(&self, token: &str) -> Result<T, JwtError>
    where
        T: DeserializeOwned + TryFrom<CognitoJwtClaims, Error = JwtError>,
    {
        // Parse header to get the token type
        let header = TokenParser::parse_token_header(token)?;

        // Extract issuer from token without full validation yet
        let issuer = TokenParser::extract_issuer(token)?;

        // Find the user pool ID based on the issuer
        let pool_id = self.find_pool_id_by_issuer(&issuer)?;

        // Verify with the specific pool
        self.verify_generic_with_pool::<T>(token, &pool_id).await
    }

    /// Find a user pool ID by issuer
    ///
    /// This method finds the user pool ID that matches the given issuer.
    ///
    /// # Parameters
    ///
    /// * `issuer` - The issuer to match
    ///
    /// # Returns
    ///
    /// Returns a `Result` containing the user pool ID if found, or a `JwtError`
    /// if no matching user pool is found.
    fn find_pool_id_by_issuer(&self, issuer: &str) -> Result<String, JwtError> {
        for (id, config) in &self.configs {
            let expected_issuer = format!(
                "https://cognito-idp.{}.amazonaws.com/{}",
                config.region, config.user_pool_id
            );
            if expected_issuer == issuer {
                return Ok(id.clone());
            }
        }

        Err(JwtError::InvalidIssuer {
            expected: "a registered Cognito user pool".to_string(),
            actual: issuer.to_string(),
        })
    }

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

        let config = self
            .configs
            .get(pool_id)
            .ok_or_else(|| JwtError::ConfigurationError {
                parameter: Some("pool_id".to_string()),
                error: format!("Configuration for user pool '{}' not found", pool_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: CognitoJwtClaims = 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
    fn validate_claims(
        &self,
        claims: &CognitoJwtClaims,
        config: &VerifierConfig,
    ) -> Result<(), JwtError> {
        // Check if token_use is empty
        if claims.token_use.is_empty() {
            return Err(JwtError::InvalidClaim {
                claim: "token_use".to_string(),
                reason: "Token use is empty".to_string(),
                value: None,
            });
        }

        // Validate token use against allowed token uses
        let token_use = match TokenUse::from_str(&claims.token_use) {
            Some(tu) => tu,
            None => {
                // If it's "refresh", explicitly reject
                if claims.token_use == "refresh" {
                    return Err(JwtError::UnsupportedTokenType {
                        token_type: "refresh".to_string(),
                    });
                } else {
                    // Otherwise, it's an invalid token use
                    return Err(JwtError::InvalidTokenUse {
                        expected: "id or access".to_string(),
                        actual: claims.token_use.clone(),
                    });
                }
            }
        };

        // Check if the token use is allowed for this provider
        if !config.allowed_token_uses.contains(&token_use) {
            let expected = config
                .allowed_token_uses
                .iter()
                .map(|t| t.as_str())
                .collect::<Vec<_>>()
                .join(" or ");

            return Err(JwtError::InvalidTokenUse {
                expected,
                actual: claims.token_use.clone(),
            });
        }

        // Create a claims validator and validate the claims
        let claims_validator = ClaimsValidator::new(Arc::new(config.clone()));
        claims_validator.validate_claims(claims)
    }
}

// Implement the JwtVerifier trait for CognitoJwtVerifier
impl JwtVerifier for CognitoJwtVerifier {
    // 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 CognitoIdTokenClaims as the target type
        let claims = self.verify::<CognitoIdTokenClaims>(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 CognitoAccessTokenClaims as the target type
        let claims = self.verify::<CognitoAccessTokenClaims>(token).await?;
        Ok(Box::new(claims))
    }
}

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

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

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

    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 CognitoAccessTokenClaims
impl AccessTokenClaims for CognitoAccessTokenClaims {
    fn get_sub(&self) -> &str {
        &self.base.sub
    }

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

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

    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> {
        Some(&self.base.client_id)
    }
}