/// JWT validation and key management logic for auth enforcement.
use std::sync::{Arc, RwLock};
use std::time::Duration;

use jsonwebtoken::{decode, decode_header, errors::ErrorKind, Algorithm, DecodingKey, Validation};

use crate::plugins::auth_enforcement::auth_result::{map_auth_result_with_message, AuthResult};

use super::configuration::JWTRuleConfiguration;

use super::jwks::{JwksProviderRef, RealJwksProvider};

/// Abstraction for components that validate JWT access tokens.
pub trait JWTChecker {
    /// Validates a JWT token against the provided rule configuration.
    fn check_jwt_auth(&self, configuration: &JWTRuleConfiguration, token: &str) -> AuthResult;
}

/// Production JWT checker backed by cached JWKS decoding keys.
#[derive(Debug)]
pub struct RealJWTChecker {
    /// Decoding keys indexed by key ID (`kid`) for signature verification.
    keys: Arc<RwLock<Vec<(String, DecodingKey)>>>,
}

/// Represents the result of attempting to decode a token with one candidate key.
enum DecodeOutcome {
    Valid(serde_json::Value),
    Expired,
    Invalid,
}

impl RealJWTChecker {
    /// Creates a checker and starts a background JWKS refresh loop.
    ///
    /// This is the production constructor: it always wires the real
    /// network-backed `fetch_jwks` implementation.
    ///
    /// `jwks_uri` is fetched immediately and then re-fetched every
    /// `lifespan_secs` seconds while the checker is alive.
    pub async fn new(
        jwks_uri: &str,
        lifespan_secs: u64,
        connect_timeout: Option<Duration>,
        request_timeout: Option<Duration>,
    ) -> Result<Self, anyhow::Error> {
        let provider: JwksProviderRef = Arc::new(RealJwksProvider);
        Self::new_with_provider(
            jwks_uri,
            lifespan_secs,
            connect_timeout,
            request_timeout,
            provider,
        )
        .await
    }

    /// Creates a checker using an injected JWKS fetcher implementation.
    ///
    /// This is primarily intended for unit tests, where callers provide a
    /// deterministic in-memory/mock provider instead of making real HTTP calls.
    /// That is why unit tests use `new_with_provider` directly rather than `new`.
    async fn new_with_provider(
        jwks_uri: &str,
        lifespan_secs: u64,
        connect_timeout: Option<Duration>,
        request_timeout: Option<Duration>,
        provider: JwksProviderRef,
    ) -> Result<Self, anyhow::Error> {
        if lifespan_secs == 0 {
            anyhow::bail!("jwks_lifespan_secs must be greater than 0");
        }

        let client = reqwest::Client::builder()
            .connect_timeout(connect_timeout.unwrap_or(Duration::from_secs(5)))
            .timeout(request_timeout.unwrap_or(Duration::from_secs(10)))
            .build()?;
        let initial_keys = provider
            .fetch_jwks(client.clone(), jwks_uri.to_string())
            .await?;
        let keys = Arc::new(RwLock::new(initial_keys));

        // Weak reference so the task exits automatically when the checker is dropped —
        // a strong clone here would keep `keys` alive indefinitely.
        let keys_weak = Arc::downgrade(&keys);
        let uri = jwks_uri.to_string();
        let bg_client = client.clone();
        let bg_provider = provider.clone();
        tokio::spawn(async move {
            loop {
                tokio::time::sleep(Duration::from_secs(lifespan_secs)).await;
                // Upgrade confirms the checker is still alive; None means it was dropped and the task should stop.
                let Some(keys_arc) = keys_weak.upgrade() else {
                    break;
                };
                match bg_provider.fetch_jwks(bg_client.clone(), uri.clone()).await {
                    Ok(new_keys) => {
                        *keys_arc.write().unwrap() = new_keys;
                    }
                    Err(e) => {
                        tracing::error!("JWKS refresh failed: {}", e);
                    }
                }
            }
        });

        Ok(Self { keys })
    }

    /// Builds the shared JWT validation rules from the configured issuers and audiences.
    fn build_validation(configuration: &JWTRuleConfiguration) -> Validation {
        let mut validation = Validation::new(Algorithm::RS256);
        validation.set_issuer(&configuration.issuers);
        validation.set_audience(&configuration.audiences);
        validation.set_required_spec_claims(&["iss", "sub", "exp"]);
        validation.validate_exp = true;
        // Validation only happens if `aud` claim is present in the token, but it is not required to be present
        validation.validate_aud = true;
        // Validation only happens if `nbf` claim is present in the token, but it is not required to be present
        validation.validate_nbf = true;
        validation.leeway = 0;
        validation
    }

    /// Extracts the token `kid`, mapping header parse failures to the configured auth result.
    fn token_kid(configuration: &JWTRuleConfiguration, token: &str) -> Result<String, AuthResult> {
        let header = decode_header(token).map_err(|_| {
            map_auth_result_with_message(configuration.invalid_jwt, "Invalid JWT - No header")
        })?;

        header.kid.ok_or_else(|| {
            map_auth_result_with_message(configuration.invalid_jwt, "Invalid JWT - Missing kid")
        })
    }

