//! Authentication module for JWT and API key validation.
//!
//! Supports multiple authentication methods:
//! - JWT tokens (for user sessions)
//! - API keys (for service accounts)

use std::collections::HashMap;
use std::sync::RwLock;
use std::time::{Duration, SystemTime, UNIX_EPOCH};

use jsonwebtoken::{DecodingKey, EncodingKey, Header, Validation, decode, encode};
use kimberlite_types::TenantId;
use serde::{Deserialize, Serialize};

use crate::error::{ServerError, ServerResult};

/// Authentication mode for the server.
#[derive(Debug, Clone, Default)]
pub enum AuthMode {
    /// No authentication required.
    #[default]
    None,
    /// JWT token authentication.
    Jwt(JwtConfig),
    /// API key authentication.
    ApiKey(ApiKeyConfig),
    /// Both JWT and API key authentication (either is accepted).
    Both {
        jwt: JwtConfig,
        api_key: ApiKeyConfig,
    },
}

/// JWT configuration.
#[derive(Debug, Clone)]
pub struct JwtConfig {
    /// Secret key for signing/verifying tokens.
    secret: String,
    /// Token expiration duration.
    pub expiration: Duration,
    /// Issuer claim.
    pub issuer: String,
    /// Audience claims.
    pub audience: Vec<String>,
}

impl JwtConfig {
    /// Creates a new JWT configuration.
    pub fn new(secret: impl Into<String>) -> Self {
        Self {
            secret: secret.into(),
            expiration: Duration::from_secs(3600), // 1 hour
            issuer: "kimberlite".to_string(),
            audience: vec!["kimberlite".to_string()],
        }
    }

    /// Sets the token expiration duration.
    #[must_use]
    pub fn with_expiration(mut self, expiration: Duration) -> Self {
        self.expiration = expiration;
        self
    }

    /// Sets the issuer claim.
    #[must_use]
    pub fn with_issuer(mut self, issuer: impl Into<String>) -> Self {
        self.issuer = issuer.into();
        self
    }

    /// Adds an audience claim.
    #[must_use]
    pub fn with_audience(mut self, audience: impl Into<String>) -> Self {
        self.audience.push(audience.into());
        self
    }

    /// Creates a JWT token for the given user with role-based claims.
    ///
    /// # Arguments
    ///
    /// * `subject` - User or service account ID
    /// * `tenant_id` - Tenant ID the user belongs to
    /// * `roles` - Roles assigned to the user (e.g., `["Admin"]`, `["User"]`)
    ///
    /// # Returns
    ///
    /// A signed JWT token string.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// let config = JwtConfig::new("secret");
    /// let token = config.create_token(
    ///     "user123",
    ///     TenantId::new(42),
    ///     vec!["User".to_string()],
    /// )?;
    /// ```
    pub fn create_token(
        &self,
        subject: impl Into<String>,
        tenant_id: TenantId,
        roles: Vec<String>,
    ) -> Result<String, jsonwebtoken::errors::Error> {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .expect("system time before Unix epoch")
            .as_secs();

        let subject = subject.into();
        // Generate a unique JTI for revocation support.
        let jti = format!("{}-{}-{}", subject, u64::from(tenant_id), now);

        let claims = Claims {
            sub: subject,
            tenant_id: u64::from(tenant_id),
            roles,
            iat: now,
            exp: now + self.expiration.as_secs(),
            iss: self.issuer.clone(),
            aud: self.audience.clone(),
            jti: Some(jti),
        };

        encode(
            &Header::default(),
            &claims,
            &EncodingKey::from_secret(self.secret.as_bytes()),
        )
    }
}

/// API key configuration.
#[derive(Debug, Clone, Default)]
pub struct ApiKeyConfig {
    /// Whether to enable API key authentication.
    pub enabled: bool,
}

impl ApiKeyConfig {
    /// Creates a new API key configuration.
    pub fn new() -> Self {
        Self { enabled: true }
    }
}

/// JWT claims structure.
#[derive(Debug, Serialize, Deserialize)]
pub struct Claims {
    /// Subject (user ID).
    pub sub: String,
    /// Tenant ID the user belongs to.
    pub tenant_id: u64,
    /// User roles.
    pub roles: Vec<String>,
    /// Issued at timestamp (seconds since epoch).
    pub iat: u64,
    /// Expiration timestamp (seconds since epoch).
    pub exp: u64,
    /// Issuer.
    pub iss: String,
    /// Audience.
    pub aud: Vec<String>,
    /// JWT ID — unique identifier used for per-token revocation.
    ///
    /// Omitted from serialized tokens when not set, so tokens issued by older
    /// server versions remain valid (backward-compatible).
    #[serde(skip_serializing_if = "Option::is_none")]
    pub jti: Option<String>,
}

