// SPDX-License-Identifier: PolyForm-Noncommercial-1.0.0

//! Token store — persistence layer for issued temporary tokens.
//!
//! The [`TokenStore`] trait abstracts over storage backends. The only current
//! implementation is [`InMemoryTokenStore`], backed by a `DashMap` with a
//! background reaper that evicts expired tokens every 60 seconds.
//!
//! # Design
//!
//! Tokens are indexed by their **opaque bearer value** (the string clients
//! send in the `Authorization: Bearer` header) for O(1) validation, *and* by
//! their **JTI** for O(1) revocation.

use std::net::IpAddr;
use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};

use dashmap::DashMap;
use rand::RngExt;
use serde::{Deserialize, Serialize};
use tracing::debug;

use super::oidc::VerifiedIdentity;

/// A temporary gateway token issued after OIDC verification.
#[derive(Clone, Serialize, Deserialize)]
pub struct TemporaryToken {
    /// Unique token identifier (used for revocation).
    pub jti: String,
    /// The opaque bearer token value (`mcpgw_<base64>`).
    pub token: String,
    /// Verified identity that requested this token.
    pub identity: VerifiedIdentity,
    /// Allowed scopes for this token.
    pub scopes: TokenScopes,
    /// Issued-at (Unix epoch seconds).
    pub iat: u64,
    /// Expires-at (Unix epoch seconds).
    pub exp: u64,
    /// Client IP at issuance (for audit).
    #[serde(skip_serializing_if = "Option::is_none")]
    pub client_ip: Option<IpAddr>,
}

// Manual `Debug` redacts the opaque bearer value (CWE-532, mirrors MIK-6733).
// A derived `Debug` would print a live, usable session credential into logs;
// the `jti` is shown instead for correlation/revocation.
impl std::fmt::Debug for TemporaryToken {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("TemporaryToken")
            .field("jti", &self.jti)
            .field("token", &"<redacted>")
            .field("identity", &self.identity)
            .field("scopes", &self.scopes)
            .field("iat", &self.iat)
            .field("exp", &self.exp)
            .field("client_ip", &self.client_ip)
            .finish()
    }
}

impl TemporaryToken {
    /// Returns `true` if the token has passed its expiry time.
    #[must_use]
    pub fn is_expired(&self) -> bool {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or(Duration::ZERO)
            .as_secs();
        now >= self.exp
    }
}

/// Scopes granted to a temporary token.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct TokenScopes {
    /// Allowed backends. Empty slice means "all backends".
    pub backends: Vec<String>,
    /// Allowed tools. Empty slice means "all tools on allowed backends".
    pub tools: Vec<String>,
    /// Per-identity rate limit (requests per minute; 0 = unlimited).
    pub rate_limit: u32,
}

/// Trait abstracting the token storage backend.
///
/// Implementations must be `Send + Sync` because the token store is shared
/// across async tasks.
#[async_trait::async_trait]
pub trait TokenStore: Send + Sync + 'static {
    /// Insert a newly-issued token.
    async fn insert(&self, token: TemporaryToken);

    /// Look up a token by its opaque bearer value.
    ///
    /// Returns `None` if the token does not exist or is expired/revoked.
    async fn get(&self, bearer: &str) -> Option<TemporaryToken>;

    /// Revoke a token by its JTI.
    ///
    /// Returns `true` if the token existed and was removed.
    async fn revoke_by_jti(&self, jti: &str) -> bool;

    /// Revoke all tokens for a given OIDC subject (e.g., on offboarding).
    async fn revoke_by_subject(&self, subject: &str) -> usize;

    /// Count active (non-expired) tokens for a given OIDC subject.
    async fn count_for_subject(&self, subject: &str) -> usize;

    /// Remove all expired tokens. Called periodically by the background reaper.
    async fn reap_expired(&self) -> usize;
}

