prisma_auth 0.2.1

// src/storage.rs
use std::collections::HashMap;
use std::sync::RwLock;
use std::time::{SystemTime, UNIX_EPOCH};

use crate::{User, Role, AuthError};

/// Trait for user storage backends
pub trait UserStorage: Send + Sync {
    /// Create a new user
    fn create_user(&self, user: User) -> Result<(), AuthError>;

    /// Get user by username
    fn get_user(&self, username: &str) -> Result<Option<User>, AuthError>;

    /// Update user
    fn update_user(&self, user: User) -> Result<(), AuthError>;

    /// Delete user
    fn delete_user(&self, username: &str) -> Result<(), AuthError>;

    /// List all users
    fn list_users(&self) -> Result<Vec<User>, AuthError>;

    /// Create a role
    fn create_role(&self, role: Role) -> Result<(), AuthError>;

    /// Get role by name
    fn get_role(&self, name: &str) -> Result<Option<Role>, AuthError>;

    /// Update role
    fn update_role(&self, role: Role) -> Result<(), AuthError>;

    /// Delete role
    fn delete_role(&self, name: &str) -> Result<(), AuthError>;

    /// List all roles
    fn list_roles(&self) -> Result<Vec<Role>, AuthError>;

    /// Check if user has permission
    fn user_has_permission(&self, username: &str, permission: &str) -> Result<bool, AuthError>;
}

/// In-memory storage implementation
pub struct InMemoryStorage {
    users: RwLock<HashMap<String, User>>,
    roles: RwLock<HashMap<String, Role>>,
}

impl InMemoryStorage {
    pub fn new() -> Self {
        Self {
            users: RwLock::new(HashMap::new()),
            roles: RwLock::new(HashMap::new()),
        }
    }

    /// Initialize with default admin role
    pub fn with_default_roles(self) -> Self {
        let admin_role = Role {
            name: "admin".to_string(),
            permissions: vec!["*".to_string()], // Wildcard permission for admin
            description: Some("Full administrative access".to_string()),
        };

        let user_role = Role {
            name: "user".to_string(),
            permissions: vec!["read".to_string()],
            description: Some("Basic user access".to_string()),
        };

        if let Ok(mut roles) = self.roles.write() {
            roles.insert("admin".to_string(), admin_role);
            roles.insert("user".to_string(), user_role);
        }

        self
    }
}

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

impl UserStorage for InMemoryStorage {
    fn create_user(&self, user: User) -> Result<(), AuthError> {
        let mut users = self.users.write()
            .map_err(|_| AuthError::StorageError("Failed to acquire write lock".to_string()))?;

        if users.contains_key(&user.username) {
            return Err(AuthError::StorageError("User already exists".to_string()));
        }

        users.insert(user.username.clone(), user);
        Ok(())
    }

    fn get_user(&self, username: &str) -> Result<Option<User>, AuthError> {
        let users = self.users.read()
            .map_err(|_| AuthError::StorageError("Failed to acquire read lock".to_string()))?;

        Ok(users.get(username).cloned())
    }

    fn update_user(&self, user: User) -> Result<(), AuthError> {
        let mut users = self.users.write()
            .map_err(|_| AuthError::StorageError("Failed to acquire write lock".to_string()))?;

        if !users.contains_key(&user.username) {
            return Err(AuthError::StorageError("User not found".to_string()));
        }

        users.insert(user.username.clone(), user);
        Ok(())
    }

    fn delete_user(&self, username: &str) -> Result<(), AuthError> {
        let mut users = self.users.write()
            .map_err(|_| AuthError::StorageError("Failed to acquire write lock".to_string()))?;

        users.remove(username)
            .ok_or_else(|| AuthError::StorageError("User not found".to_string()))?;

        Ok(())
    }

    fn list_users(&self) -> Result<Vec<User>, AuthError> {
        let users = self.users.read()
            .map_err(|_| AuthError::StorageError("Failed to acquire read lock".to_string()))?;

        Ok(users.values().cloned().collect())
    }