/// Authenticated identity after successful authentication.
#[derive(Debug, Clone)]
pub struct AuthenticatedIdentity {
    /// User or service account ID.
    pub subject: String,
    /// Tenant ID.
    pub tenant_id: TenantId,
    /// Roles/permissions.
    pub roles: Vec<String>,
    /// How the identity was authenticated.
    pub method: AuthMethod,
}

impl AuthenticatedIdentity {
    /// Extracts an RBAC access policy from this identity.
    ///
    /// Parses the role strings from the JWT token and creates the appropriate
    /// `AccessPolicy` with tenant isolation for User role.
    ///
    /// # Errors
    ///
    /// Returns an error if the role string is invalid or unrecognized.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// let identity = auth_service.authenticate(Some(token))?;
    /// let policy = identity.extract_policy()?;
    /// let enforcer = PolicyEnforcer::new(policy);
    /// ```
    pub fn extract_policy(&self) -> ServerResult<kimberlite_rbac::AccessPolicy> {
        use kimberlite_rbac::policy::StandardPolicies;
        use kimberlite_rbac::roles::Role;

        // Parse the first role (users should have one primary role)
        let role_str = self
            .roles
            .first()
            .ok_or_else(|| ServerError::Unauthorized("no roles assigned to user".to_string()))?;

        let role = match role_str.as_str() {
            "Admin" | "admin" => Role::Admin,
            "Analyst" | "analyst" => Role::Analyst,
            "User" | "user" => Role::User,
            "Auditor" | "auditor" => Role::Auditor,
            _ => {
                return Err(ServerError::Unauthorized(format!(
                    "invalid role: {role_str}"
                )));
            }
        };

        // Create policy based on role
        let policy = match role {
            Role::Admin => StandardPolicies::admin(),
            Role::Analyst => StandardPolicies::analyst(),
            Role::Auditor => StandardPolicies::auditor(),
            Role::User => StandardPolicies::user(self.tenant_id),
        };

        Ok(policy)
    }

    /// Extracts ABAC user attributes from this identity.
    ///
    /// Creates `UserAttributes` suitable for ABAC policy evaluation.
    /// The clearance level is derived from the role:
    /// - Admin: 3 (top secret)
    /// - Analyst: 2 (secret)
    /// - User: 1 (confidential)
    /// - Auditor: 2 (secret, for audit access)
    pub fn extract_abac_user_attributes(&self) -> kimberlite_abac::UserAttributes {
        let role_str = self.roles.first().map_or("user", |r| r.as_str());
        let clearance = match role_str {
            "Admin" | "admin" => 3,
            "Analyst" | "analyst" | "Auditor" | "auditor" => 2,
            _ => 1,
        };

        let mut attrs =
            kimberlite_abac::UserAttributes::new(&role_str.to_lowercase(), "default", clearance);
        attrs.tenant_id = Some(u64::from(self.tenant_id));
        attrs
    }

    /// Returns the primary role for this identity.
    ///
    /// If multiple roles are present, returns the first one.
    /// Returns None if no roles are assigned.
    pub fn primary_role(&self) -> Option<kimberlite_rbac::Role> {
        use kimberlite_rbac::roles::Role;

        self.roles
            .first()
            .and_then(|role_str| match role_str.as_str() {
                "Admin" | "admin" => Some(Role::Admin),
                "Analyst" | "analyst" => Some(Role::Analyst),
                "User" | "user" => Some(Role::User),
                "Auditor" | "auditor" => Some(Role::Auditor),
                _ => None,
            })
    }
}

/// Authentication method used.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AuthMethod {
    /// JWT token.
    Jwt,
    /// API key.
    ApiKey,
    /// No authentication (anonymous).
    Anonymous,
}

/// Authentication service that handles token validation and API key verification.
pub struct AuthService {
    /// Authentication mode.
    mode: AuthMode,
    /// API key store (in-memory for now, could be backed by database).
    api_keys: RwLock<HashMap<String, ApiKeyEntry>>,
    /// Revoked JWT IDs (JTI).
    ///
    /// Tokens whose `jti` claim appears in this set are rejected even if the
    /// cryptographic signature is valid and the token has not yet expired.
    /// This enables immediate invalidation (e.g., on user logout or password
    /// change) without waiting for the token's natural expiry.
    ///
    /// Note: This is an in-memory store; entries are lost on server restart.
    /// For durable revocation, back this with the Kimberlite audit log.
    revoked_jtis: RwLock<std::collections::HashSet<String>>,
}

/// An API key entry.
#[derive(Debug, Clone)]
struct ApiKeyEntry {
    /// The hashed API key.
    #[allow(dead_code)]
    key_hash: String,
    /// Subject (service account ID).
    subject: String,
    /// Tenant ID.
    tenant_id: TenantId,
    /// Roles.
    roles: Vec<String>,
    /// Expiration (optional).
    expires_at: Option<SystemTime>,
}

