jwt-verify 0.1.3

use crate::common::error::JwtError;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::HashMap;
use std::fmt;

/// Base JWT claims for all OIDC tokens
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OidcJwtClaims {
    /// Subject (user identifier)
    pub sub: String,
    /// Issuer
    pub iss: String,
    /// Audience
    #[serde(alias = "client_id")]
    pub aud: String,
    /// Expiration time
    pub exp: u64,
    /// Issued at time
    pub iat: u64,
    /// Authentication time (optional)
    pub auth_time: Option<u64>,
    /// Nonce (optional)
    pub nonce: Option<String>,
    /// Authentication Context Class Reference (optional)
    pub acr: Option<String>,
    /// Authentication Methods References (optional)
    pub amr: Option<Vec<String>>,
    /// Authorized party (optional)
    pub azp: Option<String>,
    /// Custom claims
    #[serde(flatten)]
    pub custom_claims: HashMap<String, Value>,
}

/// ID token specific claims
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OidcIdTokenClaims {
    /// Base claims
    #[serde(flatten)]
    pub base: OidcJwtClaims,
    /// Email
    pub email: Option<String>,
    /// Email verified
    pub email_verified: Option<bool>,
    /// Name
    pub name: Option<String>,
    /// Preferred username
    pub preferred_username: Option<String>,
    /// Given name
    pub given_name: Option<String>,
    /// Family name
    pub family_name: Option<String>,
    /// Locale
    pub locale: Option<String>,
    /// Picture URL
    pub picture: Option<String>,
}

/// Access token specific claims
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OidcAccessTokenClaims {
    /// Base claims
    #[serde(flatten)]
    pub base: OidcJwtClaims,
    /// Scope
    pub scope: Option<String>,
    /// Client ID
    pub client_id: Option<String>,
}

impl OidcJwtClaims {
    /// Validate that the token has the expected issuer
    pub fn validate_issuer(&self, expected_issuer: &str) -> bool {
        self.iss == expected_issuer
    }

    /// Validate that the token has one of the expected client IDs
    pub fn validate_client_id(&self, expected_client_ids: &[String]) -> bool {
        if expected_client_ids.is_empty() {
            return true;
        }

        // Check if aud is a string that matches one of the client IDs
        if expected_client_ids.contains(&self.aud) {
            return true;
        }

        // Check if azp is present and matches one of the client IDs
        if let Some(azp) = &self.azp {
            if expected_client_ids.contains(azp) {
                return true;
            }
        }

        false
    }

    /// Validate that the token has the expected token use
    pub fn validate_token_use(&self, expected_token_use: &str) -> bool {
        match self.get_custom_claim_string("token_use") {
            Some(token_use) => token_use == expected_token_use,
            None => false,
        }
    }

    /// Get a custom claim as a string
    pub fn get_custom_claim_string(&self, claim_name: &str) -> Option<String> {
        self.custom_claims
            .get(claim_name)
            .and_then(|v| v.as_str().map(|s| s.to_string()))
    }

    /// Get a custom claim as a number
    pub fn get_custom_claim_number(&self, claim_name: &str) -> Option<f64> {
        self.custom_claims.get(claim_name).and_then(|v| v.as_f64())
    }

    /// Get a custom claim as a boolean
    pub fn get_custom_claim_bool(&self, claim_name: &str) -> Option<bool> {
        self.custom_claims.get(claim_name).and_then(|v| v.as_bool())
    }

    /// Get a custom claim as an array of strings
    pub fn get_custom_claim_string_array(&self, claim_name: &str) -> Option<Vec<String>> {
        self.custom_claims.get(claim_name).and_then(|v| {
            v.as_array().map(|arr| {
                arr.iter()
                    .filter_map(|item| item.as_str().map(|s| s.to_string()))
                    .collect()
            })
        })
    }
}

impl fmt::Display for OidcJwtClaims {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "OIDC JWT Claims (sub={}, iss={}, exp={})",
            self.sub, self.iss, self.exp
        )
    }
}

impl OidcIdTokenClaims {
    /// Get the user's email if available
    pub fn get_email(&self) -> Option<&str> {
        self.email.as_deref()
    }

    /// Check if the user's email is verified
    pub fn is_email_verified(&self) -> bool {
        self.email_verified.unwrap_or(false)
    }

