//! Token types for IAM
//!
//! Defines token claims for both external (JWT/OIDC) and internal tokens.

use serde::{Deserialize, Serialize};
use std::collections::HashMap;

use crate::principal::PrincipalKind;
use crate::scope::Scope;

/// Claims from an external JWT/OIDC token
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JwtClaims {
    /// Subject (user identifier from IdP)
    pub sub: String,

    /// Issuer (IdP URL)
    pub iss: String,

    /// Audience (our service)
    pub aud: Vec<String>,

    /// Expiration timestamp (Unix seconds)
    pub exp: u64,

    /// Issued at timestamp (Unix seconds)
    pub iat: u64,

    /// Not before timestamp (Unix seconds)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub nbf: Option<u64>,

    /// JWT ID (unique identifier)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub jti: Option<String>,

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

    /// Email verified flag
    #[serde(skip_serializing_if = "Option::is_none")]
    pub email_verified: Option<bool>,

    /// User's name
    #[serde(skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,

    /// Groups the user belongs to
    #[serde(default)]
    pub groups: Vec<String>,

    /// Organization ID (custom claim)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub org_id: Option<String>,

    /// Project ID (custom claim)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub project_id: Option<String>,

    /// Additional claims
    #[serde(flatten)]
    pub extra: HashMap<String, serde_json::Value>,
}

impl JwtClaims {
    /// Check if the token is expired
    pub fn is_expired(&self, now_secs: u64) -> bool {
        now_secs > self.exp
    }

    /// Check if the token is not yet valid
    pub fn is_not_yet_valid(&self, now_secs: u64) -> bool {
        self.nbf.is_some_and(|nbf| now_secs < nbf)
    }

    /// Validate basic claims
    pub fn validate(
        &self,
        expected_issuer: &str,
        expected_audience: &str,
        now_secs: u64,
    ) -> Result<(), TokenValidationError> {
        if self.iss != expected_issuer {
            return Err(TokenValidationError::InvalidIssuer {
                expected: expected_issuer.to_string(),
                actual: self.iss.clone(),
            });
        }

        if !self.aud.iter().any(|a| a == expected_audience) {
            return Err(TokenValidationError::InvalidAudience {
                expected: expected_audience.to_string(),
                actual: self.aud.clone(),
            });
        }

        if self.is_expired(now_secs) {
            return Err(TokenValidationError::Expired);
        }

        if self.is_not_yet_valid(now_secs) {
            return Err(TokenValidationError::NotYetValid);
        }

        Ok(())
    }
}

/// Token validation errors
#[derive(Debug, Clone)]
pub enum TokenValidationError {
    /// Token issuer doesn't match
    InvalidIssuer { expected: String, actual: String },
    /// Token audience doesn't match
    InvalidAudience {
        expected: String,
        actual: Vec<String>,
    },
    /// Token is expired
    Expired,
    /// Token is not yet valid
    NotYetValid,
    /// Invalid signature
    InvalidSignature,
    /// Malformed token
    Malformed(String),
}

impl std::fmt::Display for TokenValidationError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::InvalidIssuer { expected, actual } => {
                write!(f, "invalid issuer: expected {}, got {}", expected, actual)
            }
            Self::InvalidAudience { expected, actual } => {
                write!(
                    f,
                    "invalid audience: expected {}, got {:?}",
                    expected, actual
                )
            }
            Self::Expired => write!(f, "token expired"),
            Self::NotYetValid => write!(f, "token not yet valid"),
            Self::InvalidSignature => write!(f, "invalid signature"),
            Self::Malformed(msg) => write!(f, "malformed token: {}", msg),
        }
    }
}

impl std::error::Error for TokenValidationError {}

/// Internal token claims (issued by IAM service)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InternalTokenClaims {
    /// Principal ID
    pub principal_id: String,

    /// Principal kind
    pub principal_kind: PrincipalKind,

    /// Display name
    pub principal_name: String,

    /// Roles assigned to this principal
    pub roles: Vec<String>,

    /// Effective scope for this token
    pub scope: Scope,

    /// Organization ID (if applicable)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub org_id: Option<String>,

    /// Project ID (if applicable)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub project_id: Option<String>,

    /// Node ID (for service accounts bound to nodes)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub node_id: Option<String>,

    /// Issued at timestamp (TSO format)
    pub iat: u64,

    /// Expiration timestamp (TSO format)
    pub exp: u64,

    /// Unique session/token ID
    pub session_id: String,

    /// Original authentication method
    pub auth_method: AuthMethod,

    /// Additional metadata
    #[serde(default)]
    pub metadata: HashMap<String, String>,
}