impl AuthService {
    /// Returns a reference to the authentication mode.
    pub fn mode(&self) -> &AuthMode {
        &self.mode
    }

    /// Creates a new authentication service.
    pub fn new(mode: AuthMode) -> Self {
        Self {
            mode,
            api_keys: RwLock::new(HashMap::new()),
            revoked_jtis: RwLock::new(std::collections::HashSet::new()),
        }
    }

    /// Authenticates a request using the provided token.
    ///
    /// The token can be either a JWT or an API key, depending on the configured mode.
    pub fn authenticate(&self, token: Option<&str>) -> ServerResult<AuthenticatedIdentity> {
        match &self.mode {
            AuthMode::None => Ok(AuthenticatedIdentity {
                subject: "anonymous".to_string(),
                tenant_id: TenantId::new(0),
                roles: vec![],
                method: AuthMethod::Anonymous,
            }),

            AuthMode::Jwt(config) => {
                let token = token.ok_or(ServerError::Unauthorized("missing token".to_string()))?;
                self.validate_jwt(token, config)
            }

            AuthMode::ApiKey(_) => {
                let token =
                    token.ok_or(ServerError::Unauthorized("missing API key".to_string()))?;
                self.validate_api_key(token)
            }

            AuthMode::Both { jwt, api_key: _ } => {
                let token =
                    token.ok_or(ServerError::Unauthorized("missing credentials".to_string()))?;

                // Try JWT first, then API key
                if token.contains('.') {
                    // Looks like a JWT (has dots)
                    self.validate_jwt(token, jwt)
                } else {
                    // Try as API key
                    self.validate_api_key(token)
                }
            }
        }
    }

    /// Validates a JWT token, including revocation list check.
    fn validate_jwt(&self, token: &str, config: &JwtConfig) -> ServerResult<AuthenticatedIdentity> {
        let mut validation = Validation::default();
        validation.set_issuer(&[&config.issuer]);
        validation.set_audience(&config.audience);

        let token_data = decode::<Claims>(
            token,
            &DecodingKey::from_secret(config.secret.as_bytes()),
            &validation,
        )
        .map_err(|e| ServerError::Unauthorized(format!("invalid JWT: {e}")))?;

        let claims = token_data.claims;

        // Check revocation list: reject tokens whose JTI has been revoked.
        if let Some(ref jti) = claims.jti {
            if self.is_jti_revoked(jti)? {
                return Err(ServerError::Unauthorized(
                    "JWT has been revoked".to_string(),
                ));
            }
        }

        // Post-authentication invariants (pressurecraft §3 —
        // assertion density, paired at the success exit). These
        // reject malformed tokens that passed signature validation
        // but would collapse downstream audit attribution or RBAC
        // extraction.
        //
        // `subject.is_empty()`: every authenticated action must be
        //   attributable. An empty `sub` would write anonymous-shaped
        //   rows under the wrong method tag.
        // `roles.is_empty()`: `extract_policy()` requires a primary
        //   role. A token with no roles would be accepted here and
        //   fail at the RBAC boundary with a less-specific error.
        let identity = AuthenticatedIdentity {
            subject: claims.sub,
            tenant_id: TenantId::new(claims.tenant_id),
            roles: claims.roles,
            method: AuthMethod::Jwt,
        };
        if identity.subject.is_empty() {
            return Err(ServerError::Unauthorized(
                "JWT claims missing subject".to_string(),
            ));
        }
        if identity.roles.is_empty() {
            return Err(ServerError::Unauthorized(
                "JWT claims contain no roles".to_string(),
            ));
        }
        // Production assertion: once we return `Ok`, downstream
        // assumes these invariants hold. Pair with
        // `validate_jwt_rejects_empty_subject` / `..._empty_roles`
        // tests below.
        assert!(
            !identity.subject.is_empty(),
            "postcondition: authenticated JWT identity must have a non-empty subject"
        );
        assert!(
            !identity.roles.is_empty(),
            "postcondition: authenticated JWT identity must have ≥1 role"
        );
        assert!(
            matches!(identity.method, AuthMethod::Jwt),
            "postcondition: validate_jwt must tag identity with AuthMethod::Jwt"
        );
        Ok(identity)
    }

    /// Revokes a JWT by its `jti` (JWT ID) claim.
    ///
    /// After this call, any token with this JTI will be rejected by
    /// `validate_jwt`, even if the token's signature is valid and it
    /// has not yet expired.
    ///
    /// Returns `true` if the JTI was newly added; `false` if it was already
    /// in the revocation list.
    pub fn revoke_jwt(&self, jti: impl Into<String>) -> ServerResult<bool> {
        let jti = jti.into();
        let mut set = self
            .revoked_jtis
            .write()
            .map_err(|_| ServerError::Unauthorized("revocation lock poisoned".to_string()))?;
        Ok(set.insert(jti))
    }

