//! Rate limiting: Cloudflare-aware IP extraction, per-app SyncKit extraction,
//! and governor config builders.

use base64::Engine;
use tower_governor::errors::GovernorError;
use tower_governor::key_extractor::KeyExtractor;

use crate::db::SyncAppId;

/// IP key extractor that trusts `CF-Connecting-IP` as the client identity.
///
/// `CF-Connecting-IP` is safe to trust because every public path to the app
/// sets it from a value the client cannot forge:
///   - makenot.work subdomains: Caddy enforces Cloudflare mTLS, so only
///     Cloudflare reaches the origin and Cloudflare sets the header itself.
///   - Custom domains (the `:443` on-demand-TLS block): Caddy overwrites
///     `CF-Connecting-IP` with the real TCP peer and strips `X-Forwarded-For`
///     before proxying, so a client-supplied value never survives.
///
/// When the header is absent — only direct dev/test access with no Caddy in
/// front — we fall back to [`PeerIpKeyExtractor`], the actual TCP peer from
/// `ConnectInfo`. We deliberately do NOT fall back to `SmartIpKeyExtractor`:
/// it trusts `X-Forwarded-For`/`X-Real-IP`, which a client can forge to mint
/// fresh rate-limit buckets and evade login/signup/reset throttles. In prod
/// the peer is always Caddy (loopback), so a missing CF header collapses every
/// requester into one bucket — the safe degraded behavior `extract_client_ip`
/// already warns about — rather than handing out spoofable per-IP buckets.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct CloudflareIpKeyExtractor;

impl KeyExtractor for CloudflareIpKeyExtractor {
    type Key = std::net::IpAddr;

    fn extract<T>(&self, req: &axum::http::Request<T>) -> Result<Self::Key, GovernorError> {
        if let Some(ip) = req
            .headers()
            .get("cf-connecting-ip")
            .and_then(|v: &axum::http::HeaderValue| v.to_str().ok())
            .and_then(|s: &str| s.trim().parse::<std::net::IpAddr>().ok())
        {
            return Ok(ip);
        }

        tower_governor::key_extractor::PeerIpKeyExtractor.extract(req)
    }
}

/// Per-SyncKit-app key extractor. Reads the `app` claim from the
/// `Authorization: Bearer <token>` header to bucket requests per app.
///
/// The signature IS verified here (HS256, against the SyncKit JWT secret).
/// Verification matters for rate limiting specifically: without it, an attacker
/// can mint unsigned tokens carrying arbitrary `app` UUIDs and spray a fresh
/// claim per request, landing each in its own fresh bucket and defeating the
/// per-app limiter entirely. By verifying, only validly-signed tokens earn
/// their real per-app bucket; everything else (missing/forged/malformed token)
/// collapses into ONE shared sentinel bucket, so a spray cannot manufacture
/// unlimited buckets. Auth itself is still enforced downstream by `SyncUser`,
/// and an IP-layer limiter backstops volume regardless.
///
/// Expiry/issuer are intentionally NOT enforced here — an expired-but-genuine
/// token should still bucket by its real app (the handler will reject it). We
/// only need proof the `app` claim wasn't fabricated.
#[derive(Debug, Clone)]
pub struct SyncAppKeyExtractor {
    /// SyncKit JWT signing secret. `None` when SyncKit isn't configured — in
    /// that case there is no auth and these routes are inert, so we fall back
    /// to the unverified payload parse to preserve prior behavior.
    secret: Option<std::sync::Arc<String>>,
}

impl SyncAppKeyExtractor {
    pub fn new(secret: Option<std::sync::Arc<String>>) -> Self {
        Self { secret }
    }

    /// Verify the HS256 signature and pull out the `app` claim. Returns `None`
    /// on any failure (bad signature, malformed token, missing claim).
    fn verify_app(secret: &str, token: &str) -> Option<SyncAppId> {
        use jsonwebtoken::{decode, Algorithm, DecodingKey, Validation};

        #[derive(serde::Deserialize)]
        struct AppClaim {
            app: SyncAppId,
        }

        let mut validation = Validation::new(Algorithm::HS256);
        // Bucket genuine-but-expired tokens by their real app; only the
        // signature must hold. Clearing required claims + disabling exp means
        // the decode succeeds purely on a valid signature carrying `app`.
        validation.validate_exp = false;
        validation.required_spec_claims.clear();

        decode::<AppClaim>(token, &DecodingKey::from_secret(secret.as_bytes()), &validation)
            .ok()
            .map(|data| data.claims.app)
    }

