//! JWT Token Management System
//!
//! Provides comprehensive JWT token generation, validation, rotation, and
//! revocation using RS256 asymmetric signing for enhanced security.
//!
//! # Features
//! - Access tokens (15 minute expiry)
//! - Refresh tokens (7 day expiry)
//! - Automatic token rotation
//! - Token revocation with blacklist
//! - Thread-safe token management
//! - Cedar permissions hash integration

use std::collections::HashMap;
use std::sync::{Arc, RwLock};

use chrono::{DateTime, Duration, Utc};
use jsonwebtoken::{decode, encode, Algorithm, DecodingKey, EncodingKey, Header, Validation};
use serde::{Deserialize, Serialize};
use uuid::Uuid;

use crate::error::{auth, infrastructure, ControlCenterError, Result};

/// Token type enumeration
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum TokenType {
    /// Short-lived access token (15 minutes)
    Access,
    /// Long-lived refresh token (7 days)
    Refresh,
}

/// JWT Claims structure with comprehensive metadata
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TokenClaims {
    /// Token ID (jti - JWT ID)
    pub jti: String,

    /// Subject (user ID)
    pub sub: String,

    /// User email address
    #[serde(skip_serializing_if = "Option::is_none")]
    pub email: Option<String>,

    /// Workspace identifier
    pub workspace: String,

    /// Cedar permissions hash for quick validation
    pub permissions_hash: String,

    /// Token type (access or refresh)
    #[serde(rename = "type")]
    pub token_type: TokenType,

    /// MFA verification status
    #[serde(skip_serializing_if = "Option::is_none")]
    pub mfa_verified: Option<bool>,

    /// Issued at timestamp
    pub iat: i64,

    /// Expiration timestamp
    pub exp: i64,

    /// Issuer
    pub iss: String,

    /// Audience
    pub aud: Vec<String>,

    /// Optional metadata
    #[serde(skip_serializing_if = "Option::is_none")]
    pub metadata: Option<HashMap<String, String>>,
}

impl TokenClaims {
    /// Check if token is expired
    pub fn is_expired(&self) -> bool {
        let now = Utc::now().timestamp();
        self.exp <= now
    }

    /// Check if token is about to expire (within 5 minutes)
    pub fn needs_rotation(&self) -> bool {
        let now = Utc::now().timestamp();
        let time_until_expiry = self.exp - now;
        // Needs rotation if within 5 minutes AND not yet expired
        time_until_expiry <= 300 && time_until_expiry > 0
    }

    /// Get remaining validity duration
    pub fn remaining_validity(&self) -> std::time::Duration {
        let now = Utc::now().timestamp();
        let remaining = self.exp - now;
        std::time::Duration::from_secs(remaining.max(0) as u64)
    }
}

/// Token pair containing both access and refresh tokens
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TokenPair {
    /// Access token (short-lived)
    pub access_token: String,

    /// Refresh token (long-lived)
    pub refresh_token: String,

    /// Token type (always "Bearer")
    pub token_type: String,

    /// Access token expiry in seconds
    pub expires_in: i64,

    /// Access token expiry timestamp
    pub expires_at: DateTime<Utc>,

    /// Refresh token expiry timestamp
    pub refresh_expires_at: DateTime<Utc>,
}

/// Revoked token entry for blacklist
/// Fields used for future cleanup and audit features
#[allow(dead_code)]
#[derive(Debug, Clone)]
struct RevokedToken {
    jti: String,
    revoked_at: DateTime<Utc>,
    expires_at: DateTime<Utc>,
}

/// JWT Service for token management
pub struct JwtService {
    /// RSA encoding key (private key)
    encoding_key: EncodingKey,

    /// RSA decoding key (public key)
    decoding_key: DecodingKey,

    /// Token issuer
    issuer: String,

    /// Token audience
    audience: Vec<String>,

    /// Access token duration (15 minutes)
    access_token_duration: Duration,

    /// Refresh token duration (7 days)
    refresh_token_duration: Duration,

    /// Token blacklist (revoked tokens)
    blacklist: Arc<RwLock<HashMap<String, RevokedToken>>>,
}

