//! JWT encoding and decoding with HS256 and RS256 support
//! 
//! Provides secure JWT token generation and validation for access tokens.

use crate::error::{AuthError, Result};
use chrono::{Duration, Utc};
use jsonwebtoken::{decode, encode, Algorithm, DecodingKey, EncodingKey, Header, Validation};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// JWT claims structure
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct JwtClaims {
    /// Subject (user ID)
    pub sub: String,
    
    /// Expiration time (Unix timestamp)
    pub exp: i64,
    
    /// Issued at (Unix timestamp)
    pub iat: i64,
    
    /// Additional custom claims
    #[serde(flatten)]
    pub custom: HashMap<String, serde_json::Value>,
}

/// Encode a JWT token
/// 
/// # Arguments
/// 
/// * `payload` - Claims to include in the token
/// * `secret` - Secret key (for HS256) or private key PEM (for RS256)
/// * `exp_seconds` - Token expiration in seconds from now
/// * `algorithm` - Algorithm to use ("HS256" or "RS256")
/// 
/// # Returns
/// 
/// Encoded JWT string
/// 
/// # Security
/// 
/// - HS256: HMAC-SHA256 (symmetric, secret key)
/// - RS256: RSA-SHA256 (asymmetric, private/public key pair)
/// - Automatic expiration enforcement
/// - Issued-at timestamp included
/// 
/// # Example
/// 
/// ```rust
/// use irene_auth::encode_jwt;
/// use std::collections::HashMap;
/// 
/// let mut custom = HashMap::new();
/// custom.insert("role".to_string(), serde_json::json!("admin"));
/// 
/// let token = encode_jwt(
///     "user123",
///     &custom,
///     b"secret_key",
///     3600, // 1 hour
///     "HS256"
/// ).unwrap();
/// ```
pub fn encode_jwt(
    subject: &str,
    custom_claims: &HashMap<String, serde_json::Value>,
    secret: &[u8],
    exp_seconds: i64,
    algorithm: &str,
) -> Result<String> {
    let now = Utc::now();
    let exp = now + Duration::seconds(exp_seconds);

    let claims = JwtClaims {
        sub: subject.to_string(),
        exp: exp.timestamp(),
        iat: now.timestamp(),
        custom: custom_claims.clone(),
    };

    let (header, encoding_key) = match algorithm.to_uppercase().as_str() {
        "HS256" => {
            let header = Header::new(Algorithm::HS256);
            let key = EncodingKey::from_secret(secret);
            (header, key)
        }
        "RS256" => {
            let header = Header::new(Algorithm::RS256);
            let key = EncodingKey::from_rsa_pem(secret)
                .map_err(|e| AuthError::JwtEncodeError(format!("Invalid RS256 key: {}", e)))?;
            (header, key)
        }
        _ => return Err(AuthError::InvalidAlgorithm(algorithm.to_string())),
    };

    encode(&header, &claims, &encoding_key)
        .map_err(|e| AuthError::JwtEncodeError(e.to_string()))
}

