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

use anyhow::Result;
use mongodb::bson::oid::ObjectId;
use tokio::sync::RwLock;

use crate::db::MongoDb;

/// How long a cached token_version is considered fresh before re-querying Mongo.
const DEFAULT_TTL: Duration = Duration::from_secs(30);

/// Short-TTL in-memory cache of `{ user_id -> token_version }`.
///
/// This lets the JWT auth middleware reject access tokens whose `token_version`
/// claim is stale (e.g. issued before a password change) without doing a Mongo
/// lookup on every request. Stale entries are evicted on credential changes via
/// [`TokenVersionCache::invalidate`], so a version bump becomes visible within
/// at most `ttl` on other instances, and immediately on the instance that
/// processed the change.
#[derive(Clone)]
pub struct TokenVersionCache {
    map: Arc<RwLock<HashMap<ObjectId, (i32, Instant)>>>,
    ttl: Duration,
}

impl TokenVersionCache {
    pub fn new(ttl: Duration) -> Self {
        Self {
            map: Arc::new(RwLock::new(HashMap::new())),
            ttl,
        }
    }

    /// Build a cache with the default 30s TTL.
    pub fn with_default_ttl() -> Self {
        Self::new(DEFAULT_TTL)
    }

    /// Return the user's current `token_version`, loading it from Mongo on a
    /// cache miss or once the cached value is older than the TTL.
    ///
    /// Returns `Ok(None)` when the user no longer exists (deleted), so callers
    /// can reject the token as unauthorized.
    pub async fn get_or_load(&self, user_id: &ObjectId, db: &MongoDb) -> Result<Option<i32>> {
        // Fast path: serve from cache if fresh.
        {
            let cache = self.map.read().await;
            if let Some((version, fetched_at)) = cache.get(user_id) {
                if fetched_at.elapsed() < self.ttl {
                    return Ok(Some(*version));
                }
            }
        }

        // Slow path: query Mongo and refresh the cache.
        let version = match db.find_user_by_id(user_id).await? {
            Some(user) => Some(user.token_version),
            None => None,
        };

        let mut cache = self.map.write().await;
        match version {
            Some(v) => {
                cache.insert(*user_id, (v, Instant::now()));
            }
            None => {
                // User was deleted: drop any stale entry so we don't serve it.
                cache.remove(user_id);
            }
        }

        Ok(version)
    }

    /// Drop the cached entry for a user. Call this whenever a user's
    /// `token_version` is bumped (password change / recovery) so that the new
    /// version becomes visible immediately instead of waiting for TTL expiry.
    pub async fn invalidate(&self, user_id: &ObjectId) {
        let mut cache = self.map.write().await;
        cache.remove(user_id);
    }
}

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

    #[tokio::test]
    async fn invalidate_removes_entry() {
        // We can't easily exercise get_or_load without Mongo, but invalidate on
        // an empty cache is a no-op and must not panic.
        let cache = TokenVersionCache::with_default_ttl();
        cache.invalidate(&ObjectId::new()).await;
        assert!(cache.map.read().await.is_empty());
    }
}