impl JwtService {
    /// Create new JWT service with RSA key pair
    ///
    /// # Arguments
    /// * `private_key_pem` - RSA private key in PEM format
    /// * `public_key_pem` - RSA public key in PEM format
    /// * `issuer` - Token issuer identifier
    /// * `audience` - List of valid audiences
    ///
    /// # Example
    /// ```no_run
    /// # use control_center::auth::jwt::JwtService;
    /// let private_key = std::fs::read("private.pem").unwrap();
    /// let public_key = std::fs::read("public.pem").unwrap();
    /// let service = JwtService::new(
    ///     &private_key,
    ///     &public_key,
    ///     "control-center",
    ///     vec!["orchestrator".to_string(), "cli".to_string()]
    /// ).unwrap();
    /// ```
    pub fn new(
        private_key_pem: &[u8],
        public_key_pem: &[u8],
        issuer: impl Into<String>,
        audience: Vec<String>,
    ) -> Result<Self> {
        let encoding_key = EncodingKey::from_rsa_pem(private_key_pem).map_err(|e| {
            ControlCenterError::Auth(auth::AuthError::Authentication(format!(
                "Failed to load RSA private key: {}",
                e
            )))
        })?;

        let decoding_key = DecodingKey::from_rsa_pem(public_key_pem).map_err(|e| {
            ControlCenterError::Auth(auth::AuthError::Authentication(format!(
                "Failed to load RSA public key: {}",
                e
            )))
        })?;

        Ok(Self {
            encoding_key,
            decoding_key,
            issuer: issuer.into(),
            audience,
            access_token_duration: Duration::minutes(15),
            refresh_token_duration: Duration::days(7),
            blacklist: Arc::new(RwLock::new(HashMap::new())),
        })
    }

    /// Generate token pair (access + refresh tokens)
    ///
    /// # Arguments
    /// * `user_id` - User identifier
    /// * `workspace` - Workspace identifier
    /// * `permissions_hash` - Cedar permissions hash
    /// * `metadata` - Optional metadata
    pub fn generate_token_pair(
        &self,
        user_id: impl Into<String>,
        workspace: impl Into<String>,
        permissions_hash: impl Into<String>,
        metadata: Option<HashMap<String, String>>,
    ) -> Result<TokenPair> {
        let user_id = user_id.into();
        let workspace = workspace.into();
        let permissions_hash = permissions_hash.into();
        let now = Utc::now();

        // Generate access token
        let access_token = self.generate_token(
            user_id.clone(),
            workspace.clone(),
            permissions_hash.clone(),
            TokenType::Access,
            now,
            metadata.clone(),
        )?;

        // Generate refresh token
        let refresh_token = self.generate_token(
            user_id,
            workspace,
            permissions_hash,
            TokenType::Refresh,
            now,
            metadata,
        )?;

        let expires_at = now + self.access_token_duration;
        let refresh_expires_at = now + self.refresh_token_duration;

        Ok(TokenPair {
            access_token,
            refresh_token,
            token_type: "Bearer".to_string(),
            expires_in: self.access_token_duration.num_seconds(),
            expires_at,
            refresh_expires_at,
        })
    }

    /// Generate single token
    fn generate_token(
        &self,
        user_id: String,
        workspace: String,
        permissions_hash: String,
        token_type: TokenType,
        issued_at: DateTime<Utc>,
        metadata: Option<HashMap<String, String>>,
    ) -> Result<String> {
        let jti = Uuid::new_v4().to_string();
        let duration = match token_type {
            TokenType::Access => self.access_token_duration,
            TokenType::Refresh => self.refresh_token_duration,
        };

        let exp = issued_at + duration;

        let claims = TokenClaims {
            jti,
            sub: user_id,
            email: None, // Set by caller if needed
            workspace,
            permissions_hash,
            token_type,
            mfa_verified: None, // Set by caller if needed
            iat: issued_at.timestamp(),
            exp: exp.timestamp(),
            iss: self.issuer.clone(),
            aud: self.audience.clone(),
            metadata,
        };

        let header = Header::new(Algorithm::RS256);
        encode(&header, &claims, &self.encoding_key).map_err(|e| {
            ControlCenterError::Auth(auth::AuthError::Authentication(format!(
                "Failed to generate token: {}",
                e
            )))
        })
    }

    /// Validate and decode token
    ///
    /// Checks:
    /// - Signature validity
    /// - Expiration
    /// - Issuer and audience
    /// - Revocation status
    pub fn validate_token(&self, token: &str) -> Result<TokenClaims> {
        // Decode and validate
        let mut validation = Validation::new(Algorithm::RS256);
        validation.set_issuer(&[&self.issuer]);
        validation.set_audience(&self.audience);

        let token_data =
            decode::<TokenClaims>(token, &self.decoding_key, &validation).map_err(|e| {
                ControlCenterError::Auth(auth::AuthError::Authentication(format!(
                    "Invalid token: {}",
                    e
                )))
            })?;

        let claims = token_data.claims;

        // Check if token is revoked
        if self.is_revoked(&claims.jti)? {
            return Err(ControlCenterError::Auth(auth::AuthError::Authentication(
                "Token has been revoked".to_string(),
            )));
        }

        // Additional expiry check
        if claims.is_expired() {
            return Err(ControlCenterError::Auth(auth::AuthError::Authentication(
                "Token has expired".to_string(),
            )));
        }

        Ok(claims)
    }