    /// Unverified payload parse — only used when no secret is configured.
    fn parse_app_unverified(token: &str) -> Option<SyncAppId> {
        let payload_b64 = token.split('.').nth(1)?;
        let payload_bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD
            .decode(payload_b64)
            .or_else(|_| base64::engine::general_purpose::STANDARD.decode(payload_b64))
            .ok()?;

        #[derive(serde::Deserialize)]
        struct AppClaim {
            app: SyncAppId,
        }
        serde_json::from_slice::<AppClaim>(&payload_bytes).ok().map(|c| c.app)
    }
}

impl KeyExtractor for SyncAppKeyExtractor {
    type Key = SyncAppId;

    fn extract<T>(&self, req: &axum::http::Request<T>) -> Result<Self::Key, GovernorError> {
        let token = match req
            .headers()
            .get("authorization")
            .and_then(|v| v.to_str().ok())
            .and_then(|s| s.strip_prefix("Bearer "))
        {
            Some(t) => t,
            // No bearer token — use a nil sentinel key so the request passes
            // through to the handler, where SyncUser will properly return 401.
            None => return Ok(SyncAppId::nil()),
        };

        let app = match &self.secret {
            Some(secret) => Self::verify_app(secret, token),
            None => Self::parse_app_unverified(token),
        };

        // A verified app gets its own bucket; anything that fails verification
        // collapses to the nil sentinel so forged tokens can't spray buckets.
        Ok(app.unwrap_or_else(SyncAppId::nil))
    }
}

// ── Config builders ──

// ── Bucket-map sweeping (Run #14 CHRONIC 1) ──

/// Type-erased `retain_recent` GC hooks, one per limiter built below. Each hook
/// sweeps one limiter's keyed GCRA store and returns its post-sweep entry count.
///
/// tower_governor's in-memory store grows one entry per unique client key for
/// process lifetime unless swept, so EVERY limiter must be registered here. The
/// three `rate_limiter_*` constructors below are the only sanctioned way to
/// build a limiter precisely because they register on construct — do not build a
/// `GovernorConfig` directly (it would leak, unswept). The registry is touched
/// only at startup (registration) and once per sweep interval, so the `Mutex` is
/// effectively uncontended; no lock is ever held across an `.await`.
static GOVERNOR_SWEEPERS: std::sync::Mutex<Vec<Box<dyn Fn() -> usize + Send + Sync>>> =
    std::sync::Mutex::new(Vec::new());

/// Register a limiter's GC hook. Monomorphized at each call site (where the
/// limiter's concrete key type is known), so this stays non-generic.
fn register_for_sweep(hook: impl Fn() -> usize + Send + Sync + 'static) {
    if let Ok(mut hooks) = GOVERNOR_SWEEPERS.lock() {
        hooks.push(Box::new(hook));
    }
}

/// Spawn the periodic task that sweeps every registered limiter's bucket map.
/// Call once at startup (guarded by `Once`, so extra calls — e.g. per-test
/// `build_app` — are no-ops). Requires a Tokio runtime.
pub fn start_governor_sweeper() {
    static STARTED: std::sync::Once = std::sync::Once::new();
    STARTED.call_once(|| {
        tokio::spawn(async {
            let interval =
                std::time::Duration::from_secs(crate::constants::GOVERNOR_SWEEP_INTERVAL_SECS);
            loop {
                tokio::time::sleep(interval).await;
                // Collect counts without holding the lock across any await.
                let (limiters, retained) = {
                    let Ok(hooks) = GOVERNOR_SWEEPERS.lock() else {
                        continue;
                    };
                    let retained: usize = hooks.iter().map(|hook| hook()).sum();
                    (hooks.len(), retained)
                };
                tracing::debug!(limiters, retained_keys = retained, "swept governor bucket maps");
            }
        });
    });
}

/// Build an IP-based rate limiter from a per-millisecond interval and burst size.
pub fn rate_limiter_ms(
    ms: u64,
    burst: u32,
) -> std::sync::Arc<
    tower_governor::governor::GovernorConfig<
        CloudflareIpKeyExtractor,
        ::governor::middleware::StateInformationMiddleware,
    >,
> {
    let config = std::sync::Arc::new(
        tower_governor::governor::GovernorConfigBuilder::default()
            .key_extractor(CloudflareIpKeyExtractor)
            .per_millisecond(ms)
            .burst_size(burst)
            .use_headers()
            .finish()
            .expect("rate limiter config"),
    );
    let limiter = config.limiter().clone();
    register_for_sweep(move || {
        limiter.retain_recent();
        limiter.len()
    });
    config
}

/// Build an IP-based rate limiter from a per-second rate and burst size.
pub fn rate_limiter_per_sec(
    per_sec: u64,
    burst: u32,
) -> std::sync::Arc<
    tower_governor::governor::GovernorConfig<
        CloudflareIpKeyExtractor,
        ::governor::middleware::StateInformationMiddleware,
    >,