/// In-memory token store backed by two `DashMap` indices.
///
/// - `by_bearer`: bearer value → `TemporaryToken`  (O(1) validation)
/// - `by_jti`:    JTI string  → bearer value       (O(1) revocation)
pub struct InMemoryTokenStore {
    by_bearer: DashMap<String, TemporaryToken>,
    by_jti: DashMap<String, String>,
}

impl InMemoryTokenStore {
    /// Create an empty token store.
    #[must_use]
    pub fn new() -> Self {
        Self {
            by_bearer: DashMap::new(),
            by_jti: DashMap::new(),
        }
    }

    /// Generate a cryptographically random opaque bearer token.
    ///
    /// Format: `mcpgw_<43-char URL-safe base64>` (256 bits of entropy).
    /// The `mcpgw_` prefix makes tokens greppable and detectable by secret
    /// scanners (e.g., `truffleHog`, GitHub secret scanning).
    #[must_use]
    pub fn generate_bearer() -> String {
        let random_bytes: [u8; 32] = rand::rng().random();
        format!(
            "mcpgw_{}",
            base64::Engine::encode(
                &base64::engine::general_purpose::URL_SAFE_NO_PAD,
                random_bytes,
            )
        )
    }

    /// Generate a UUID v4 JTI.
    #[must_use]
    pub fn generate_jti() -> String {
        uuid::Uuid::new_v4().to_string()
    }
}

impl Default for InMemoryTokenStore {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait::async_trait]
impl TokenStore for InMemoryTokenStore {
    async fn insert(&self, token: TemporaryToken) {
        let bearer = token.token.clone();
        let jti = token.jti.clone();
        self.by_bearer.insert(bearer.clone(), token);
        self.by_jti.insert(jti, bearer);
    }

    async fn get(&self, bearer: &str) -> Option<TemporaryToken> {
        let entry = self.by_bearer.get(bearer)?;
        let token = entry.clone();
        drop(entry);

        if token.is_expired() {
            // Lazy eviction: remove on access
            self.by_bearer.remove(bearer);
            self.by_jti.remove(&token.jti);
            debug!(jti = %token.jti, "Lazy-evicted expired token");
            return None;
        }

        Some(token)
    }

    async fn revoke_by_jti(&self, jti: &str) -> bool {
        if let Some((_, bearer)) = self.by_jti.remove(jti) {
            self.by_bearer.remove(&bearer);
            true
        } else {
            false
        }
    }

    async fn revoke_by_subject(&self, subject: &str) -> usize {
        let mut jtis_to_revoke = Vec::new();

        for entry in &self.by_bearer {
            if entry.value().identity.subject == subject {
                jtis_to_revoke.push(entry.value().jti.clone());
            }
        }

        let count = jtis_to_revoke.len();
        for jti in jtis_to_revoke {
            self.revoke_by_jti(&jti).await;
        }
        count
    }

    async fn count_for_subject(&self, subject: &str) -> usize {
        self.by_bearer
            .iter()
            .filter(|e| e.value().identity.subject == subject && !e.value().is_expired())
            .count()
    }

    async fn reap_expired(&self) -> usize {
        let expired_bearers: Vec<String> = self
            .by_bearer
            .iter()
            .filter(|e| e.value().is_expired())
            .map(|e| e.key().clone())
            .collect();

        let count = expired_bearers.len();
        for bearer in expired_bearers {
            if let Some((_, token)) = self.by_bearer.remove(&bearer) {
                self.by_jti.remove(&token.jti);
                debug!(jti = %token.jti, "Reaped expired token");
            }
        }
        count
    }
}