    /// Rotate access token using refresh token
    ///
    /// Generates new token pair if refresh token is valid
    pub fn rotate_token(&self, refresh_token: &str) -> Result<TokenPair> {
        // Validate refresh token
        let claims = self.validate_token(refresh_token)?;

        // Ensure it's a refresh token
        if claims.token_type != TokenType::Refresh {
            return Err(ControlCenterError::Auth(auth::AuthError::Authentication(
                "Invalid token type for rotation".to_string(),
            )));
        }

        // Revoke old refresh token
        self.revoke_token(&claims.jti, claims.exp)?;

        // Generate new token pair
        self.generate_token_pair(
            claims.sub,
            claims.workspace,
            claims.permissions_hash,
            claims.metadata,
        )
    }

    /// Revoke token by adding to blacklist
    pub fn revoke_token(&self, jti: &str, expires_at: i64) -> Result<()> {
        let mut blacklist = self.blacklist.write().map_err(|e| {
            ControlCenterError::Infrastructure(infrastructure::InfrastructureError::Internal(
                format!("Failed to acquire write lock: {}", e),
            ))
        })?;

        let revoked = RevokedToken {
            jti: jti.to_string(),
            revoked_at: Utc::now(),
            expires_at: DateTime::from_timestamp(expires_at, 0).unwrap_or_else(Utc::now),
        };

        blacklist.insert(jti.to_string(), revoked);
        Ok(())
    }

    /// Check if token is revoked
    pub fn is_revoked(&self, jti: &str) -> Result<bool> {
        let blacklist = self.blacklist.read().map_err(|e| {
            ControlCenterError::Infrastructure(infrastructure::InfrastructureError::Internal(
                format!("Failed to acquire read lock: {}", e),
            ))
        })?;

        Ok(blacklist.contains_key(jti))
    }

    /// Cleanup expired tokens from blacklist
    ///
    /// Should be called periodically to prevent memory growth
    pub fn cleanup_expired_tokens(&self) -> Result<usize> {
        let mut blacklist = self.blacklist.write().map_err(|e| {
            ControlCenterError::Infrastructure(infrastructure::InfrastructureError::Internal(
                format!("Failed to acquire write lock: {}", e),
            ))
        })?;

        let now = Utc::now();
        let initial_count = blacklist.len();

        blacklist.retain(|_, token| token.expires_at > now);

        let removed_count = initial_count - blacklist.len();
        Ok(removed_count)
    }

    /// Get blacklist statistics
    pub fn blacklist_stats(&self) -> Result<BlacklistStats> {
        let blacklist = self.blacklist.read().map_err(|e| {
            ControlCenterError::Infrastructure(infrastructure::InfrastructureError::Internal(
                format!("Failed to acquire read lock: {}", e),
            ))
        })?;

        let now = Utc::now();
        let total = blacklist.len();
        let expired = blacklist
            .values()
            .filter(|token| token.expires_at <= now)
            .count();

        Ok(BlacklistStats {
            total_revoked: total,
            expired_tokens: expired,
            active_revocations: total - expired,
        })
    }

    /// Extract token from Authorization header
    ///
    /// Expects format: "Bearer <token>"
    pub fn extract_token_from_header(header: &str) -> Result<&str> {
        let parts: Vec<&str> = header.split_whitespace().collect();

        if parts.len() != 2 || parts[0] != "Bearer" {
            return Err(ControlCenterError::Auth(auth::AuthError::Authentication(
                "Invalid Authorization header format".to_string(),
            )));
        }

        Ok(parts[1])
    }
}