    /// Returns only the cached decoding keys whose `kid` matches the token header.
    fn matching_keys<'a>(
        keys: &'a [(String, DecodingKey)],
        kid: &str,
        configuration: &JWTRuleConfiguration,
    ) -> Result<Vec<&'a DecodingKey>, AuthResult> {
        let keys_to_try: Vec<&DecodingKey> = keys
            .iter()
            .filter(|(known_kid, _)| known_kid == kid)
            .map(|(_, key)| key)
            .collect();

        if keys_to_try.is_empty() {
            return Err(map_auth_result_with_message(
                configuration.invalid_jwt,
                "Invalid JWT - Unknown kid",
            ));
        }

        Ok(keys_to_try)
    }

    /// Attempts JWT decode with one candidate key and maps the result into `DecodeOutcome`.
    ///
    /// On success, returns `DecodeOutcome::Valid` with the decoded claims payload
    /// (`serde_json::Value`) so downstream checks can validate claim semantics.
    /// Expired signatures map to `DecodeOutcome::Expired`, and all other decode
    /// errors map to `DecodeOutcome::Invalid`.
    fn decode_with_key(token: &str, key: &DecodingKey, validation: &Validation) -> DecodeOutcome {
        match decode::<serde_json::Value>(token, key, validation) {
            Ok(token_data) => DecodeOutcome::Valid(token_data.claims),
            Err(error) if matches!(error.kind(), ErrorKind::ExpiredSignature) => {
                DecodeOutcome::Expired
            }
            Err(_) => DecodeOutcome::Invalid,
        }
    }

    /// Applies claim-specific checks after signature validation succeeds for a matching key.
    fn validate_claims(
        configuration: &JWTRuleConfiguration,
        claims: &serde_json::Value,
        leeway: u64,
    ) -> AuthResult {
        if claims.get("azp").is_none() {
            return map_auth_result_with_message(
                configuration.invalid_jwt,
                "Invalid JWT - Missing required claim: azp",
            );
        }

        if claims.get("iat").is_none() {
            return map_auth_result_with_message(
                configuration.invalid_jwt,
                "Invalid JWT - Missing required claim: iat",
            );
        }

        let Some(iat) = claims.get("iat").and_then(|value| value.as_i64()) else {
            return map_auth_result_with_message(
                configuration.invalid_jwt,
                "Invalid JWT - Invalid iat claim",
            );
        };

        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_secs() as i64;
        if iat > now + (leeway as i64) {
            return map_auth_result_with_message(
                configuration.invalid_jwt,
                "Invalid JWT - Token issued in the future (iat)",
            );
        }

        if let Some(iss) = claims.get("iss").and_then(|value| value.as_str()) {
            tracing::Span::current().record("jwt.issuer", iss);
        }
        if let Some(sub) = claims.get("sub").and_then(|value| value.as_str()) {
            tracing::Span::current().record("jwt.subject", sub);
        }

        AuthResult::Ok
    }
}

impl JWTChecker for RealJWTChecker {
    /// Validates token claims and signature against configured issuers, audiences,
    /// and required claims.
    fn check_jwt_auth(&self, configuration: &JWTRuleConfiguration, token: &str) -> AuthResult {
        let validation = Self::build_validation(configuration);
        let keys_guard = self.keys.read().unwrap();
        let kid = match Self::token_kid(configuration, token) {
            Ok(kid) => kid,
            Err(result) => return result,
        };
        let keys_to_try = match Self::matching_keys(&keys_guard, &kid, configuration) {
            Ok(keys_to_try) => keys_to_try,
            Err(result) => return result,
        };

        for key in keys_to_try {
            match Self::decode_with_key(token, key, &validation) {
                DecodeOutcome::Valid(claims) => {
                    return Self::validate_claims(configuration, &claims, validation.leeway);
                }
                DecodeOutcome::Expired => {
                    return map_auth_result_with_message(configuration.expired_jwt, "Expired JWT");
                }
                DecodeOutcome::Invalid => {}
            }
        }

        map_auth_result_with_message(configuration.invalid_jwt, "Invalid JWT")
    }
}

#[cfg(test)]
/// Test doubles for JWT checker behavior; it makes it easier for another
/// module (e.g. check_auth.rs) to inject a simple mock checker without
/// needing to construct a full RealJWTChecker with keys and a fetcher.
pub mod mock {
    use super::*;

    /// Mock checker that always returns a preset authentication result.
    pub struct MockJWTChecker {
        result: AuthResult,
    }

    impl MockJWTChecker {
        /// Creates a mock checker with the provided result.
        pub fn new(result: AuthResult) -> MockJWTChecker {
            MockJWTChecker { result }
        }
    }

    impl JWTChecker for MockJWTChecker {
        /// Returns the preset result, ignoring configuration and token values.
        fn check_jwt_auth(
            &self,
            _configuration: &JWTRuleConfiguration,
            _token: &str,
        ) -> AuthResult {
            self.result
        }
    }
}

#[cfg(test)]
/// Unit tests covering JWT validation and result mapping behavior.
mod test {
    use jsonwebtoken::{encode, EncodingKey, Header};
    use serde_json::json;