/// Decode and validate a JWT token
/// 
/// # Arguments
/// 
/// * `token` - JWT string to decode
/// * `secret` - Secret key (for HS256) or public key PEM (for RS256)
/// * `algorithms` - List of allowed algorithms (e.g., ["HS256"], ["RS256"], or both)
/// 
/// # Returns
/// 
/// Decoded claims if valid
/// 
/// # Security
/// 
/// - Validates signature
/// - Checks expiration
/// - Enforces algorithm whitelist
/// - Returns error for any validation failure
/// 
/// # Example
/// 
/// ```rust
/// use irene_auth::{encode_jwt, decode_jwt};
/// use std::collections::HashMap;
/// 
/// let custom = HashMap::new();
/// let token = encode_jwt("user123", &custom, b"secret", 3600, "HS256").unwrap();
/// 
/// let claims = decode_jwt(&token, b"secret", &["HS256"]).unwrap();
/// assert_eq!(claims.sub, "user123");
/// ```
pub fn decode_jwt(
    token: &str,
    secret: &[u8],
    algorithms: &[&str],
) -> Result<JwtClaims> {
    // Convert algorithm strings to Algorithm enum
    let allowed_algorithms: Vec<Algorithm> = algorithms
        .iter()
        .filter_map(|alg| match alg.to_uppercase().as_str() {
            "HS256" => Some(Algorithm::HS256),
            "RS256" => Some(Algorithm::RS256),
            _ => None,
        })
        .collect();

    if allowed_algorithms.is_empty() {
        return Err(AuthError::InvalidAlgorithm(
            "No valid algorithms specified".to_string(),
        ));
    }

    let mut validation = Validation::new(allowed_algorithms[0]);
    validation.algorithms = allowed_algorithms;
    validation.validate_exp = true;

    // Try to decode with appropriate key based on first algorithm
    let decoding_key = if algorithms.contains(&"RS256") || algorithms.contains(&"rs256") {
        DecodingKey::from_rsa_pem(secret)
            .map_err(|e| AuthError::JwtDecodeError(format!("Invalid RS256 key: {}", e)))?
    } else {
        DecodingKey::from_secret(secret)
    };

    let token_data = decode::<JwtClaims>(token, &decoding_key, &validation)
        .map_err(|e| match e.kind() {
            jsonwebtoken::errors::ErrorKind::ExpiredSignature => AuthError::TokenExpired,
            jsonwebtoken::errors::ErrorKind::InvalidToken => AuthError::InvalidToken,
            _ => AuthError::JwtDecodeError(e.to_string()),
        })?;

    Ok(token_data.claims)
}

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

    #[test]
    fn test_encode_decode_hs256() {
        let secret = b"test_secret_key_for_hs256";
        let mut custom = HashMap::new();
        custom.insert("role".to_string(), serde_json::json!("admin"));
        custom.insert("email".to_string(), serde_json::json!("user@example.com"));

        // Encode
        let token = encode_jwt("user123", &custom, secret, 3600, "HS256").unwrap();

        // Decode
        let claims = decode_jwt(&token, secret, &["HS256"]).unwrap();

        assert_eq!(claims.sub, "user123");
        assert_eq!(claims.custom.get("role").unwrap(), &serde_json::json!("admin"));
        assert_eq!(claims.custom.get("email").unwrap(), &serde_json::json!("user@example.com"));
    }

    #[test]
    fn test_token_expiration() {
        let secret = b"test_secret";
        let custom = HashMap::new();

        // Create token that expires in -1 seconds (already expired)
        let token = encode_jwt("user123", &custom, secret, -1, "HS256").unwrap();

        // Try to decode expired token
        let result = decode_jwt(&token, secret, &["HS256"]);
        assert!(matches!(result, Err(AuthError::TokenExpired)));
    }

    #[test]
    fn test_invalid_signature() {
        let secret1 = b"secret1";
        let secret2 = b"secret2";
        let custom = HashMap::new();

        // Encode with one secret
        let token = encode_jwt("user123", &custom, secret1, 3600, "HS256").unwrap();

        // Try to decode with different secret
        let result = decode_jwt(&token, secret2, &["HS256"]);
        assert!(result.is_err());
    }

    #[test]
    fn test_algorithm_validation() {
        let secret = b"test_secret";
        let custom = HashMap::new();

        // Encode with HS256
        let token = encode_jwt("user123", &custom, secret, 3600, "HS256").unwrap();

        // Try to decode with only RS256 allowed (should fail)
        let result = decode_jwt(&token, secret, &["RS256"]);
        assert!(result.is_err());
    }

    #[test]
    fn test_invalid_algorithm() {
        let secret = b"secret";
        let custom = HashMap::new();

        let result = encode_jwt("user123", &custom, secret, 3600, "INVALID");
        assert!(matches!(result, Err(AuthError::InvalidAlgorithm(_))));
    }

    #[test]
    fn test_empty_custom_claims() {
        let secret = b"secret";
        let custom = HashMap::new();

        let token = encode_jwt("user123", &custom, secret, 3600, "HS256").unwrap();
        let claims = decode_jwt(&token, secret, &["HS256"]).unwrap();

        assert_eq!(claims.sub, "user123");
        assert!(claims.custom.is_empty());
    }

    #[test]
    fn test_multiple_algorithms_allowed() {
        let secret = b"secret";
        let custom = HashMap::new();

        let token = encode_jwt("user123", &custom, secret, 3600, "HS256").unwrap();
        
        // Should work when HS256 is in the allowed list
        let claims = decode_jwt(&token, secret, &["HS256", "RS256"]).unwrap();
        assert_eq!(claims.sub, "user123");
    }
}
