use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _};
use std::collections::{HashMap, HashSet};
use std::sync::RwLock;
use uuid::Uuid;

pub struct SessionStore(RwLock<HashMap<Uuid, HashSet<String>>>);

impl SessionStore {
    pub fn new() -> Self {
        SessionStore(RwLock::new(HashMap::new()))
    }

    pub fn generate_secret() -> String {
        let mut bytes = [0u8; 32];
        getrandom::getrandom(&mut bytes).expect("Failed to generate random bytes");
        BASE64.encode(bytes)
    }

    pub fn store(&self, user_id: Uuid, secret: String) {
        let mut store = self.0.write().unwrap();
        store
            .entry(user_id)
            .or_insert_with(HashSet::new)
            .insert(secret);
    }

    pub fn verify(&self, user_id: Uuid, secret: &str) -> bool {
        let store = self.0.read().unwrap();
        store
            .get(&user_id)
            .map_or(false, |secrets| secrets.contains(secret))
    }

    pub fn remove(&self, user_id: Uuid, secret: &str) {
        let mut store = self.0.write().unwrap();
        if let Some(secrets) = store.get_mut(&user_id) {
            secrets.remove(secret);
            // Clean up the user entry if no sessions remain
            if secrets.is_empty() {
                store.remove(&user_id);
            }
        }
    }
}