    /// Inside mod test, super is check (this file/module), so jwks is in super::super.
    use super::super::jwks::{
        mock::{MockJwksProvider, MockJwksResponse},
        JwksError,
    };
    use crate::plugins::auth_enforcement::auth_result::AuthResultConfiguration;

    use super::*;

    const TEST_PRIVATE_KEY: &[u8] = b"-----BEGIN RSA PRIVATE KEY-----
MIIEowIBAAKCAQEArguts8kRMaf55JSXIfp0Nt1GYOEZeCaAwGiegSPZuKOh7Oqz
jhGCBayaNmoPEB86gRGKzpW9cdsSHAholAsQXiQYrMNZMgcDLT3TE/4n5KeL3zEs
FL+lsV9QndM2chpZ3/U/jzreOMz/U0JbBS6WsNU2HYwsFW+vU69fTVqChKf4SfbM
9rfV1I7zLZxEpYKyVIoSMXydTE90TuoPtXMtOLgRuBjp5En/szHqY3nkF+VbB9Fq
1mhu1LsnuynbxTWErwWf43rqGt9BFyvouwyzdfOD9H69hxw5i09De9dXLAU7nl5C
9AsOs/67DY4VFznMWRQ6AoFEhd+j5NzopVa3nQIDAQABAoIBABCwNfujL+7e+Gse
hE9RwIryNKldbB/FMVtBrNEdKuc1aXVaG9VngFwi7LaC/ObC181AbVrZqMyeAOKH
O9/WH7+nHmaSZJ8TZQ6VewPTvueB2TuRb3Sd27liNkyrz/Co7i03Stzk/CHeJi1J
k4ivqutRxlqW8jMr/le73RuBrjoh6CpnfWP4FqRPumDhBmbuaArAkW6hsaif8k2P
/Nh7mNlrB1lHRtMmNNjP3IPDgB6y150SpnG9MvZ9wKxytJjU+3aGHRQ8daZbEuJs
y3wfv5fQ/xbRDkin0dA9yTbeYbcDs/HDQ58tjYxLrBmk8ouEjqsfzHjZcITIskC2
UV3o2FECgYEA8RL1UcgS1RjgUKSluj7XMAmfcqJ69lNJV9i0NN+C0erQZ3ofrkA9
kIbxWvhGOcCPBcUMqxj3KLKmGZ+ttwgiNbeWokvleEoxVRtSBRb4py8RTI7VoxMw
1H6CginOA17Spp7dDZgkkIbub9kWDS1AkkpIdfpDx052XnbJ9SOkoxsCgYEAuNJQ
haIeTABaHYpKfbFFDZIaMrz+tQ4kdtbpVN5k+iSFDABu4Zd1ihmliZFVE1tKV98c
mcGH5VnyN5Ij8W1TqcB0OKYGIbC2BNrI3T+RNTIqH50PEUH4VJbWeRr5f0DWOtM2
rqjt9Yuwn7llck282Ex0cGkeoprtI83eEc6FA6cCgYEAikAlD9F9e1fYzGaf5D9F
iPb7yfIEMl9xID/WmLvVeiz/d1hB8txEci3xHAplu5kCbHyk1zpqA5zwRKGeLeks
NUIj0M9VubZStEwGrYSO9NG5Sf5f5jWLO1GR5rOywZwPkh5pBvLJhVjcRKhqTQQ2
k66l/4KWfCYzG9lj2IGoWSECgYBq1B348BklAvsTohCmkGWCj5WXCEkNbCxFiLl7
9cRWFe/VXRa6TKqC9YoKdze+pYDWQ78rXILpXdE0r8MY1IGMWBKVf9mBdbvwGGoM
AKc1IVbZyRWhewlIE5vKin32QQrTG6y6ziglSuX/nND3fdmPCLoxE4HCqVwy2tbm
YgHedwKBgAZhFjisfWO17TZ3BhbHD/F1+65w755s+7PE9mrhBVzgGaZROT+YhMoU
cM7Ou4PTrve2hRLSBYWD//aIIduRHNRhG+Sz4Nt8YpmIaJMJ+RpIy+sfWq3HgW9L
HTx1cTXnWNP8+KB7AZ2DUBk0tUEn3E8xEQCs/LnT16YTgUeFoXZS
-----END RSA PRIVATE KEY-----";

    const TEST_PUBLIC_KEY: &[u8] = b"-----BEGIN PUBLIC KEY-----
MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEArguts8kRMaf55JSXIfp0
Nt1GYOEZeCaAwGiegSPZuKOh7OqzjhGCBayaNmoPEB86gRGKzpW9cdsSHAholAsQ
XiQYrMNZMgcDLT3TE/4n5KeL3zEsFL+lsV9QndM2chpZ3/U/jzreOMz/U0JbBS6W
sNU2HYwsFW+vU69fTVqChKf4SfbM9rfV1I7zLZxEpYKyVIoSMXydTE90TuoPtXMt
OLgRuBjp5En/szHqY3nkF+VbB9Fq1mhu1LsnuynbxTWErwWf43rqGt9BFyvouwyz
dfOD9H69hxw5i09De9dXLAU7nl5C9AsOs/67DY4VFznMWRQ6AoFEhd+j5NzopVa3
nQIDAQAB
-----END PUBLIC KEY-----";

    const TEST_KID: &str = "test-key-1";