impl InternalTokenClaims {
    /// Create new internal token claims
    pub fn new(
        principal_id: impl Into<String>,
        principal_kind: PrincipalKind,
        principal_name: impl Into<String>,
        scope: Scope,
        session_id: impl Into<String>,
    ) -> Self {
        Self {
            principal_id: principal_id.into(),
            principal_kind,
            principal_name: principal_name.into(),
            roles: Vec::new(),
            scope,
            org_id: None,
            project_id: None,
            node_id: None,
            iat: 0,
            exp: 0,
            session_id: session_id.into(),
            auth_method: AuthMethod::Unknown,
            metadata: HashMap::new(),
        }
    }

    /// Set roles
    pub fn with_roles(mut self, roles: Vec<String>) -> Self {
        self.roles = roles;
        self
    }

    /// Set organization ID
    pub fn with_org_id(mut self, org_id: impl Into<String>) -> Self {
        self.org_id = Some(org_id.into());
        self
    }

    /// Set project ID
    pub fn with_project_id(mut self, project_id: impl Into<String>) -> Self {
        self.project_id = Some(project_id.into());
        self
    }

    /// Set node ID
    pub fn with_node_id(mut self, node_id: impl Into<String>) -> Self {
        self.node_id = Some(node_id.into());
        self
    }

    /// Set timestamps
    pub fn with_timestamps(mut self, iat: u64, exp: u64) -> Self {
        self.iat = iat;
        self.exp = exp;
        self
    }

    /// Set auth method
    pub fn with_auth_method(mut self, method: AuthMethod) -> Self {
        self.auth_method = method;
        self
    }

    /// Check if the token is expired (TSO timestamp comparison)
    pub fn is_expired(&self, now_tso: u64) -> bool {
        now_tso > self.exp
    }

    /// Get remaining validity in TSO units
    pub fn remaining_validity(&self, now_tso: u64) -> Option<u64> {
        if now_tso >= self.exp {
            None
        } else {
            Some(self.exp - now_tso)
        }
    }
}

/// Authentication method used
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AuthMethod {
    /// JWT/OIDC authentication
    Jwt,
    /// mTLS certificate authentication
    Mtls,
    /// API key authentication
    ApiKey,
    /// Internal service-to-service
    Internal,
    /// Unknown/unspecified
    Unknown,
}

impl std::fmt::Display for AuthMethod {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            AuthMethod::Jwt => write!(f, "jwt"),
            AuthMethod::Mtls => write!(f, "mtls"),
            AuthMethod::ApiKey => write!(f, "api_key"),
            AuthMethod::Internal => write!(f, "internal"),
            AuthMethod::Unknown => write!(f, "unknown"),
        }
    }
}

/// Token type enumeration
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum TokenType {
    /// Access token for API calls
    Access,
    /// Refresh token for obtaining new access tokens
    Refresh,
    /// ID token containing user identity info
    Id,
}

/// Token metadata for storage/tracking
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TokenMetadata {
    /// Token ID (hash of the token)
    pub token_id: String,

    /// Principal who owns this token
    pub principal_id: String,

    /// Token type
    pub token_type: TokenType,

    /// Issued at (TSO)
    pub issued_at: u64,

    /// Expires at (TSO)
    pub expires_at: u64,

    /// Whether the token has been revoked
    pub revoked: bool,

    /// Revocation timestamp (TSO)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub revoked_at: Option<u64>,

    /// Revocation reason
    #[serde(skip_serializing_if = "Option::is_none")]
    pub revocation_reason: Option<String>,

    /// Source IP that requested the token
    #[serde(skip_serializing_if = "Option::is_none")]
    pub source_ip: Option<String>,

    /// User agent that requested the token
    #[serde(skip_serializing_if = "Option::is_none")]
    pub user_agent: Option<String>,
}