    fn create_role(&self, role: Role) -> Result<(), AuthError> {
        let mut roles = self.roles.write()
            .map_err(|_| AuthError::StorageError("Failed to acquire write lock".to_string()))?;

        if roles.contains_key(&role.name) {
            return Err(AuthError::StorageError("Role already exists".to_string()));
        }

        roles.insert(role.name.clone(), role);
        Ok(())
    }

    fn get_role(&self, name: &str) -> Result<Option<Role>, AuthError> {
        let roles = self.roles.read()
            .map_err(|_| AuthError::StorageError("Failed to acquire read lock".to_string()))?;

        Ok(roles.get(name).cloned())
    }

    fn update_role(&self, role: Role) -> Result<(), AuthError> {
        let mut roles = self.roles.write()
            .map_err(|_| AuthError::StorageError("Failed to acquire write lock".to_string()))?;

        if !roles.contains_key(&role.name) {
            return Err(AuthError::StorageError("Role not found".to_string()));
        }

        roles.insert(role.name.clone(), role);
        Ok(())
    }

    fn delete_role(&self, name: &str) -> Result<(), AuthError> {
        let mut roles = self.roles.write()
            .map_err(|_| AuthError::StorageError("Failed to acquire write lock".to_string()))?;

        roles.remove(name)
            .ok_or_else(|| AuthError::StorageError("Role not found".to_string()))?;

        Ok(())
    }

    fn list_roles(&self) -> Result<Vec<Role>, AuthError> {
        let roles = self.roles.read()
            .map_err(|_| AuthError::StorageError("Failed to acquire read lock".to_string()))?;

        Ok(roles.values().cloned().collect())
    }

    fn user_has_permission(&self, username: &str, permission: &str) -> Result<bool, AuthError> {
        let user = self.get_user(username)?
            .ok_or_else(|| AuthError::StorageError("User not found".to_string()))?;

        // Check user's direct permissions first
        if user.permissions.contains(&permission.to_string()) || user.permissions.contains(&"*".to_string()) {
            return Ok(true);
        }

        // Check role permissions
        if let Ok(Some(role)) = self.get_role(&user.role) {
            if role.permissions.contains(&permission.to_string()) || role.permissions.contains(&"*".to_string()) {
                return Ok(true);
            }
        }

        Ok(false)
    }
}

/// Helper function to get current timestamp
pub fn current_timestamp() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs()
}

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

    #[test]
    fn test_in_memory_storage() {
        let storage = InMemoryStorage::new();

        let user = User {
            username: "test".to_string(),
            password_hash: "hash".to_string(),
            role: "user".to_string(),
            permissions: vec!["read".to_string()],
            created_at: current_timestamp(),
            last_login: None,
        };

        // Test user creation
        storage.create_user(user.clone()).unwrap();

        // Test user retrieval
        let retrieved = storage.get_user("test").unwrap().unwrap();
        assert_eq!(retrieved.username, "test");

        // Test duplicate user creation fails
        assert!(storage.create_user(user).is_err());
    }

    #[test]
    fn test_role_management() {
        let storage = InMemoryStorage::new();

        let role = Role {
            name: "editor".to_string(),
            permissions: vec!["read".to_string(), "write".to_string()],
            description: Some("Can read and write".to_string()),
        };

        // Test role creation
        storage.create_role(role.clone()).unwrap();

        // Test role retrieval
        let retrieved = storage.get_role("editor").unwrap().unwrap();
        assert_eq!(retrieved.name, "editor");
        assert_eq!(retrieved.permissions.len(), 2);

        // Test duplicate role creation fails
        assert!(storage.create_role(role).is_err());
    }

    #[test]
    fn test_permissions() {
        let storage = InMemoryStorage::new().with_default_roles();

        let user = User {
            username: "admin".to_string(),
            password_hash: "hash".to_string(),
            role: "admin".to_string(),
            permissions: vec![],
            created_at: current_timestamp(),
            last_login: None,
        };

        storage.create_user(user).unwrap();

        // Admin should have all permissions due to wildcard
        assert!(storage.user_has_permission("admin", "any_permission").unwrap());
    }
}