    const TEST_PRIVATE_KEY_2: &[u8] = b"-----BEGIN RSA PRIVATE KEY-----
MIIEogIBAAKCAQEAwJ4xrjVd0QyTCYIxAfzgIAekthY0Y1Ei03DSwQ1TRWh5+zqU
e3mBq0gX8nhNZznHBgT7J2Bp7URMID9RIDzgstIUd9kOMWBxbqv5kIq7Uar1yV0r
C7TTw7rPRbBJsM3k0zRogO29ziR6mSzHcvGnV4jar2TwYCkfahfSt5VaAx5In0Fj
hUeq8KUsohLsNHD/S3N3MMNS6gHSB/YrahYMY2cWPZBET3/wK1xiOYAo9RA8+t68
MJdLL1VkXE5FS/8OkCY2KSfned9fpVSMY+o64vJRqW8YLG9LYUMVS0U5jGBMZGYi
iAuY5UhspCSeCz/PB0b8Gl8OWIj2yoTw2yINDQIDAQABAoIBABeEWPl1j45B9EPR
bu1Ys/6FIGGQ5XWIxvE80aZ7W4LaZiY7uTw4EhJY+3+Bm6WyQq/j9MwSc+xhjc/U
GzIN7Ls4dgJJ8Hcr2L2T5qT0ozkt39Nleid0PPSlPsA57U1WXg1BVU8Ed8uO3Z8M
/g71HSLNVTHb55M+AhVgcjwy0moV1n1d1D9MTsoDF1Fq299Tf1Ljxhjmi1IYGwIg
6EHzSrMW3VZ7VgjRh3UBwPPboetnBYnKCG/sX0UMq7u5UjgPLQnLcAWVYtm1qi4M
diCbobqltlHx423h4SgCPNlYmB8jD6LfBvJx92Fwjm3Pec19Svcviln92el6shWt
cr6rcNECgYEA9CGKy88pQoJCyyx7X4avYr3xAZYACG8aGT+iS4v6dGez7q/klzS5
/lA6/ux9WueD/NbShzz61Gp+CG/yUJgB1ARHjzFkW7WYXkMkVQfZWxWzenCvVkap
29bvYiC52p2zDu5VeY3YWtT9DZANg6yrwhKUdNJWNKPdXyqHmcWG8RkCgYEAyfuE
+3bx4q+ft/iJq70ue3l+hkipZpUYOK/jDSCvi46tnzMrhhEkGCK566vzT5hzohCD
WGLojjT/hLwok0kBLUURWkLz5Blda7Vyvfz4sSVDSfaQEwqqm4UJC+IEDaDQP6/b
ZVfJD1bkJZ8bc1iLXZ0Im2urXfVXu+TrbQ99NhUCgYA81o8MM1gI5wVWNIOaTqou
ZT9GMedYMrfCSwVa5ktH7v/8H4FL3SvsoHnTRYQHv2I5BB/668tmXBOEbsFoLDlo
aJicqqZj5GFQPFM1jXsxI7tSa/qF+Z4i1vFY9ubaACq7MdHDNQIRwLo/QwLbKZC0
EIVGh5RMEfn82VBplbwHIQKBgGbpsKgvwq8trMOf7wWCLtza2oyxGR8beGvpslZO
ASOqXuxutg5pTWq+xtw9kItNsh0xLLssKjHK7yZCeKoS2Haoal8DHiFqix8GJKCI
9G3N/qsA5gfNt3/mGOXlEwuelHzerXrFb7xRGY3Rs7rhsqDQEkKeaYJGL52c1heL
fNDBAoGAD8eTruhWutxZJtRb7hC2P8tQrG2JdnRn4UUDH2G8ZNaHJhblK2RElR35
Fs7dVVpoPRf0Mh67pGonxXmFBsZeuhTfS4BiCk1NVSQWIkxAB1mUQThuC5ofUwOk
cI6yfrP+ThZh/7hHVzUmG4d4M2R8npeBpiuYYFoHnKAkk2Jl8ZM=
-----END RSA PRIVATE KEY-----";

    const TEST_PUBLIC_KEY_2: &[u8] = b"-----BEGIN PUBLIC KEY-----
MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAwJ4xrjVd0QyTCYIxAfzg
IAekthY0Y1Ei03DSwQ1TRWh5+zqUe3mBq0gX8nhNZznHBgT7J2Bp7URMID9RIDzg
stIUd9kOMWBxbqv5kIq7Uar1yV0rC7TTw7rPRbBJsM3k0zRogO29ziR6mSzHcvGn
V4jar2TwYCkfahfSt5VaAx5In0FjhUeq8KUsohLsNHD/S3N3MMNS6gHSB/YrahYM
Y2cWPZBET3/wK1xiOYAo9RA8+t68MJdLL1VkXE5FS/8OkCY2KSfned9fpVSMY+o6
4vJRqW8YLG9LYUMVS0U5jGBMZGYiiAuY5UhspCSeCz/PB0b8Gl8OWIj2yoTw2yIN
DQIDAQAB
-----END PUBLIC KEY-----";

    const TEST_KID_2: &str = "test-key-2";

    mod test_real_jwt_checker_construction {
        use super::*;