impl TokenMetadata {
    /// Create new token metadata
    pub fn new(
        token_id: impl Into<String>,
        principal_id: impl Into<String>,
        token_type: TokenType,
        issued_at: u64,
        expires_at: u64,
    ) -> Self {
        Self {
            token_id: token_id.into(),
            principal_id: principal_id.into(),
            token_type,
            issued_at,
            expires_at,
            revoked: false,
            revoked_at: None,
            revocation_reason: None,
            source_ip: None,
            user_agent: None,
        }
    }

    /// Mark the token as revoked
    pub fn revoke(&mut self, now_tso: u64, reason: impl Into<String>) {
        self.revoked = true;
        self.revoked_at = Some(now_tso);
        self.revocation_reason = Some(reason.into());
    }

    /// Check if the token is valid (not expired and not revoked)
    pub fn is_valid(&self, now_tso: u64) -> bool {
        !self.revoked && now_tso <= self.expires_at
    }

    /// Storage key for token metadata
    pub fn storage_key(&self) -> String {
        format!("iam/tokens/{}/{}", self.principal_id, self.token_id)
    }
}

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

    #[test]
    fn test_jwt_claims_validation() {
        let claims = JwtClaims {
            sub: "user123".into(),
            iss: "https://auth.example.com".into(),
            aud: vec!["my-service".into()],
            exp: 2000,
            iat: 1000,
            nbf: Some(1000),
            jti: None,
            email: Some("user@example.com".into()),
            email_verified: Some(true),
            name: Some("Test User".into()),
            groups: vec!["developers".into()],
            org_id: Some("org-1".into()),
            project_id: None,
            extra: HashMap::new(),
        };

        // Valid
        assert!(claims
            .validate("https://auth.example.com", "my-service", 1500)
            .is_ok());

        // Expired
        assert!(matches!(
            claims.validate("https://auth.example.com", "my-service", 2500),
            Err(TokenValidationError::Expired)
        ));

        // Wrong issuer
        assert!(matches!(
            claims.validate("https://other.com", "my-service", 1500),
            Err(TokenValidationError::InvalidIssuer { .. })
        ));

        // Wrong audience
        assert!(matches!(
            claims.validate("https://auth.example.com", "other-service", 1500),
            Err(TokenValidationError::InvalidAudience { .. })
        ));
    }

    #[test]
    fn test_internal_token_claims() {
        let claims = InternalTokenClaims::new(
            "user-123",
            PrincipalKind::User,
            "Alice",
            Scope::project("proj-1", "org-1"),
            "session-abc",
        )
        .with_roles(vec!["roles/ProjectAdmin".into()])
        .with_org_id("org-1")
        .with_timestamps(1000, 2000)
        .with_auth_method(AuthMethod::Jwt);

        assert_eq!(claims.principal_id, "user-123");
        assert_eq!(claims.roles.len(), 1);
        assert!(!claims.is_expired(1500));
        assert!(claims.is_expired(2500));
        assert_eq!(claims.remaining_validity(1500), Some(500));
        assert_eq!(claims.remaining_validity(2500), None);
    }

    #[test]
    fn test_token_metadata() {
        let mut meta =
            TokenMetadata::new("token-hash-123", "user-456", TokenType::Access, 1000, 2000);

        assert!(meta.is_valid(1500));
        assert!(!meta.is_valid(2500)); // expired

        meta.revoke(1600, "user logout");
        assert!(!meta.is_valid(1700)); // revoked
        assert!(meta.revoked);
        assert_eq!(meta.revoked_at, Some(1600));
    }

    #[test]
    fn test_auth_method_display() {
        assert_eq!(AuthMethod::Jwt.to_string(), "jwt");
        assert_eq!(AuthMethod::Mtls.to_string(), "mtls");
        assert_eq!(AuthMethod::ApiKey.to_string(), "api_key");
    }

    #[test]
    fn test_token_serialization() {
        let claims = InternalTokenClaims::new(
            "sa-compute",
            PrincipalKind::ServiceAccount,
            "Compute Agent",
            Scope::System,
            "sess-1",
        )
        .with_node_id("node-001")
        .with_auth_method(AuthMethod::Mtls);

        let json = serde_json::to_string(&claims).unwrap();
        let parsed: InternalTokenClaims = serde_json::from_str(&json).unwrap();

        assert_eq!(parsed.principal_id, "sa-compute");
        assert_eq!(parsed.node_id, Some("node-001".into()));
        assert_eq!(parsed.auth_method, AuthMethod::Mtls);
    }
}