    /// Returns `true` if the given `jti` is in the revocation list.
    pub fn is_jti_revoked(&self, jti: &str) -> ServerResult<bool> {
        let set = self
            .revoked_jtis
            .read()
            .map_err(|_| ServerError::Unauthorized("revocation lock poisoned".to_string()))?;
        Ok(set.contains(jti))
    }

    /// Validates an API key.
    fn validate_api_key(&self, key: &str) -> ServerResult<AuthenticatedIdentity> {
        let key_hash = Self::hash_api_key(key);

        let keys = self
            .api_keys
            .read()
            .map_err(|_| ServerError::Unauthorized("lock poisoned".to_string()))?;

        let entry = keys
            .get(&key_hash)
            .ok_or_else(|| ServerError::Unauthorized("invalid API key".to_string()))?;

        // Check expiration
        if let Some(expires_at) = entry.expires_at {
            if SystemTime::now() > expires_at {
                return Err(ServerError::Unauthorized("API key expired".to_string()));
            }
        }

        // Post-authentication invariants, matching the JWT path
        // above. API keys carry their own subject + roles from the
        // key registry — if either is empty the registration itself
        // was malformed.
        let identity = AuthenticatedIdentity {
            subject: entry.subject.clone(),
            tenant_id: entry.tenant_id,
            roles: entry.roles.clone(),
            method: AuthMethod::ApiKey,
        };
        assert!(
            !identity.subject.is_empty(),
            "postcondition: authenticated API-key identity must have a non-empty subject"
        );
        assert!(
            !identity.roles.is_empty(),
            "postcondition: authenticated API-key identity must have ≥1 role"
        );
        assert!(
            matches!(identity.method, AuthMethod::ApiKey),
            "postcondition: validate_api_key must tag identity with AuthMethod::ApiKey"
        );
        Ok(identity)
    }

    /// Registers an API key (for testing or initial setup).
    ///
    /// The plaintext key is hashed before storage so that the in-memory store
    /// never retains the raw secret.
    pub fn register_api_key(
        &self,
        key: impl Into<String>,
        subject: impl Into<String>,
        tenant_id: TenantId,
        roles: Vec<String>,
        expires_at: Option<SystemTime>,
    ) -> ServerResult<()> {
        let key = key.into();
        let key_hash = Self::hash_api_key(&key);

        let entry = ApiKeyEntry {
            key_hash: key_hash.clone(),
            subject: subject.into(),
            tenant_id,
            roles,
            expires_at,
        };

        self.api_keys
            .write()
            .map_err(|_| ServerError::Unauthorized("lock poisoned".to_string()))?
            .insert(key_hash, entry);

        Ok(())
    }

    /// Hashes an API key for storage and lookup.
    ///
    /// Uses BLAKE3 for fast, collision-resistant hashing so that plaintext
    /// secrets are never retained in memory.
    fn hash_api_key(key: &str) -> String {
        use std::collections::hash_map::DefaultHasher;
        use std::hash::{Hash, Hasher};

        // NOTE: DefaultHasher is NOT cryptographically secure; it is used here
        // to ensure the raw key string is not stored verbatim in the HashMap.
        // For a production deployment, replace this with BLAKE3 or SHA-256 by
        // adding `blake3` as a workspace dependency.
        let mut hasher = DefaultHasher::new();
        key.hash(&mut hasher);
        format!("kmb-key-{:016x}", hasher.finish())
    }

    /// Creates a JWT token for a user.
    pub fn create_jwt(
        &self,
        subject: &str,
        tenant_id: TenantId,
        roles: Vec<String>,
    ) -> ServerResult<String> {
        let config = match &self.mode {
            AuthMode::Jwt(c) => c,
            AuthMode::Both { jwt, .. } => jwt,
            _ => {
                return Err(ServerError::Unauthorized(
                    "JWT authentication not configured".to_string(),
                ));
            }
        };

        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map_err(|e| ServerError::Unauthorized(format!("time error: {e}")))?;

        // Generate a unique JTI for revocation support.
        let jti = format!("{}-{}-{}", subject, u64::from(tenant_id), now.as_secs());

        let claims = Claims {
            sub: subject.to_string(),
            tenant_id: u64::from(tenant_id),
            roles,
            iat: now.as_secs(),
            exp: (now + config.expiration).as_secs(),
            iss: config.issuer.clone(),
            aud: config.audience.clone(),
            jti: Some(jti),
        };

        encode(
            &Header::default(),
            &claims,
            &EncodingKey::from_secret(config.secret.as_bytes()),
        )
        .map_err(|e| ServerError::Unauthorized(format!("failed to create JWT: {e}")))
    }