        #[tokio::test]
        async fn test_zero_lifespan_secs_is_rejected() {
            let provider: JwksProviderRef =
                Arc::new(MockJwksProvider::new(vec![
                    Arc::new(|| Ok(vec![])) as MockJwksResponse
                ]));
            let result =
                RealJWTChecker::new_with_provider("http://ignored", 0, None, None, provider).await;
            assert!(result.is_err());
            assert!(result
                .err()
                .unwrap()
                .to_string()
                .contains("jwks_lifespan_secs must be greater than 0"));
        }

        #[tokio::test]
        async fn test_refresh_success_installs_new_keys() {
            let provider: JwksProviderRef = Arc::new(MockJwksProvider::new(vec![
                Arc::new(|| {
                    // Initial fetch: return an unrelated key that cannot verify our token.
                    let unrelated_key = DecodingKey::from_rsa_pem(TEST_PUBLIC_KEY_2).unwrap();
                    Ok(vec![("other-key".to_string(), unrelated_key)])
                }) as MockJwksResponse,
                Arc::new(|| {
                    // Refresh and later calls: install/keep the real key.
                    let real_key = DecodingKey::from_rsa_pem(TEST_PUBLIC_KEY).unwrap();
                    Ok(vec![(TEST_KID.to_string(), real_key)])
                }) as MockJwksResponse,
            ]));

            let checker =
                RealJWTChecker::new_with_provider("http://ignored", 1, None, None, provider)
                    .await
                    .unwrap();

            let config = JWTRuleConfiguration::all_warning();
            let token = sign_token(
                &valid_claims(&config.issuers[0], &config.audiences[0]),
                Some(TEST_KID),
            );

            // Before refresh: the token should fail because the real key isn't loaded yet.
            assert_eq!(
                checker.check_jwt_auth(&config, &token),
                AuthResult::Warning {
                    message: "Invalid JWT - Unknown kid"
                }
            );

            // Wait for the background refresh to fire and install the real key.
            tokio::time::sleep(Duration::from_millis(1500)).await;

            // After refresh: the same token should now validate.
            assert_eq!(checker.check_jwt_auth(&config, &token), AuthResult::Ok);
        }

        #[tokio::test]
        /// Verifies that a failed JWKS refresh leaves the previously fetched keys intact.
        async fn test_refresh_failure_retains_prior_keys() {
            let provider: JwksProviderRef = Arc::new(MockJwksProvider::new(vec![
                Arc::new(|| {
                    let key = DecodingKey::from_rsa_pem(TEST_PUBLIC_KEY).unwrap();
                    Ok(vec![(TEST_KID.to_string(), key)])
                }) as MockJwksResponse,
                Arc::new(|| {
                    Err(JwksError::ErrorStatus(
                        reqwest::StatusCode::INTERNAL_SERVER_ERROR,
                    ))
                }) as MockJwksResponse,
            ]));

            // lifespan_secs=1 so the refresh loop fires quickly.
            let checker =
                RealJWTChecker::new_with_provider("http://ignored", 1, None, None, provider)
                    .await
                    .unwrap();

            // Wait long enough for the refresh loop to fire and fail.
            tokio::time::sleep(Duration::from_millis(1500)).await;

            // Keys from the initial fetch should still be in place.
            let config = JWTRuleConfiguration::all_warning();
            let token = sign_token(
                &valid_claims(&config.issuers[0], &config.audiences[0]),
                Some(TEST_KID),
            );
            assert_eq!(checker.check_jwt_auth(&config, &token), AuthResult::Ok);
        }
    }

    /// Builds a checker with a single in-memory test decoding key.
    fn test_checker() -> RealJWTChecker {
        let key = DecodingKey::from_rsa_pem(TEST_PUBLIC_KEY).unwrap();
        RealJWTChecker {
            keys: Arc::new(RwLock::new(vec![(TEST_KID.to_string(), key)])),
        }
    }

    /// Signs a JWT with test RSA credentials and an optional `kid` header.
    fn sign_token(claims: &serde_json::Value, kid: Option<&str>) -> String {
        let encoding_key = EncodingKey::from_rsa_pem(TEST_PRIVATE_KEY).unwrap();
        let mut header = Header::new(Algorithm::RS256);
        header.kid = kid.map(|s| s.to_string());
        encode(&header, claims, &encoding_key).unwrap()
    }

    /// Returns the current Unix epoch timestamp in seconds for time-based claims.
    fn now_epoch_secs() -> i64 {
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_secs() as i64
    }

    /// Builds the baseline valid JWT claims used by most tests.
    ///
    /// Individual tests should mutate this shape via `token_with_claim_overrides`
    /// when validating claim-specific behavior.
    fn default_claims(config: &JWTRuleConfiguration) -> serde_json::Map<String, serde_json::Value> {
        let now = now_epoch_secs();
        let mut claims = serde_json::Map::new();
        claims.insert("iss".to_string(), json!(config.issuers[0]));
        claims.insert("sub".to_string(), json!("user|123"));
        claims.insert("aud".to_string(), json!(config.audiences[0]));
        claims.insert("exp".to_string(), json!(now + 3600));
        claims.insert("iat".to_string(), json!(now));
        claims.insert("azp".to_string(), json!("some-client"));
        claims
    }