/// Blacklist statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BlacklistStats {
    /// Total number of revoked tokens in blacklist
    pub total_revoked: usize,

    /// Number of expired tokens (can be cleaned up)
    pub expired_tokens: usize,

    /// Number of active revocations
    pub active_revocations: usize,
}

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

    fn generate_test_keys() -> (Vec<u8>, Vec<u8>) {
        // Generate fresh RSA keys for testing
        use crate::services::jwt::generate_rsa_key_pair;

        let keys = generate_rsa_key_pair().expect("Failed to generate test RSA keys");

        (
            keys.private_key_pem.into_bytes(),
            keys.public_key_pem.into_bytes(),
        )
    }

    fn create_test_service() -> JwtService {
        let (private_key, public_key) = generate_test_keys();
        JwtService::new(
            &private_key,
            &public_key,
            "test-issuer",
            vec!["test-audience".to_string()],
        )
        .unwrap()
    }

    #[test]
    fn test_generate_token_pair() {
        let service = create_test_service();
        let result = service.generate_token_pair("user123", "workspace1", "perm_hash_abc", None);

        assert!(result.is_ok());
        let pair = result.unwrap();
        assert_eq!(pair.token_type, "Bearer");
        assert!(pair.expires_in > 0);
        assert!(!pair.access_token.is_empty());
        assert!(!pair.refresh_token.is_empty());
    }

    #[test]
    fn test_validate_token() {
        let service = create_test_service();
        let pair = service
            .generate_token_pair("user123", "workspace1", "perm_hash_abc", None)
            .unwrap();

        let claims = service.validate_token(&pair.access_token);
        assert!(claims.is_ok());

        let claims = claims.unwrap();
        assert_eq!(claims.sub, "user123");
        assert_eq!(claims.workspace, "workspace1");
        assert_eq!(claims.token_type, TokenType::Access);
    }

    #[test]
    fn test_token_revocation() {
        let service = create_test_service();
        let pair = service
            .generate_token_pair("user123", "workspace1", "perm_hash_abc", None)
            .unwrap();

        // Validate token works initially
        let claims = service.validate_token(&pair.access_token).unwrap();
        assert!(!service.is_revoked(&claims.jti).unwrap());

        // Revoke token
        service.revoke_token(&claims.jti, claims.exp).unwrap();
        assert!(service.is_revoked(&claims.jti).unwrap());

        // Validation should fail
        let result = service.validate_token(&pair.access_token);
        assert!(result.is_err());
    }

    #[test]
    fn test_token_rotation() {
        let service = create_test_service();
        let pair = service
            .generate_token_pair("user123", "workspace1", "perm_hash_abc", None)
            .unwrap();

        // Rotate using refresh token
        let new_pair = service.rotate_token(&pair.refresh_token);
        assert!(new_pair.is_ok());

        let new_pair = new_pair.unwrap();
        assert_ne!(new_pair.access_token, pair.access_token);
        assert_ne!(new_pair.refresh_token, pair.refresh_token);

        // Old refresh token should be revoked
        let result = service.validate_token(&pair.refresh_token);
        assert!(result.is_err());
    }

    #[test]
    fn test_extract_token_from_header() {
        let result = JwtService::extract_token_from_header("Bearer abc123xyz");
        assert!(result.is_ok());
        assert_eq!(result.unwrap(), "abc123xyz");

        let result = JwtService::extract_token_from_header("Invalid abc123xyz");
        assert!(result.is_err());

        let result = JwtService::extract_token_from_header("Bearer");
        assert!(result.is_err());
    }

    #[test]
    fn test_blacklist_cleanup() {
        let service = create_test_service();

        // Add some tokens to blacklist
        service
            .revoke_token("token1", Utc::now().timestamp() - 1000)
            .unwrap();
        service
            .revoke_token("token2", Utc::now().timestamp() + 1000)
            .unwrap();

        let stats = service.blacklist_stats().unwrap();
        assert_eq!(stats.total_revoked, 2);

        // Cleanup should remove expired token
        let removed = service.cleanup_expired_tokens().unwrap();
        assert_eq!(removed, 1);

        let stats = service.blacklist_stats().unwrap();
        assert_eq!(stats.total_revoked, 1);
    }

    #[test]
    fn test_claims_expiry_check() {
        let claims = TokenClaims {
            jti: "test".to_string(),
            sub: "user1".to_string(),
            email: None,
            workspace: "ws1".to_string(),
            permissions_hash: "hash".to_string(),
            token_type: TokenType::Access,
            mfa_verified: None,
            iat: Utc::now().timestamp(),
            exp: Utc::now().timestamp() - 100, // Expired
            iss: "issuer".to_string(),
            aud: vec!["audience".to_string()],
            metadata: None,
        };

        assert!(claims.is_expired());
        assert!(!claims.needs_rotation());
    }

    #[test]
    fn test_token_metadata() {
        let service = create_test_service();
        let mut metadata = HashMap::new();
        metadata.insert("ip_address".to_string(), "192.168.1.1".to_string());
        metadata.insert("user_agent".to_string(), "test-client".to_string());

        let pair = service
            .generate_token_pair(
                "user123",
                "workspace1",
                "perm_hash_abc",
                Some(metadata.clone()),
            )
            .unwrap();

        let claims = service.validate_token(&pair.access_token).unwrap();
        assert!(claims.metadata.is_some());

        let token_metadata = claims.metadata.unwrap();
        assert_eq!(
            token_metadata.get("ip_address"),
            Some(&"192.168.1.1".to_string())
        );
    }
}