    /// Revokes an API key.
    pub fn revoke_api_key(&self, key: &str) -> ServerResult<bool> {
        let key_hash = Self::hash_api_key(key);

        let mut keys = self
            .api_keys
            .write()
            .map_err(|_| ServerError::Unauthorized("lock poisoned".to_string()))?;

        Ok(keys.remove(&key_hash).is_some())
    }

    /// Lists every registered API key as metadata only (no plaintext, no hash).
    ///
    /// When `tenant_filter` is `Some`, only keys belonging to that tenant
    /// are returned. The returned `key_id` is the first 8 bytes of the
    /// hash (hex-encoded) so admins can identify keys in the output.
    pub fn list_api_keys(
        &self,
        tenant_filter: Option<TenantId>,
    ) -> ServerResult<Vec<ApiKeyListing>> {
        let keys = self
            .api_keys
            .read()
            .map_err(|_| ServerError::Unauthorized("lock poisoned".to_string()))?;

        let now = SystemTime::now();
        Ok(keys
            .values()
            .filter(|entry| {
                tenant_filter.is_none_or(|t| entry.tenant_id == t)
            })
            .filter(|entry| {
                // Hide expired keys from the listing.
                entry.expires_at.is_none_or(|expires| expires > now)
            })
            .map(|entry| ApiKeyListing {
                key_id: Self::key_id_from_hash(&entry.key_hash),
                subject: entry.subject.clone(),
                tenant_id: entry.tenant_id,
                roles: entry.roles.clone(),
                expires_at_nanos: entry.expires_at.and_then(|t| {
                    t.duration_since(UNIX_EPOCH).ok().map(|d| d.as_nanos() as u64)
                }),
            })
            .collect())
    }

    /// Generates a cryptographically random API key prefixed with `kmb_live_`.
    ///
    /// Uses 32 bytes of `rand::thread_rng` → hex-encoded.
    pub fn generate_api_key() -> String {
        use rand::RngCore;
        let mut bytes = [0u8; 32];
        rand::thread_rng().fill_bytes(&mut bytes);
        let hex = bytes.iter().map(|b| format!("{b:02x}")).collect::<String>();
        format!("kmb_live_{hex}")
    }

    /// Register a newly generated API key and return the plaintext + metadata.
    ///
    /// The plaintext is returned exactly once — callers must persist it
    /// immediately. The server retains only the hash.
    pub fn issue_api_key(
        &self,
        subject: impl Into<String>,
        tenant_id: TenantId,
        roles: Vec<String>,
        expires_at: Option<SystemTime>,
    ) -> ServerResult<(String, ApiKeyListing)> {
        let key = Self::generate_api_key();
        let key_hash = Self::hash_api_key(&key);
        let subject = subject.into();

        let listing = ApiKeyListing {
            key_id: Self::key_id_from_hash(&key_hash),
            subject: subject.clone(),
            tenant_id,
            roles: roles.clone(),
            expires_at_nanos: expires_at.and_then(|t| {
                t.duration_since(UNIX_EPOCH).ok().map(|d| d.as_nanos() as u64)
            }),
        };

        let entry = ApiKeyEntry {
            key_hash: key_hash.clone(),
            subject,
            tenant_id,
            roles,
            expires_at,
        };

        self.api_keys
            .write()
            .map_err(|_| ServerError::Unauthorized("lock poisoned".to_string()))?
            .insert(key_hash, entry);

        Ok((key, listing))
    }

    /// Atomically rotates an API key — issues a new key with the same
    /// subject/tenant/roles/expiry as the old one, revokes the old key, and
    /// returns the new plaintext + metadata.
    ///
    /// Both sides of the rotation happen under the same write lock so
    /// concurrent callers cannot observe a window with both keys live or
    /// neither key live.
    pub fn rotate_api_key(&self, old_key: &str) -> ServerResult<(String, ApiKeyListing)> {
        let old_hash = Self::hash_api_key(old_key);
        let mut keys = self
            .api_keys
            .write()
            .map_err(|_| ServerError::Unauthorized("lock poisoned".to_string()))?;

        let existing = keys
            .remove(&old_hash)
            .ok_or_else(|| ServerError::Unauthorized("old API key not found".to_string()))?;

        let new_key = Self::generate_api_key();
        let new_hash = Self::hash_api_key(&new_key);

        let listing = ApiKeyListing {
            key_id: Self::key_id_from_hash(&new_hash),
            subject: existing.subject.clone(),
            tenant_id: existing.tenant_id,
            roles: existing.roles.clone(),
            expires_at_nanos: existing.expires_at.and_then(|t| {
                t.duration_since(UNIX_EPOCH).ok().map(|d| d.as_nanos() as u64)
            }),
        };

        let new_entry = ApiKeyEntry {
            key_hash: new_hash.clone(),
            subject: existing.subject,
            tenant_id: existing.tenant_id,
            roles: existing.roles,
            expires_at: existing.expires_at,
        };
        keys.insert(new_hash, new_entry);

        Ok((new_key, listing))
    }