    /// Creates a signed token from default claims after applying a targeted mutation.
    ///
    /// This keeps tests focused on the one claim being changed instead of rebuilding
    /// an entire payload each time.
    fn token_with_claim_overrides(
        config: &JWTRuleConfiguration,
        mutate_claims: impl FnOnce(&mut serde_json::Map<String, serde_json::Value>),
    ) -> String {
        let mut claims = default_claims(config);
        mutate_claims(&mut claims);
        sign_token(&serde_json::Value::Object(claims), Some(TEST_KID))
    }

    /// Asserts checker output using an `all_ok` config after applying caller overrides.
    fn assert_with_all_ok_config(
        configure: impl FnOnce(&mut JWTRuleConfiguration),
        token_builder: impl FnOnce(&JWTRuleConfiguration) -> String,
        expected: AuthResult,
    ) {
        let mut config = JWTRuleConfiguration::all_ok();
        configure(&mut config);
        let checker = test_checker();
        let token = token_builder(&config);
        assert_eq!(checker.check_jwt_auth(&config, &token), expected);
    }

    /// Converts an enforcement config enum into the expected auth result for a message.
    fn auth_result_from_config(
        result: &AuthResultConfiguration,
        message: &'static str,
    ) -> AuthResult {
        match result {
            AuthResultConfiguration::Ok => AuthResult::Ok,
            AuthResultConfiguration::Warning => AuthResult::Warning { message },
            AuthResultConfiguration::Block => AuthResult::Block { message },
        }
    }

    /// Verifies that `Ok`, `Warning`, and `Block` enforcement settings are all respected.
    ///
    /// The provided rule setter selects which config field to drive (e.g. `invalid_jwt`
    /// or `expired_jwt`), and `token_builder` defines the scenario under test.
    fn assert_respects_enforcement_configs(
        set_rule: impl Fn(&mut JWTRuleConfiguration, AuthResultConfiguration),
        token_builder: impl Fn(&JWTRuleConfiguration) -> String,
        message: &'static str,
    ) {
        for result in [
            AuthResultConfiguration::Ok,
            AuthResultConfiguration::Warning,
            AuthResultConfiguration::Block,
        ] {
            assert_with_all_ok_config(
                |config| set_rule(config, result),
                |config| token_builder(config),
                auth_result_from_config(&result, message),
            );
        }
    }

    /// Generates a token payload with an expiration in the future.
    fn valid_claims(iss: &str, aud: &str) -> serde_json::Value {
        let now = now_epoch_secs();
        json!({
            "iss": iss,
            "sub": "user|123",
            "aud": aud,
            "exp": now + 3600,
            "iat": now,
            "azp": "some-client"
        })
    }

    #[test]
    /// Verifies a valid token passes validation.
    fn test_valid_jwt_ok() {
        let config = JWTRuleConfiguration::all_warning();
        let checker = test_checker();
        let token = sign_token(
            &valid_claims(&config.issuers[0], &config.audiences[0]),
            Some(TEST_KID),
        );
        assert_eq!(checker.check_jwt_auth(&config, &token), AuthResult::Ok);
    }

    #[test]
    /// Verifies that non-JWT input fails at header decoding and maps through
    /// `invalid_jwt` with the specific `No header` message.
    fn test_invalid_jwt_no_header() {
        let config = JWTRuleConfiguration::all_warning();
        let checker = test_checker();
        assert_eq!(
            checker.check_jwt_auth(&config, "not.a.jwt"),
            AuthResult::Warning {
                message: "Invalid JWT - No header"
            }
        );
    }

    fn invalid_token_with_decodable_header(config: &JWTRuleConfiguration) -> String {
        let encoding_key = EncodingKey::from_rsa_pem(TEST_PRIVATE_KEY_2).unwrap();
        let mut header = Header::new(Algorithm::RS256);
        // Use a known kid so the checker attempts verification and fails on signature.
        header.kid = Some(TEST_KID.to_string());
        let claims = valid_claims(&config.issuers[0], &config.audiences[0]);
        encode(&header, &claims, &encoding_key).unwrap()
    }

    #[test]
    fn test_invalid_jwt_respects_enforcement_configs() {
        assert_respects_enforcement_configs(
            |config, result| config.invalid_jwt = result,
            invalid_token_with_decodable_header,
            "Invalid JWT",
        );
    }

    fn expired_token(config: &JWTRuleConfiguration) -> String {
        token_with_claim_overrides(config, |claims| {
            let now = now_epoch_secs();
            claims.insert("exp".to_string(), json!(now - 3600));
            claims.insert("iat".to_string(), json!(now - 7200));
        })
    }

    #[test]
    fn test_expired_jwt_respects_enforcement_configs() {
        assert_respects_enforcement_configs(
            |config, result| config.expired_jwt = result,
            expired_token,
            "Expired JWT",
        );
    }

    mod test_kid_matching {
        use super::*;

        #[test]
        /// Verifies tokens without `kid` are rejected immediately.
        fn test_no_kid_in_header_respects_enforcement_configs() {
            assert_respects_enforcement_configs(
                |config, result| config.invalid_jwt = result,
                |config| {
                    sign_token(
                        &valid_claims(&config.issuers[0], &config.audiences[0]),
                        None,
                    )
                },
                "Invalid JWT - Missing kid",
            );
        }