    /// Get the user's name if available
    pub fn get_name(&self) -> Option<&str> {
        self.name.as_deref()
    }

    /// Get the user's preferred username if available
    pub fn get_preferred_username(&self) -> Option<&str> {
        self.preferred_username.as_deref()
    }

    /// Get the user's given name if available
    pub fn get_given_name(&self) -> Option<&str> {
        self.given_name.as_deref()
    }

    /// Get the user's family name if available
    pub fn get_family_name(&self) -> Option<&str> {
        self.family_name.as_deref()
    }

    /// Get the user's locale if available
    pub fn get_locale(&self) -> Option<&str> {
        self.locale.as_deref()
    }

    /// Get the user's picture URL if available
    pub fn get_picture(&self) -> Option<&str> {
        self.picture.as_deref()
    }
}

impl fmt::Display for OidcIdTokenClaims {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "OIDC ID Token Claims (sub={}, iss={}, exp={})",
            self.base.sub, self.base.iss, self.base.exp
        )
    }
}

impl TryFrom<OidcJwtClaims> for OidcIdTokenClaims {
    type Error = JwtError;

    fn try_from(claims: OidcJwtClaims) -> Result<Self, Self::Error> {
        match claims.get_custom_claim_string("token_use") {
            Some(token_use) => {
                if token_use != "id" {
                    return Err(JwtError::InvalidTokenUse {
                        expected: "id".to_string(),
                        actual: token_use.clone(),
                    });
                }
            }
            None => {
                return Err(JwtError::InvalidTokenUse {
                    expected: "id".to_string(),
                    actual: "None".to_string(),
                });
            }
        }

        // For a proper implementation, we would deserialize the additional ID token fields
        // from the custom_claims map. For now, we'll create a minimal implementation.
        let email = claims.get_custom_claim_string("email");
        let email_verified = claims.get_custom_claim_bool("email_verified");
        let name = claims.get_custom_claim_string("name");
        let preferred_username = claims.get_custom_claim_string("preferred_username");
        let given_name = claims.get_custom_claim_string("given_name");
        let family_name = claims.get_custom_claim_string("family_name");
        let locale = claims.get_custom_claim_string("locale");
        let picture = claims.get_custom_claim_string("picture");

        Ok(Self {
            base: claims,
            email,
            email_verified,
            name,
            preferred_username,
            given_name,
            family_name,
            locale,
            picture,
        })
    }
}

impl OidcAccessTokenClaims {
    /// Get the token's scope as a list of individual scopes
    pub fn get_scopes(&self) -> Vec<String> {
        match &self.scope {
            Some(scope) => scope.split_whitespace().map(|s| s.to_string()).collect(),
            None => Vec::new(),
        }
    }

    /// Check if the token has a specific scope
    pub fn has_scope(&self, scope: &str) -> bool {
        self.get_scopes().contains(&scope.to_string())
    }

    /// Get the client ID if available
    pub fn get_client_id(&self) -> Option<&str> {
        self.client_id.as_deref()
    }
}

impl fmt::Display for OidcAccessTokenClaims {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let scope_display = match &self.scope {
            Some(s) => s,
            None => "none",
        };

        write!(
            f,
            "OIDC Access Token Claims (sub={}, iss={}, exp={}, scope={})",
            self.base.sub, self.base.iss, self.base.exp, scope_display
        )
    }
}

impl TryFrom<OidcJwtClaims> for OidcAccessTokenClaims {
    type Error = JwtError;

    fn try_from(claims: OidcJwtClaims) -> Result<Self, Self::Error> {
        match claims.get_custom_claim_string("token_use") {
            Some(token_use) => {
                if token_use != "access" {
                    return Err(JwtError::InvalidTokenUse {
                        expected: "access".to_string(),
                        actual: token_use.clone(),
                    });
                }
            }
            None => {
                return Err(JwtError::InvalidTokenUse {
                    expected: "access".to_string(),
                    actual: "None".to_string(),
                });
            }
        }
        // Extract the scope from the claims
        let scope = claims.get_custom_claim_string("scope");

        // Extract the client_id from the claims
        let client_id = claims.get_custom_claim_string("client_id");

        Ok(Self {
            base: claims,
            scope,
            client_id,
        })
    }
}