> {
    let config = std::sync::Arc::new(
        tower_governor::governor::GovernorConfigBuilder::default()
            .key_extractor(CloudflareIpKeyExtractor)
            .per_second(per_sec)
            .burst_size(burst)
            .use_headers()
            .finish()
            .expect("rate limiter config"),
    );
    let limiter = config.limiter().clone();
    register_for_sweep(move || {
        limiter.retain_recent();
        limiter.len()
    });
    config
}

/// Build a per-SyncKit-app rate limiter from a per-millisecond interval and burst size.
///
/// `secret` is the SyncKit JWT signing secret; when present the extractor
/// verifies token signatures (so forged `app` claims can't mint fresh buckets).
pub fn synckit_app_rate_limiter_ms(
    secret: Option<std::sync::Arc<String>>,
    ms: u64,
    burst: u32,
) -> std::sync::Arc<
    tower_governor::governor::GovernorConfig<
        SyncAppKeyExtractor,
        ::governor::middleware::StateInformationMiddleware,
    >,
> {
    let config = std::sync::Arc::new(
        tower_governor::governor::GovernorConfigBuilder::default()
            .key_extractor(SyncAppKeyExtractor::new(secret))
            .per_millisecond(ms)
            .burst_size(burst)
            .use_headers()
            .finish()
            .expect("synckit app rate limiter config"),
    );
    let limiter = config.limiter().clone();
    register_for_sweep(move || {
        limiter.retain_recent();
        limiter.len()
    });
    config
}

#[cfg(test)]
mod tests {
    use super::*;
    use axum::http::Request;
    use tower_governor::key_extractor::KeyExtractor;