        #[test]
        /// Verifies unknown `kid` values fail validation as invalid JWTs.
        fn test_unknown_kid_respects_enforcement_configs() {
            assert_respects_enforcement_configs(
                |config, result| config.invalid_jwt = result,
                |config| {
                    sign_token(
                        &valid_claims(&config.issuers[0], &config.audiences[0]),
                        Some("unknown-key"),
                    )
                },
                "Invalid JWT - Unknown kid",
            );
        }

        #[test]
        /// Verifies that a token signed with an untrusted key and carrying an unknown
        /// `kid` is rejected before signature verification, returning `Unknown kid`.
        ///
        /// This specifically exercises key selection behavior: the checker only tries
        /// decoding keys whose `kid` matches a cached JWKS entry.
        fn test_wrong_signature_key_respects_enforcement_configs() {
            assert_respects_enforcement_configs(
                |config, result| config.invalid_jwt = result,
                |config| {
                    // Sign with key-2, but checker only knows key-1.
                    let encoding_key = EncodingKey::from_rsa_pem(TEST_PRIVATE_KEY_2).unwrap();
                    let mut header = Header::new(Algorithm::RS256);
                    header.kid = Some(TEST_KID_2.to_string());
                    let claims = valid_claims(&config.issuers[0], &config.audiences[0]);
                    encode(&header, &claims, &encoding_key).unwrap()
                },
                "Invalid JWT - Unknown kid",
            );
        }

        #[test]
        fn test_second_matching_kid_key_validates_when_first_fails() {
            let config = JWTRuleConfiguration::all_warning();
            // Defensive malformed-JWKS scenario: duplicate `kid` entries should be rare,
            // but we still verify the checker can continue past a bad match and succeed.
            // Store a decoy key first with the same kid, then the real key second.
            let decoy_key = DecodingKey::from_rsa_pem(TEST_PUBLIC_KEY_2).unwrap();
            let real_key = DecodingKey::from_rsa_pem(TEST_PUBLIC_KEY).unwrap();
            let checker = RealJWTChecker {
                keys: Arc::new(RwLock::new(vec![
                    (TEST_KID.to_string(), decoy_key),
                    (TEST_KID.to_string(), real_key),
                ])),
            };
            // Token's kid matches both entries; checker should continue to the second key.
            let token = sign_token(
                &valid_claims(&config.issuers[0], &config.audiences[0]),
                Some(TEST_KID),
            );
            assert_eq!(checker.check_jwt_auth(&config, &token), AuthResult::Ok);
        }

        #[test]
        fn test_kid_match_skips_other_keys() {
            let config = JWTRuleConfiguration::all_warning();
            let real_key = DecodingKey::from_rsa_pem(TEST_PUBLIC_KEY).unwrap();
            let other_key = DecodingKey::from_rsa_pem(TEST_PUBLIC_KEY_2).unwrap();
            let checker = RealJWTChecker {
                keys: Arc::new(RwLock::new(vec![
                    (TEST_KID.to_string(), real_key),
                    ("other-key".to_string(), other_key),
                ])),
            };
            // Token's kid matches TEST_KID exactly; only that key should be tried.
            let token = sign_token(
                &valid_claims(&config.issuers[0], &config.audiences[0]),
                Some(TEST_KID),
            );
            assert_eq!(checker.check_jwt_auth(&config, &token), AuthResult::Ok);
        }
    }

    mod test_iss_claim {
        use super::*;

        #[test]
        fn test_missing_iss_claim_respects_enforcement_configs() {
            assert_respects_enforcement_configs(
                |config, result| config.invalid_jwt = result,
                |config| {
                    token_with_claim_overrides(config, |claims| {
                        claims.remove("iss");
                    })
                },
                "Invalid JWT",
            );
        }

        #[test]
        fn test_wrong_issuer_respects_enforcement_configs() {
            assert_respects_enforcement_configs(
                |config, result| config.invalid_jwt = result,
                |config| {
                    sign_token(
                        &valid_claims(
                            "https://untrusted-issuer.example.com/",
                            &config.audiences[0],
                        ),
                        Some(TEST_KID),
                    )
                },
                "Invalid JWT",
            );
        }

        #[test]
        fn test_second_issuer_in_config_is_accepted() {
            let mut config = JWTRuleConfiguration::all_warning();
            config.issuers = vec![
                "https://primary.example.com/".to_string(),
                "https://secondary.example.com/".to_string(),
            ];
            let checker = test_checker();
            let token = sign_token(
                &valid_claims("https://secondary.example.com/", &config.audiences[0]),
                Some(TEST_KID),
            );
            assert_eq!(checker.check_jwt_auth(&config, &token), AuthResult::Ok);
        }

        #[test]
        fn test_issuer_not_in_multi_issuer_config_respects_enforcement_configs() {
            assert_respects_enforcement_configs(
                |config, result| {
                    config.issuers = vec![
                        "https://primary.example.com/".to_string(),
                        "https://secondary.example.com/".to_string(),
                    ];
                    config.invalid_jwt = result;
                },
                |config| {
                    sign_token(
                        &valid_claims("https://unlisted.example.com/", &config.audiences[0]),
                        Some(TEST_KID),
                    )
                },
                "Invalid JWT",
            );
        }
    }

