use crate::auth::domain::AuthClaims;
use crate::auth::traits::TokenStore;
use crate::common::errors::AuthError;
use std::collections::HashMap;
use std::sync::Arc;

pub struct InMemoryTokenStore {
    tokens: Arc<parking_lot::RwLock<HashMap<String, AuthClaims>>>,
    revoked: Arc<parking_lot::RwLock<HashMap<String, i64>>>,
}

impl InMemoryTokenStore {
    pub fn new() -> Self {
        Self {
            tokens: Arc::new(parking_lot::RwLock::new(HashMap::new())),
            revoked: Arc::new(parking_lot::RwLock::new(HashMap::new())),
        }
    }

    pub fn cleanup_expired(&self) {
        let now = chrono::Utc::now().timestamp();
        let mut revoked = self.revoked.write();
        revoked.retain(|_, exp| *exp > now);
    }
}

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

#[async_trait::async_trait]
impl TokenStore for InMemoryTokenStore {
    async fn store_token(&self, token: &str, claims: &AuthClaims) -> Result<(), AuthError> {
        let mut tokens = self.tokens.write();
        tokens.insert(token.to_string(), claims.clone());
        Ok(())
    }

    async fn get_token(&self, token: &str) -> Result<Option<AuthClaims>, AuthError> {
        let tokens = self.tokens.read();
        Ok(tokens.get(token).cloned())
    }

    async fn revoke_token(&self, token: &str) -> Result<(), AuthError> {
        let mut revoked = self.revoked.write();
        let exp = chrono::Utc::now().timestamp() + 3600;
        revoked.insert(token.to_string(), exp);
        Ok(())
    }

    async fn is_revoked(&self, token: &str) -> Result<bool, AuthError> {
        self.cleanup_expired();
        let revoked = self.revoked.read();
        Ok(revoked.contains_key(token))
    }
}