//! Per-user registry of outstanding viewer tokens.
//!
//! Minted viewer tokens ([`super::jwt::ViewerClaims`]) are short-lived
//! (`VIEWER_TOKEN_TTL_SECS`, 5 min) and session-scoped. The login-JWT blacklist
//! revokes the LOGIN token on logout, but a viewer token minted before logout
//! stays valid for the rest of its TTL — letting a just-logged-out user keep a
//! live remote-control plane until natural expiry.
//!
//! This registry closes that gap with the least-invasive mechanism: when a
//! viewer token is minted it is registered here under the minting user's `sub`;
//! on logout, the user's outstanding viewer tokens are taken from the registry
//! and added to the SAME exact-string blacklist the viewer WebSocket already
//! checks (`relay::viewer_ws_handler` -> `TokenBlacklist::is_revoked`). No
//! change to the WS path is required.
//!
//! Design notes:
//! - The registry holds bearer-token strings transiently in memory. The server
//!   already handles these tokens; they are never logged or otherwise exposed
//!   by this module.
//! - Entries are pruned-on-access (insert and take) by expiry, so the map can
//!   never grow unbounded even if a user never logs out — expired viewer tokens
//!   self-evict on the next operation that touches the map.
//! - Locking is a poison-tolerant `std::sync::Mutex` (matching the rate-limiter
//!   pattern: `lock().unwrap_or_else(|e| e.into_inner())`); a panic while the
//!   lock is held must not wedge logout/mint. No `.unwrap()` on lock paths.

use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

/// A single registered viewer token: the token string and the instant at which
/// it expires (mint time + TTL).
#[derive(Clone)]
struct Entry {
    token: String,
    expires_at: Instant,
}

/// Per-user registry of outstanding (un-expired, un-revoked) viewer tokens.
///
/// Cheap to clone (shares the inner map via `Arc`), so it can live in
/// `AppState` and be cloned with it like the other shared state.
#[derive(Clone, Default)]
pub struct ViewerTokenRegistry {
    /// Map of user `sub` -> that user's outstanding viewer tokens.
    inner: Arc<Mutex<HashMap<String, Vec<Entry>>>>,
}

impl ViewerTokenRegistry {
    /// Create an empty registry.
    pub fn new() -> Self {
        Self::default()
    }

    /// Register a freshly minted viewer token under its owning user (`sub`),
    /// expiring `ttl` from now.
    ///
    /// Prunes expired entries for this user on insert so the per-user vector and
    /// the overall map stay bounded. The token string is stored verbatim (it is
    /// the exact string the viewer WS blacklist-checks); it is never logged.
    pub fn register(&self, sub: &str, token: &str, ttl: Duration) {
        let now = Instant::now();
        let expires_at = now + ttl;
        let mut map = self.inner.lock().unwrap_or_else(|e| e.into_inner());

        let entries = map.entry(sub.to_string()).or_default();
        // Prune-on-insert: drop this user's already-expired tokens.
        entries.retain(|e| e.expires_at > now);
        entries.push(Entry {
            token: token.to_string(),
            expires_at,
        });
    }

    /// Remove and return all currently-valid viewer tokens for `sub`.
    ///
    /// Expired tokens are pruned (not returned) — they are already invalid at
    /// the JWT layer, so there is no value in blacklisting them. The user's
    /// entry is cleared from the map regardless (taking everything out). Returns
    /// an empty vec if the user has no registered tokens.
    pub fn take_for_user(&self, sub: &str) -> Vec<String> {
        let now = Instant::now();
        let mut map = self.inner.lock().unwrap_or_else(|e| e.into_inner());

        match map.remove(sub) {
            Some(entries) => entries
                .into_iter()
                .filter(|e| e.expires_at > now)
                .map(|e| e.token)
                .collect(),
            None => Vec::new(),
        }
    }

    /// Number of users currently tracked (test/observability helper). Does not
    /// prune; reflects the raw map size.
    #[cfg(test)]
    fn user_count(&self) -> usize {
        let map = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        map.len()
    }

    /// Number of registered (un-pruned) tokens for a user (test helper).
    #[cfg(test)]
    fn token_count(&self, sub: &str) -> usize {
        let map = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        map.get(sub).map(|v| v.len()).unwrap_or(0)
    }
}

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

    #[test]
    fn register_then_take_returns_and_clears() {
        let reg = ViewerTokenRegistry::new();
        reg.register("user-a", "tok-1", Duration::from_secs(300));
        reg.register("user-a", "tok-2", Duration::from_secs(300));

        assert_eq!(reg.token_count("user-a"), 2);

        let mut taken = reg.take_for_user("user-a");
        taken.sort();
        assert_eq!(taken, vec!["tok-1".to_string(), "tok-2".to_string()]);

        // Taking clears the user's entry.
        assert_eq!(reg.token_count("user-a"), 0);
        assert_eq!(reg.user_count(), 0);
        assert!(reg.take_for_user("user-a").is_empty());
    }

    #[test]
    fn take_for_unknown_user_is_empty() {
        let reg = ViewerTokenRegistry::new();
        assert!(reg.take_for_user("nobody").is_empty());
    }

    #[test]
    fn tokens_are_scoped_per_user() {
        let reg = ViewerTokenRegistry::new();
        reg.register("user-a", "a-tok", Duration::from_secs(300));
        reg.register("user-b", "b-tok", Duration::from_secs(300));

        assert_eq!(reg.take_for_user("user-a"), vec!["a-tok".to_string()]);
        // user-b is untouched.
        assert_eq!(reg.take_for_user("user-b"), vec!["b-tok".to_string()]);
    }

    #[test]
    fn expired_tokens_are_not_returned_by_take() {
        let reg = ViewerTokenRegistry::new();
        // Zero TTL -> already expired by the time take runs.
        reg.register("user-a", "expired", Duration::from_secs(0));
        reg.register("user-a", "live", Duration::from_secs(300));

        let taken = reg.take_for_user("user-a");
        assert_eq!(taken, vec!["live".to_string()]);
    }

    #[test]
    fn expired_tokens_are_pruned_on_insert() {
        let reg = ViewerTokenRegistry::new();
        // Register an already-expired token, then a live one for the same user.
        reg.register("user-a", "expired", Duration::from_secs(0));
        // First insert leaves 1 entry; the second insert prunes the expired one
        // before pushing, so only the live token remains.
        reg.register("user-a", "live", Duration::from_secs(300));

        assert_eq!(reg.token_count("user-a"), 1);
        assert_eq!(reg.take_for_user("user-a"), vec!["live".to_string()]);
    }

    #[test]
    fn multiple_logins_accumulate_then_all_taken() {
        // A user with several concurrent logins mints several viewer tokens;
        // logout must revoke ALL of them at once.
        let reg = ViewerTokenRegistry::new();
        for i in 0..5 {
            reg.register("user-a", &format!("tok-{i}"), Duration::from_secs(300));
        }
        assert_eq!(reg.token_count("user-a"), 5);
        assert_eq!(reg.take_for_user("user-a").len(), 5);
        assert_eq!(reg.user_count(), 0);
    }
}