    mod test_aud_claim {
        use super::*;

        #[test]
        fn test_missing_aud_claim_is_valid() {
            let config = JWTRuleConfiguration::all_warning();
            let checker = test_checker();
            let token = token_with_claim_overrides(&config, |claims| {
                claims.remove("aud");
            });
            assert_eq!(checker.check_jwt_auth(&config, &token), AuthResult::Ok);
        }

        #[test]
        fn test_wrong_audience_respects_enforcement_configs() {
            assert_respects_enforcement_configs(
                |config, result| config.invalid_jwt = result,
                |config| {
                    sign_token(
                        &valid_claims(&config.issuers[0], "https://wrong-audience.example.com"),
                        Some(TEST_KID),
                    )
                },
                "Invalid JWT",
            );
        }

        #[test]
        fn test_second_audience_in_config_is_accepted() {
            let mut config = JWTRuleConfiguration::all_warning();
            config.audiences = vec![
                "https://api-primary.example.com".to_string(),
                "https://api-secondary.example.com".to_string(),
            ];
            let checker = test_checker();
            let token = sign_token(
                &valid_claims(&config.issuers[0], "https://api-secondary.example.com"),
                Some(TEST_KID),
            );
            assert_eq!(checker.check_jwt_auth(&config, &token), AuthResult::Ok);
        }

        #[test]
        fn test_audience_not_in_multi_audience_config_respects_enforcement_configs() {
            assert_respects_enforcement_configs(
                |config, result| {
                    config.audiences = vec![
                        "https://api-primary.example.com".to_string(),
                        "https://api-secondary.example.com".to_string(),
                    ];
                    config.invalid_jwt = result;
                },
                |config| {
                    sign_token(
                        &valid_claims(&config.issuers[0], "https://api-unlisted.example.com"),
                        Some(TEST_KID),
                    )
                },
                "Invalid JWT",
            );
        }
    }

    #[test]
    fn test_future_nbf_claim_respects_enforcement_configs() {
        assert_respects_enforcement_configs(
            |config, result| config.invalid_jwt = result,
            |config| {
                token_with_claim_overrides(config, |claims| {
                    claims.insert("nbf".to_string(), json!(now_epoch_secs() + 600));
                })
            },
            "Invalid JWT",
        );
    }

    #[test]
    fn test_invalid_nbf_claim_type_respects_enforcement_configs() {
        assert_respects_enforcement_configs(
            |config, result| config.invalid_jwt = result,
            |config| {
                token_with_claim_overrides(config, |claims| {
                    claims.insert("nbf".to_string(), json!("soon"));
                })
            },
            "Invalid JWT",
        );
    }

    #[test]
    fn test_missing_iat_claim_respects_enforcement_configs() {
        assert_respects_enforcement_configs(
            |config, result| config.invalid_jwt = result,
            |config| {
                token_with_claim_overrides(config, |claims| {
                    claims.remove("iat");
                })
            },
            "Invalid JWT - Missing required claim: iat",
        );
    }

    #[test]
    fn test_future_iat_claim_respects_enforcement_configs() {
        assert_respects_enforcement_configs(
            |config, result| config.invalid_jwt = result,
            |config| {
                token_with_claim_overrides(config, |claims| {
                    claims.insert("iat".to_string(), json!(now_epoch_secs() + 600));
                })
            },
            "Invalid JWT - Token issued in the future (iat)",
        );
    }

    #[test]
    fn test_invalid_iat_claim_type_respects_enforcement_configs() {
        assert_respects_enforcement_configs(
            |config, result| config.invalid_jwt = result,
            |config| {
                token_with_claim_overrides(config, |claims| {
                    claims.insert("iat".to_string(), json!("later"));
                })
            },
            "Invalid JWT - Invalid iat claim",
        );
    }

    #[test]
    fn test_missing_sub_claim_respects_enforcement_configs() {
        assert_respects_enforcement_configs(
            |config, result| config.invalid_jwt = result,
            |config| {
                token_with_claim_overrides(config, |claims| {
                    claims.remove("sub");
                })
            },
            "Invalid JWT",
        );
    }

    #[test]
    fn test_missing_azp_claim_respects_enforcement_configs() {
        assert_respects_enforcement_configs(
            |config, result| config.invalid_jwt = result,
            |config| {
                token_with_claim_overrides(config, |claims| {
                    claims.remove("azp");
                })
            },
            "Invalid JWT - Missing required claim: azp",
        );
    }
}

#[cfg(test)]
mod test_mock {
    use crate::plugins::auth_enforcement::auth_result::AuthResult;
    use crate::plugins::auth_enforcement::token::jwt::configuration::JWTRuleConfiguration;

    use super::mock::MockJWTChecker;
    use super::JWTChecker;

    #[test]
    fn test_mock_jwt_checker() {
        let checker = MockJWTChecker::new(AuthResult::Warning { message: "mock" });
        let config = JWTRuleConfiguration::all_ok();
        assert_eq!(
            checker.check_jwt_auth(&config, "any-token"),
            AuthResult::Warning { message: "mock" }
        );
    }
}