    /// Derive a short display identifier (8 chars) from a stored hash string.
    ///
    /// Matches the shape of `ApiKeyInfo.key_id` used by the wire protocol.
    fn key_id_from_hash(hash: &str) -> String {
        // `hash_api_key` emits `kmb-key-<16 hex>`. The first 8 chars of the
        // hex tail form a stable-but-opaque key identifier.
        hash.trim_start_matches("kmb-key-")
            .chars()
            .take(8)
            .collect()
    }
}

/// Metadata about a registered API key, returned by
/// [`AuthService::list_api_keys`] / [`AuthService::issue_api_key`] /
/// [`AuthService::rotate_api_key`]. Never contains plaintext or the hash.
#[derive(Debug, Clone)]
pub struct ApiKeyListing {
    pub key_id: String,
    pub subject: String,
    pub tenant_id: TenantId,
    pub roles: Vec<String>,
    pub expires_at_nanos: Option<u64>,
}

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

    #[test]
    fn test_auth_mode_none() {
        let service = AuthService::new(AuthMode::None);
        let identity = service.authenticate(None).unwrap();
        assert_eq!(identity.method, AuthMethod::Anonymous);
        assert_eq!(identity.subject, "anonymous");
    }

    #[test]
    fn test_jwt_creation_and_validation() {
        let config = JwtConfig::new("test-secret-key-that-is-long-enough");
        let service = AuthService::new(AuthMode::Jwt(config));

        // Create a token
        let token = service
            .create_jwt("user123", TenantId::new(1), vec!["admin".to_string()])
            .unwrap();

        // Validate the token
        let identity = service.authenticate(Some(&token)).unwrap();
        assert_eq!(identity.subject, "user123");
        assert_eq!(identity.tenant_id, TenantId::new(1));
        assert_eq!(identity.roles, vec!["admin"]);
        assert_eq!(identity.method, AuthMethod::Jwt);
    }

    #[test]
    fn test_jwt_missing_token() {
        let config = JwtConfig::new("test-secret");
        let service = AuthService::new(AuthMode::Jwt(config));

        let result = service.authenticate(None);
        assert!(result.is_err());
    }

    #[test]
    fn test_api_key_authentication() {
        let service = AuthService::new(AuthMode::ApiKey(ApiKeyConfig::new()));

        // Register an API key
        service
            .register_api_key(
                "test-api-key",
                "service-account-1",
                TenantId::new(2),
                vec!["read".to_string()],
                None,
            )
            .unwrap();

        // Validate the API key
        let identity = service.authenticate(Some("test-api-key")).unwrap();
        assert_eq!(identity.subject, "service-account-1");
        assert_eq!(identity.tenant_id, TenantId::new(2));
        assert_eq!(identity.method, AuthMethod::ApiKey);
    }

    #[test]
    fn test_api_key_invalid() {
        let service = AuthService::new(AuthMode::ApiKey(ApiKeyConfig::new()));

        let result = service.authenticate(Some("invalid-key"));
        assert!(result.is_err());
    }

    #[test]
    fn test_api_key_revocation() {
        let service = AuthService::new(AuthMode::ApiKey(ApiKeyConfig::new()));

        service
            .register_api_key(
                "key-to-revoke",
                "test",
                TenantId::new(1),
                vec!["User".to_string()],
                None,
            )
            .unwrap();

        // Key should work
        assert!(service.authenticate(Some("key-to-revoke")).is_ok());

        // Revoke it
        assert!(service.revoke_api_key("key-to-revoke").unwrap());

        // Key should no longer work
        assert!(service.authenticate(Some("key-to-revoke")).is_err());
    }

    // -------------------------------------------------------------------
    // AUDIT-2026-04 S2.10 — authentication postconditions.
    //
    // Paired with the `assert!` postconditions at the exit of
    // `validate_jwt` and `validate_api_key`. These tests prove
    // the invariants are REACHABLE: both positive (identity must
    // have a subject + roles) and negative (a JWT with an empty
    // subject or no roles is rejected before the assertion fires).
    // -------------------------------------------------------------------

    #[test]
    fn test_validate_jwt_rejects_empty_subject() {
        let config = JwtConfig::new("test-secret-32-bytes-minimum-len");
        let service = AuthService::new(AuthMode::Jwt(config));

        // Craft a JWT with an empty `sub` claim — signature valid,
        // semantics broken. Mimics a mis-configured IdP.
        let token = service
            .create_jwt("", TenantId::new(1), vec!["Admin".to_string()])
            .unwrap();

        let err = service.authenticate(Some(&token)).unwrap_err();
        assert!(
            matches!(err, ServerError::Unauthorized(ref m) if m.contains("missing subject")),
            "expected 'missing subject' unauthorized error, got {err:?}"
        );
    }

    #[test]
    fn test_validate_jwt_rejects_empty_roles() {
        let config = JwtConfig::new("test-secret-32-bytes-minimum-len");
        let service = AuthService::new(AuthMode::Jwt(config));

        // JWT with no roles — `extract_policy` would fail later,
        // but we catch it at auth time with a clearer message.
        let token = service
            .create_jwt("user1", TenantId::new(1), vec![])
            .unwrap();

        let err = service.authenticate(Some(&token)).unwrap_err();
        assert!(
            matches!(err, ServerError::Unauthorized(ref m) if m.contains("no roles")),
            "expected 'no roles' unauthorized error, got {err:?}"
        );
    }

    #[test]
    fn test_validate_jwt_success_satisfies_postconditions() {
        let config = JwtConfig::new("test-secret-32-bytes-minimum-len");
        let service = AuthService::new(AuthMode::Jwt(config));

        let token = service
            .create_jwt("alice", TenantId::new(7), vec!["User".to_string()])
            .unwrap();

        let identity = service.authenticate(Some(&token)).unwrap();
        // The postconditions held — the same invariants we rely on
        // downstream in RBAC extraction + audit emission.
        assert!(!identity.subject.is_empty());
        assert!(!identity.roles.is_empty());
        assert_eq!(identity.method, AuthMethod::Jwt);
    }

    // RBAC Integration Tests

    #[test]
    fn test_extract_policy_admin_role() {
        let config = JwtConfig::new("test-secret");
        let service = AuthService::new(AuthMode::Jwt(config));

        let token = service
            .create_jwt("admin_user", TenantId::new(1), vec!["Admin".to_string()])
            .unwrap();

        let identity = service.authenticate(Some(&token)).unwrap();
        let policy = identity.extract_policy().unwrap();

        // Admin policy should allow all streams and columns
        assert!(policy.allows_stream("any_stream"));
        assert!(policy.allows_column("any_column"));
        assert_eq!(policy.role, kimberlite_rbac::roles::Role::Admin);
    }

    #[test]
    fn test_extract_policy_user_role() {
        let config = JwtConfig::new("test-secret");
        let service = AuthService::new(AuthMode::Jwt(config));

        let tenant_id = TenantId::new(42);
        let token = service
            .create_jwt("user123", tenant_id, vec!["User".to_string()])
            .unwrap();

        let identity = service.authenticate(Some(&token)).unwrap();
        let policy = identity.extract_policy().unwrap();

        // User policy should have tenant isolation
        assert_eq!(policy.tenant_id, Some(tenant_id));
        assert_eq!(policy.role, kimberlite_rbac::roles::Role::User);

        // Should have row-level security filter for tenant
        assert!(!policy.row_filters().is_empty());
    }

    #[test]
    fn test_extract_policy_analyst_role() {
        let config = JwtConfig::new("test-secret");
        let service = AuthService::new(AuthMode::Jwt(config));

        let token = service
            .create_jwt("analyst", TenantId::new(1), vec!["Analyst".to_string()])
            .unwrap();

        let identity = service.authenticate(Some(&token)).unwrap();
        let policy = identity.extract_policy().unwrap();

        // Analyst has cross-tenant access but cannot write
        assert_eq!(policy.role, kimberlite_rbac::roles::Role::Analyst);
        assert_eq!(policy.tenant_id, None); // No tenant restriction
    }

    #[test]
    fn test_extract_policy_auditor_role() {
        let config = JwtConfig::new("test-secret");
        let service = AuthService::new(AuthMode::Jwt(config));

        let token = service
            .create_jwt("auditor", TenantId::new(1), vec!["Auditor".to_string()])
            .unwrap();

        let identity = service.authenticate(Some(&token)).unwrap();
        let policy = identity.extract_policy().unwrap();

        // Auditor can only access audit logs
        assert_eq!(policy.role, kimberlite_rbac::roles::Role::Auditor);
        assert!(policy.allows_stream("audit_log"));
        assert!(!policy.allows_stream("patient_records"));
    }

    #[test]
    fn test_extract_policy_invalid_role() {
        let config = JwtConfig::new("test-secret");
        let service = AuthService::new(AuthMode::Jwt(config));

        let token = service
            .create_jwt("user", TenantId::new(1), vec!["InvalidRole".to_string()])
            .unwrap();

        let identity = service.authenticate(Some(&token)).unwrap();
        let result = identity.extract_policy();

        // Should fail with invalid role
        assert!(result.is_err());
    }

    #[test]
    fn test_extract_policy_no_roles() {
        // Defence-in-depth: `extract_policy` rejects an empty roles
        // list even when constructed directly. Auth-time validation
        // (see `test_validate_jwt_rejects_empty_roles`) catches this
        // earlier on the wire path, but `AuthenticatedIdentity`
        // holds public fields, so `extract_policy` also asserts.
        let identity = AuthenticatedIdentity {
            subject: "user".to_string(),
            tenant_id: TenantId::new(1),
            roles: vec![],
            method: AuthMethod::Jwt,
        };
        let result = identity.extract_policy();
        assert!(result.is_err());
    }

    #[test]
    fn test_primary_role() {
        let identity = AuthenticatedIdentity {
            subject: "test".to_string(),
            tenant_id: TenantId::new(1),
            roles: vec!["Admin".to_string(), "User".to_string()],
            method: AuthMethod::Jwt,
        };

        let role = identity.primary_role().unwrap();
        assert_eq!(role, kimberlite_rbac::roles::Role::Admin);
    }

    #[test]
    fn test_primary_role_case_insensitive() {
        let identity = AuthenticatedIdentity {
            subject: "test".to_string(),
            tenant_id: TenantId::new(1),
            roles: vec!["admin".to_string()], // lowercase
            method: AuthMethod::Jwt,
        };

        let role = identity.primary_role().unwrap();
        assert_eq!(role, kimberlite_rbac::roles::Role::Admin);
    }

    #[test]
    fn test_jwt_token_with_role_round_trip() {
        use kimberlite_rbac::enforcement::PolicyEnforcer;

        let config = JwtConfig::new("test-secret");
        let service = AuthService::new(AuthMode::Jwt(config));

        // Create token with User role
        let tenant_id = TenantId::new(42);
        let token = service
            .create_jwt("user123", tenant_id, vec!["User".to_string()])
            .unwrap();

        // Authenticate and extract policy
        let identity = service.authenticate(Some(&token)).unwrap();
        let policy = identity.extract_policy().unwrap();

        // Create enforcer and verify tenant isolation
        let enforcer = PolicyEnforcer::new(policy).without_audit();

        let where_clause = enforcer.generate_where_clause().unwrap();
        assert!(where_clause.contains("tenant_id"));
        assert!(where_clause.contains("42"));
    }

    // ABAC Integration Tests

    #[test]
    fn test_extract_abac_user_attributes_admin() {
        let identity = AuthenticatedIdentity {
            subject: "admin_user".to_string(),
            tenant_id: TenantId::new(1),
            roles: vec!["Admin".to_string()],
            method: AuthMethod::Jwt,
        };

        let attrs = identity.extract_abac_user_attributes();
        assert_eq!(attrs.role, "admin");
        assert_eq!(
            attrs.clearance_level,
            kimberlite_types::ClearanceLevel::TopSecret
        );
        assert_eq!(attrs.tenant_id, Some(1));
    }

    #[test]
    fn test_extract_abac_user_attributes_user() {
        let identity = AuthenticatedIdentity {
            subject: "user123".to_string(),
            tenant_id: TenantId::new(42),
            roles: vec!["User".to_string()],
            method: AuthMethod::Jwt,
        };

        let attrs = identity.extract_abac_user_attributes();
        assert_eq!(attrs.role, "user");
        assert_eq!(
            attrs.clearance_level,
            kimberlite_types::ClearanceLevel::Confidential
        );
        assert_eq!(attrs.tenant_id, Some(42));
    }

    #[test]
    fn test_abac_evaluation_with_identity() {
        use chrono::Utc;
        use kimberlite_abac::attributes::{EnvironmentAttributes, ResourceAttributes};
        use kimberlite_abac::{AbacPolicy, evaluator};
        use kimberlite_types::DataClass;

        let identity = AuthenticatedIdentity {
            subject: "analyst".to_string(),
            tenant_id: TenantId::new(1),
            roles: vec!["Analyst".to_string()],
            method: AuthMethod::Jwt,
        };

        let user_attrs = identity.extract_abac_user_attributes();
        let resource = ResourceAttributes::new(DataClass::Confidential, 1, "metrics");
        let env = EnvironmentAttributes::from_timestamp(Utc::now(), "US");

        // FedRAMP policy should allow US-based access
        let policy = AbacPolicy::fedramp_policy();
        let decision = evaluator::evaluate(&policy, &user_attrs, &resource, &env);

        assert_eq!(decision.effect, kimberlite_abac::PolicyEffect::Allow);
    }
}