/// Spawn a background task that reaps expired tokens every `interval`.
///
/// The task exits when the `shutdown` receiver fires.
pub fn spawn_reaper(
    store: Arc<dyn TokenStore>,
    interval: Duration,
    mut shutdown: tokio::sync::broadcast::Receiver<()>,
) {
    tokio::spawn(async move {
        let mut ticker = tokio::time::interval(interval);
        loop {
            tokio::select! {
                _ = ticker.tick() => {
                    let reaped = store.reap_expired().await;
                    if reaped > 0 {
                        debug!(count = reaped, "Reaped expired temporary tokens");
                    }
                }
                _ = shutdown.recv() => {
                    debug!("Token reaper shutting down");
                    break;
                }
            }
        }
    });
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::key_server::oidc::VerifiedIdentity;

    fn make_identity(subject: &str, email: &str) -> VerifiedIdentity {
        VerifiedIdentity {
            subject: subject.to_string(),
            email: email.to_string(),
            name: None,
            groups: Vec::new(),
            issuer: "https://accounts.google.com".to_string(),
        }
    }

    fn make_token(subject: &str, email: &str, exp_offset_secs: i64) -> TemporaryToken {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_secs();

        #[allow(clippy::cast_sign_loss)]
        let exp = if exp_offset_secs >= 0 {
            now + exp_offset_secs as u64
        } else {
            now.saturating_sub((-exp_offset_secs) as u64)
        };

        TemporaryToken {
            jti: InMemoryTokenStore::generate_jti(),
            token: InMemoryTokenStore::generate_bearer(),
            identity: make_identity(subject, email),
            scopes: TokenScopes {
                backends: vec!["tavily".to_string()],
                tools: Vec::new(),
                rate_limit: 100,
            },
            iat: now,
            exp,
            client_ip: None,
        }
    }

    #[test]
    fn debug_output_redacts_bearer_value() {
        // CWE-532 / MIK-6733 sibling: {:?} must never leak the live bearer.
        let token = make_token("user-1", "u1@test.local", 3600);
        let raw = token.token.clone();
        let dbg = format!("{token:?}");
        assert!(!dbg.contains(&raw), "Debug leaked the bearer value: {dbg}");
        assert!(
            dbg.contains("<redacted>"),
            "expected redaction marker: {dbg}"
        );
        assert!(dbg.contains(&token.jti), "jti stays visible for revocation");
    }

    #[tokio::test]
    async fn insert_and_get_valid_token() {
        // GIVEN: a token store with one valid token
        let store = InMemoryTokenStore::new();
        let token = make_token("sub123", "alice@company.com", 3600);
        let bearer = token.token.clone();

        // WHEN: we insert then look it up by bearer
        store.insert(token).await;
        let found = store.get(&bearer).await;

        // THEN: the token is returned
        assert!(found.is_some());
        assert_eq!(found.unwrap().identity.email, "alice@company.com");
    }

    #[tokio::test]
    async fn get_returns_none_for_unknown_token() {
        // GIVEN: an empty token store
        let store = InMemoryTokenStore::new();

        // WHEN: we look up a non-existent bearer
        let found = store.get("mcpgw_nonexistent").await;

        // THEN: None is returned
        assert!(found.is_none());
    }

    #[tokio::test]
    async fn get_lazy_evicts_expired_token() {
        // GIVEN: a store with an already-expired token (exp = 1 second ago)
        let store = InMemoryTokenStore::new();
        let token = make_token("sub123", "alice@company.com", -1);
        let bearer = token.token.clone();

        store.insert(token).await;

        // WHEN: we try to retrieve it
        let found = store.get(&bearer).await;

        // THEN: it is evicted and None is returned
        assert!(found.is_none());
        // Both indices cleaned up
        assert_eq!(store.by_bearer.len(), 0);
        assert_eq!(store.by_jti.len(), 0);
    }

    #[tokio::test]
    async fn revoke_by_jti_removes_token() {
        // GIVEN: a store with one valid token
        let store = InMemoryTokenStore::new();
        let token = make_token("sub123", "alice@company.com", 3600);
        let jti = token.jti.clone();
        let bearer = token.token.clone();
        store.insert(token).await;

        // WHEN: we revoke by JTI
        let removed = store.revoke_by_jti(&jti).await;

        // THEN: returns true and token is gone
        assert!(removed);
        assert!(store.get(&bearer).await.is_none());
        assert_eq!(store.by_jti.len(), 0);
    }

    #[tokio::test]
    async fn revoke_by_jti_returns_false_for_unknown() {
        // GIVEN: empty store
        let store = InMemoryTokenStore::new();

        // WHEN: revoke non-existent JTI
        let removed = store.revoke_by_jti("nonexistent-jti").await;

        // THEN: returns false
        assert!(!removed);
    }

    #[tokio::test]
    async fn revoke_by_subject_removes_all_tokens() {
        // GIVEN: two tokens for alice, one for bob
        let store = InMemoryTokenStore::new();
        let t1 = make_token("alice-sub", "alice@company.com", 3600);
        let t2 = make_token("alice-sub", "alice@company.com", 3600);
        let t3 = make_token("bob-sub", "bob@company.com", 3600);

        store.insert(t1).await;
        store.insert(t2).await;
        store.insert(t3).await;

        // WHEN: revoke all for alice
        let count = store.revoke_by_subject("alice-sub").await;

        // THEN: 2 removed, bob's token untouched
        assert_eq!(count, 2);
        assert_eq!(store.by_bearer.len(), 1); // bob's token remains
    }

    #[tokio::test]
    async fn count_for_subject_counts_active_tokens() {
        // GIVEN: two valid + one expired token for the same subject
        let store = InMemoryTokenStore::new();
        let t1 = make_token("alice-sub", "alice@company.com", 3600);
        let t2 = make_token("alice-sub", "alice@company.com", 3600);
        let t_expired = make_token("alice-sub", "alice@company.com", -1);

        store.insert(t1).await;
        store.insert(t2).await;
        store.insert(t_expired).await;

        // WHEN: count active tokens for alice
        let count = store.count_for_subject("alice-sub").await;

        // THEN: only 2 active tokens (expired one not counted)
        assert_eq!(count, 2);
    }

    #[tokio::test]
    async fn reap_expired_removes_only_expired() {
        // GIVEN: one valid + two expired tokens
        let store = InMemoryTokenStore::new();
        let valid = make_token("sub1", "alice@company.com", 3600);
        let expired1 = make_token("sub2", "bob@company.com", -1);
        let expired2 = make_token("sub3", "carol@company.com", -10);

        store.insert(valid).await;
        store.insert(expired1).await;
        store.insert(expired2).await;

        // WHEN: reap expired
        let reaped = store.reap_expired().await;

        // THEN: 2 removed, 1 remains
        assert_eq!(reaped, 2);
        assert_eq!(store.by_bearer.len(), 1);
    }

    #[tokio::test]
    async fn generate_bearer_has_correct_prefix() {
        // GIVEN/WHEN: generate a bearer token
        let bearer = InMemoryTokenStore::generate_bearer();

        // THEN: it starts with the required prefix
        assert!(bearer.starts_with("mcpgw_"));
        // And has 256 bits of entropy (32 bytes = 43 base64url chars)
        assert!(bearer.len() > 40);
    }

    #[test]
    fn is_expired_returns_true_for_past_exp() {
        // GIVEN: token with exp in the past
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_secs();
        let token = TemporaryToken {
            jti: "test".to_string(),
            token: "mcpgw_test".to_string(),
            identity: make_identity("sub", "test@example.com"),
            scopes: TokenScopes::default(),
            iat: now - 7200,
            exp: now - 1,
            client_ip: None,
        };

        // THEN: is_expired returns true
        assert!(token.is_expired());
    }

    #[test]
    fn is_expired_returns_false_for_future_exp() {
        // GIVEN: token with exp in the future
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_secs();
        let token = TemporaryToken {
            jti: "test".to_string(),
            token: "mcpgw_test".to_string(),
            identity: make_identity("sub", "test@example.com"),
            scopes: TokenScopes::default(),
            iat: now,
            exp: now + 3600,
            client_ip: None,
        };

        // THEN: is_expired returns false
        assert!(!token.is_expired());
    }
}
