markbase/markbase-core/src/auth.rs

use bcrypt::{hash, verify, DEFAULT_COST};
use chrono::{Duration, Utc};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use uuid::Uuid;

use crate::provider::DataProvider;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct User {
    pub user_id: String,
    pub username: String,
    pub password_hash: String,
    pub created_at: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Session {
    pub token: String,
    pub user_id: String,
    pub username: String,
    pub created_at: String,
    pub expires_at: String,
    pub groups: Vec<String>,
    pub permissions: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LoginRequest {
    pub username: String,
    pub password: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LoginResponse {
    pub token: String,
    pub expires_at: String,
    pub user_id: String,
    pub groups: Vec<String>,
    pub permissions: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AdminLoginRequest {
    pub username: String,
    pub password: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AdminSession {
    pub token: String,
    pub username: String,
    pub created_at: String,
    pub expires_at: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AdminLoginResponse {
    pub token: String,
    pub expires_at: String,
    pub username: String,
}

#[derive(Clone)]
pub struct AuthState {
    pub sessions: Arc<Mutex<HashMap<String, Session>>>,
    pub users: Arc<Mutex<HashMap<String, User>>>,
    pub auth_db: Option<crate::sync::AuthDb>,
    pub admin_sessions: Arc<Mutex<HashMap<String, AdminSession>>>,
    pub provider: Option<Arc<dyn DataProvider>>,
}

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

impl AuthState {
    pub fn new() -> Self {
        let mut users = HashMap::new();

        let password_hash = hash("demo123", DEFAULT_COST).unwrap();
        users.insert(
            "demo".to_string(),
            User {
                user_id: "demo".to_string(),
                username: "demo".to_string(),
                password_hash,
                created_at: Utc::now().format("%Y-%m-%dT%H:%M:%SZ").to_string(),
            },
        );

        AuthState {
            sessions: Arc::new(Mutex::new(HashMap::new())),
            users: Arc::new(Mutex::new(users)),
            auth_db: None,
            admin_sessions: Arc::new(Mutex::new(HashMap::new())),
            provider: None,
        }
    }

    pub fn with_sync(auth_db_path: &str) -> Self {
        let auth_db = crate::sync::AuthDb::new(auth_db_path).ok();

        AuthState {
            sessions: Arc::new(Mutex::new(HashMap::new())),
            users: Arc::new(Mutex::new(HashMap::new())),
            auth_db,
            admin_sessions: Arc::new(Mutex::new(HashMap::new())),
            provider: None,
        }
    }

    pub fn with_provider(provider: Box<dyn DataProvider>) -> Self {
        AuthState {
            sessions: Arc::new(Mutex::new(HashMap::new())),
            users: Arc::new(Mutex::new(HashMap::new())),
            auth_db: None,
            admin_sessions: Arc::new(Mutex::new(HashMap::new())),
            provider: Some(Arc::from(provider)),
        }
    }

    pub fn login(&self, username: &str, password: &str) -> Option<LoginResponse> {
        let users = self.users.lock().unwrap();
        let user = users.get(username)?;

        if verify(password, &user.password_hash).unwrap_or(false) {
            let token = Uuid::new_v4().to_string();
            let now = Utc::now();
            let expires_at = now + Duration::hours(24);

            let session = Session {
                token: token.clone(),
                user_id: user.user_id.clone(),
                username: user.username.clone(),
                created_at: now.format("%Y-%m-%dT%H:%M:%SZ").to_string(),
                expires_at: expires_at.format("%Y-%m-%dT%H:%M:%SZ").to_string(),
                groups: vec![],
                permissions: "{}".to_string(),
            };

            let mut sessions = self.sessions.lock().unwrap();
            sessions.insert(token.clone(), session);

            Some(LoginResponse {
                token,
                expires_at: expires_at.format("%Y-%m-%dT%H:%M:%SZ").to_string(),
                user_id: user.user_id.clone(),
                groups: vec![],
                permissions: "{}".to_string(),
            })
        } else {
            None
        }
    }

    pub fn admin_login(&self, username: &str, password: &str) -> Option<AdminLoginResponse> {
        if let Some(auth_db) = &self.auth_db {
            match auth_db.get_admin(username) {
                Ok(Some(admin)) if admin.status == 1 => {
                    if verify(password, &admin.password_hash).unwrap_or(false) {
                        let token = Uuid::new_v4().to_string();
                        let now = Utc::now();
                        let expires_at = now + Duration::hours(24);

                        let session = AdminSession {
                            token: token.clone(),
                            username: username.to_string(),
                            created_at: now.format("%Y-%m-%dT%H:%M:%SZ").to_string(),
                            expires_at: expires_at.format("%Y-%m-%dT%H:%M:%SZ").to_string(),
                        };

                        let mut admin_sessions = self.admin_sessions.lock().unwrap();
                        admin_sessions.insert(token.clone(), session);

                        log::info!("Admin {} logged in successfully", username);

                        Some(AdminLoginResponse {
                            token,
                            expires_at: expires_at.format("%Y-%m-%dT%H:%M:%SZ").to_string(),
                            username: username.to_string(),
                        })
                    } else {
                        log::warn!("Invalid password for admin {}", username);
                        None
                    }
                }
                Ok(Some(_)) => {
                    log::warn!("Admin {} is not active", username);
                    None
                }
                Ok(None) => {
                    log::warn!("Admin {} not found", username);
                    None
                }
                Err(e) => {
                    log::error!("Failed to get admin {}: {}", username, e);
                    None
                }
            }
        } else {
            log::warn!("Auth DB not available for admin login");
            None
        }
    }

    pub fn verify_admin_token(&self, token: &str) -> Option<AdminSession> {
        let admin_sessions = self.admin_sessions.lock().unwrap();

        if let Some(session) = admin_sessions.get(token) {
            let expires_at = chrono::DateTime::parse_from_rfc3339(&session.expires_at)
                .ok()
                .map(|dt| dt.with_timezone(&Utc));

            if let Some(exp) = expires_at {
                if Utc::now() < exp {
                    return Some(session.clone());
                } else {
                    log::warn!("Admin token {} has expired", token);
                }
            }
        }

        None
    }

    pub fn login_with_sync(&self, username: &str, password: &str) -> Option<LoginResponse> {
        // Prefer provider over auth_db
        if let Some(provider) = &self.provider {
            return self.login_with_provider(&**provider, username, password);
        }

        if let Some(auth_db) = &self.auth_db {
            let user = match auth_db.get_user(username) {
                Ok(Some(user)) => user,
                Ok(None) => {
                    log::warn!("User {} not found in auth database", username);
                    return None;
                }
                Err(e) => {
                    log::error!("Failed to get user {}: {}", username, e);
                    return None;
                }
            };

            if user.status != 1 {
                log::warn!("User {} is disabled", username);
                return None;
            }

            if verify(password, &user.password_hash).unwrap_or(false) {
                let groups = auth_db.get_user_groups(username).unwrap_or_default();
                let permissions = user.permissions.clone();

                let token = Uuid::new_v4().to_string();
                let now = Utc::now();
                let expires_at = now + Duration::hours(24);

                let session = Session {
                    token: token.clone(),
                    user_id: username.to_string(),
                    username: username.to_string(),
                    created_at: now.format("%Y-%m-%dT%H:%M:%SZ").to_string(),
                    expires_at: expires_at.format("%Y-%m-%dT%H:%M:%SZ").to_string(),
                    groups: groups.clone(),
                    permissions: permissions.clone(),
                };

                let mut sessions = self.sessions.lock().unwrap();
                sessions.insert(token.clone(), session);

                log::info!("User {} logged in successfully", username);

                Some(LoginResponse {
                    token,
                    expires_at: expires_at.format("%Y-%m-%dT%H:%M:%SZ").to_string(),
                    user_id: username.to_string(),
                    groups,
                    permissions,
                })
            } else {
                log::warn!("Invalid password for user {}", username);
                None
            }
        } else {
            self.login(username, password)
        }
    }

    fn login_with_provider(
        &self,
        provider: &dyn DataProvider,
        username: &str,
        password: &str,
    ) -> Option<LoginResponse> {
        match provider.get_user(username) {
            Ok(Some(user)) => {
                if user.status != 1 {
                    log::warn!("User {} is disabled or not found", username);
                    return None;
                }

                match provider.check_password(username, password) {
                    Ok(true) => {
                        let groups = provider.get_user_groups(username).unwrap_or_default();

                        let token = Uuid::new_v4().to_string();
                        let now = Utc::now();
                        let expires_at = now + Duration::hours(24);

                        let session = Session {
                            token: token.clone(),
                            user_id: username.to_string(),
                            username: username.to_string(),
                            created_at: now.format("%Y-%m-%dT%H:%M:%SZ").to_string(),
                            expires_at: expires_at.format("%Y-%m-%dT%H:%M:%SZ").to_string(),
                            groups: groups.clone(),
                            permissions: user.permissions.clone(),
                        };

                        let mut sessions = self.sessions.lock().unwrap();
                        sessions.insert(token.clone(), session);

                        log::info!("User {} logged in via DataProvider", username);

                        Some(LoginResponse {
                            token,
                            expires_at: expires_at.format("%Y-%m-%dT%H:%M:%SZ").to_string(),
                            user_id: username.to_string(),
                            groups,
                            permissions: user.permissions,
                        })
                    }
                    Ok(false) => {
                        log::warn!("Invalid password for user {}", username);
                        None
                    }
                    Err(e) => {
                        log::error!("Password check error for {}: {}", username, e);
                        None
                    }
                }
            }
            Ok(None) => {
                log::warn!("User {} not found", username);
                None
            }
            Err(e) => {
                log::error!("Provider error for {}: {}", username, e);
                None
            }
        }
    }

    pub fn verify_token(&self, token: &str) -> Option<Session> {
        let sessions = self.sessions.lock().unwrap();
        let session = sessions.get(token)?;

        let expires_at = chrono::DateTime::parse_from_rfc3339(&session.expires_at)
            .ok()?
            .with_timezone(&Utc);

        if Utc::now() > expires_at {
            return None;
        }

        Some(session.clone())
    }

    pub fn logout(&self, token: &str) -> bool {
        let mut sessions = self.sessions.lock().unwrap();
        sessions.remove(token).is_some()
    }

    pub fn create_user(&self, username: &str, password: &str) -> Result<String, String> {
        let mut users = self.users.lock().unwrap();

        if users.contains_key(username) {
            return Err("User already exists".to_string());
        }

        let password_hash = hash(password, DEFAULT_COST).map_err(|e| e.to_string())?;

        let user_id = Uuid::new_v4().to_string();
        let user = User {
            user_id: user_id.clone(),
            username: username.to_string(),
            password_hash,
            created_at: Utc::now().format("%Y-%m-%dT%H:%M:%SZ").to_string(),
        };

        users.insert(username.to_string(), user);
        Ok(user_id)
    }
}

// Authorization header parser
pub fn parse_auth_header(header: &str) -> Option<String> {
    if header.starts_with("Bearer ") {
        Some(header.trim_start_matches("Bearer ").to_string())
    } else {
        None
    }
}