    /// Build a fake JWT with the given app ID in the payload (no signature verification).
    fn fake_jwt(app_id: &SyncAppId) -> String {
        let header = base64::engine::general_purpose::URL_SAFE_NO_PAD
            .encode(r#"{"alg":"HS256","typ":"JWT"}"#);
        let payload_json = serde_json::json!({
            "sub": "00000000-0000-0000-0000-000000000001",
            "app": app_id,
            "iss": "makenotwork-synckit",
            "exp": 9999999999_i64,
            "iat": 1000000000_i64,
        });
        let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD
            .encode(payload_json.to_string());
        let sig = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode("fakesig");
        format!("{header}.{payload}.{sig}")
    }

    /// Sign a real HS256 token carrying the given app id, using `secret`.
    fn signed_jwt(secret: &str, app_id: &SyncAppId) -> String {
        use jsonwebtoken::{encode, Algorithm, EncodingKey, Header};
        let claims = serde_json::json!({
            "sub": "00000000-0000-0000-0000-000000000001",
            "app": app_id,
            "iss": "makenotwork-synckit",
            "exp": 9999999999_i64,
            "iat": 1000000000_i64,
        });
        encode(
            &Header::new(Algorithm::HS256),
            &claims,
            &EncodingKey::from_secret(secret.as_bytes()),
        )
        .unwrap()
    }

    #[test]
    fn unverified_extracts_app_id_from_jwt_when_no_secret() {
        // No secret configured (SyncKit off) → fall back to unverified parse.
        let app_id = SyncAppId::new();
        let token = fake_jwt(&app_id);

        let req = Request::builder()
            .header("authorization", format!("Bearer {token}"))
            .body(())
            .unwrap();

        let extracted = SyncAppKeyExtractor::new(None).extract(&req).unwrap();
        assert_eq!(extracted, app_id);
    }

    #[test]
    fn verified_extracts_app_id_from_validly_signed_jwt() {
        let secret = "test-secret-key-for-synckit-jwt".to_string();
        let app_id = SyncAppId::new();
        let token = signed_jwt(&secret, &app_id);

        let req = Request::builder()
            .header("authorization", format!("Bearer {token}"))
            .body(())
            .unwrap();

        let extractor = SyncAppKeyExtractor::new(Some(std::sync::Arc::new(secret)));
        assert_eq!(extractor.extract(&req).unwrap(), app_id);
    }

    #[test]
    fn verified_forged_token_collapses_to_nil_bucket() {
        // A token whose `app` claim is attacker-chosen but whose signature does
        // NOT match the configured secret must NOT earn its own bucket — it
        // collapses to the nil sentinel so a spray of fresh `app` claims can't
        // manufacture unlimited buckets.
        let secret = "test-secret-key-for-synckit-jwt".to_string();
        let attacker_app = SyncAppId::new();
        let forged = fake_jwt(&attacker_app); // signed with "fakesig", not `secret`

        let req = Request::builder()
            .header("authorization", format!("Bearer {forged}"))
            .body(())
            .unwrap();

        let extractor = SyncAppKeyExtractor::new(Some(std::sync::Arc::new(secret)));
        assert_eq!(extractor.extract(&req).unwrap(), SyncAppId::nil());
    }

    #[test]
    fn verified_spray_of_forged_apps_all_share_one_bucket() {
        // The core anti-spray property: many DISTINCT forged app claims must all
        // map to the SAME (nil) key, not N distinct keys.
        let secret = std::sync::Arc::new("test-secret-key-for-synckit-jwt".to_string());
        let extractor = SyncAppKeyExtractor::new(Some(secret));
        for _ in 0..5 {
            let forged = fake_jwt(&SyncAppId::new());
            let req = Request::builder()
                .header("authorization", format!("Bearer {forged}"))
                .body(())
                .unwrap();
            assert_eq!(extractor.extract(&req).unwrap(), SyncAppId::nil());
        }
    }

    #[test]
    fn missing_auth_header_returns_nil_sentinel() {
        let req = Request::builder().body(()).unwrap();
        let key = SyncAppKeyExtractor::new(None).extract(&req).unwrap();
        assert_eq!(key, SyncAppId::nil());
    }

    #[test]
    fn non_bearer_auth_returns_nil_sentinel() {
        let req = Request::builder()
            .header("authorization", "Basic dXNlcjpwYXNz")
            .body(())
            .unwrap();
        let key = SyncAppKeyExtractor::new(None).extract(&req).unwrap();
        assert_eq!(key, SyncAppId::nil());
    }

    #[test]
    fn malformed_jwt_collapses_to_nil_sentinel() {
        // Previously this returned an extractor error; now a malformed token
        // collapses to the shared nil bucket (same anti-spray treatment as a
        // forged one) rather than erroring the request out of the limiter.
        let req = Request::builder()
            .header("authorization", "Bearer not-a-jwt")
            .body(())
            .unwrap();
        let key = SyncAppKeyExtractor::new(None).extract(&req).unwrap();
        assert_eq!(key, SyncAppId::nil());
    }

    #[test]
    fn jwt_missing_app_claim_collapses_to_nil_sentinel() {
        let header = base64::engine::general_purpose::URL_SAFE_NO_PAD
            .encode(r#"{"alg":"HS256"}"#);
        let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD
            .encode(r#"{"sub":"user","iss":"test"}"#);
        let sig = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode("sig");
        let token = format!("{header}.{payload}.{sig}");

        let req = Request::builder()
            .header("authorization", format!("Bearer {token}"))
            .body(())
            .unwrap();
        let key = SyncAppKeyExtractor::new(None).extract(&req).unwrap();
        assert_eq!(key, SyncAppId::nil());
    }

    #[test]
    fn cf_connecting_ip_is_used_when_present() {
        let req = Request::builder()
            .header("cf-connecting-ip", "203.0.113.7")
            .body(())
            .unwrap();
        let key = CloudflareIpKeyExtractor.extract(&req).unwrap();
        assert_eq!(key, "203.0.113.7".parse::<std::net::IpAddr>().unwrap());
    }

    #[test]
    fn forged_x_forwarded_for_is_not_trusted() {
        // No cf-connecting-ip and no ConnectInfo (the prod case where Caddy is
        // bypassed). A client-supplied X-Forwarded-For must NOT become the key —
        // the fallback is the TCP peer (PeerIpKeyExtractor), which errors here
        // because no ConnectInfo extension is present. The key assertion is that
        // we do NOT return the spoofed 1.2.3.4.
        let req = Request::builder()
            .header("x-forwarded-for", "1.2.3.4")
            .header("x-real-ip", "1.2.3.4")
            .body(())
            .unwrap();
        let result = CloudflareIpKeyExtractor.extract(&req);
        assert!(
            result.is_err(),
            "forged XFF/X-Real-IP must not yield a per-IP bucket; got {result:?}"
        );
    }

    #[test]
    fn different_apps_get_different_keys() {
        let app1 = SyncAppId::new();
        let app2 = SyncAppId::new();

        let req1 = Request::builder()
            .header("authorization", format!("Bearer {}", fake_jwt(&app1)))
            .body(())
            .unwrap();
        let req2 = Request::builder()
            .header("authorization", format!("Bearer {}", fake_jwt(&app2)))
            .body(())
            .unwrap();

        let extractor = SyncAppKeyExtractor::new(None);
        let key1 = extractor.extract(&req1).unwrap();
        let key2 = extractor.extract(&req2).unwrap();
        assert_ne!(key1, key2);
    }
}