momentry_studio/src-tauri/src/db.rs

use rusqlite::Connection;
use serde::{Deserialize, Serialize};
use std::path::PathBuf;
use std::sync::Mutex;
use argon2::{self, Algorithm, Argon2, Params, PasswordHash, PasswordHasher, PasswordVerifier, Version};
use argon2::password_hash::{SaltString, rand_core::OsRng};

const APP_TABLES: &str = "
CREATE TABLE IF NOT EXISTS search_history (
    id TEXT PRIMARY KEY,
    query TEXT NOT NULL,
    title TEXT NOT NULL,
    chat_state TEXT,
    mode TEXT DEFAULT 'keyword',
    pinned INTEGER DEFAULT 0,
    created_at TEXT DEFAULT (datetime('now')),
    updated_at TEXT DEFAULT (datetime('now'))
);

CREATE TABLE IF NOT EXISTS bookmarks (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    label TEXT NOT NULL,
    history_id TEXT,
    created_at TEXT DEFAULT (datetime('now'))
);

CREATE TABLE IF NOT EXISTS app_users (
    id TEXT PRIMARY KEY,
    username TEXT NOT NULL UNIQUE,
    password_hash TEXT,
    display_name TEXT,
    role TEXT DEFAULT 'user',
    created_at TEXT DEFAULT (datetime('now')),
    updated_at TEXT DEFAULT (datetime('now'))
);

CREATE INDEX IF NOT EXISTS idx_search_history_sort ON search_history(pinned DESC, updated_at DESC);
CREATE INDEX IF NOT EXISTS idx_bookmarks_created ON bookmarks(created_at DESC);
";

static DB_CONN: Mutex<Option<Connection>> = Mutex::new(None);

fn db_path() -> PathBuf {
    let base = std::env::var("MOMENTRY_DATA_DIR").unwrap_or_else(|_| "../data/users".to_string());
    std::fs::create_dir_all(&base).ok();
    PathBuf::from(base).join("demo.sqlite")
}

pub fn init_db() -> Result<(), String> {
    let path = db_path();
    let conn = Connection::open(&path).map_err(|e| format!("Failed to open db: {}", e))?;
    conn.execute_batch("PRAGMA journal_mode=WAL;").map_err(|e| format!("WAL mode error: {}", e))?;
    conn.execute_batch(APP_TABLES).map_err(|e| format!("Failed to create tables: {}", e))?;

    // Add password_hash column if not exists (migration)
    conn.execute("ALTER TABLE app_users ADD COLUMN password_hash TEXT", []).ok();

    let default_id = "demo";
    let default_user = "demo";
    let default_password = "demo123";
    
    // Check if demo user exists and has password
    let password_hash: Option<String> = conn
        .query_row(
            "SELECT password_hash FROM app_users WHERE id = ?1",
            [default_id],
            |row| row.get(0),
        )
        .unwrap_or(None);

    if password_hash.is_none() || password_hash.as_ref().map_or(true, |h| h.is_empty()) {
        let password_hash = hash_password(default_password);
        conn.execute(
            "INSERT OR REPLACE INTO app_users (id, username, password_hash, display_name, role) VALUES (?1, ?2, ?3, ?4, 'admin')",
            rusqlite::params![default_id, default_user, password_hash, "Demo User"],
        )
        .map_err(|e| format!("Failed to insert default user: {}", e))?;
    }

    let mut lock = DB_CONN.lock().map_err(|e| format!("DB lock error: {}", e))?;
    *lock = Some(conn);
    Ok(())
}

fn get_conn() -> Result<Connection, String> {
    let path = db_path();
    let conn = Connection::open(&path).map_err(|e| format!("Failed to open db: {}", e))?;
    conn.execute_batch("PRAGMA journal_mode=WAL;").map_err(|e| format!("WAL mode error: {}", e))?;
    Ok(conn)
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct HistoryItem {
    pub id: String,
    pub query: String,
    pub title: String,
    pub chat_state: Option<String>,
    pub mode: Option<String>,
    pub pinned: bool,
    pub created_at: Option<String>,
    pub updated_at: Option<String>,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct BookmarkItem {
    pub id: i64,
    pub label: String,
    pub history_id: Option<String>,
    pub created_at: Option<String>,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct User {
    pub id: String,
    pub username: String,
    pub display_name: Option<String>,
    pub role: String,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct LoginResult {
    pub success: bool,
    pub user: Option<User>,
    pub token: Option<String>,
    pub error: Option<String>,
}

fn hash_password(password: &str) -> String {
    let salt = SaltString::generate(&mut OsRng);
    let argon2 = Argon2::new(Algorithm::Argon2id, Version::V0x13, Params::default());
    argon2.hash_password(password.as_bytes(), &salt).unwrap().to_string()
}

fn verify_password(hash: &str, password: &str) -> bool {
    let parsed_hash = PasswordHash::new(hash);
    match parsed_hash {
        Ok(h) => {
            let argon2 = Argon2::default();
            argon2.verify_password(password.as_bytes(), &h).is_ok()
        }
        Err(_) => false,
    }
}

#[tauri::command(rename_all = "camelCase")]
pub fn get_search_history(limit: Option<u32>) -> Result<Vec<HistoryItem>, String> {
    let limit = limit.unwrap_or(30).min(30);
    let conn = get_conn()?;
    let mut stmt = conn
        .prepare("SELECT id, query, title, chat_state, mode, pinned, created_at, updated_at FROM search_history ORDER BY pinned DESC, updated_at DESC LIMIT ?1")
        .map_err(|e| format!("Prepare error: {}", e))?;
    let items = stmt
        .query_map([limit], |row| {
            Ok(HistoryItem {
                id: row.get(0)?,
                query: row.get(1)?,
                title: row.get(2)?,
                chat_state: row.get(3)?,
                mode: row.get(4)?,
                pinned: row.get::<_, i64>(5)? != 0,
                created_at: row.get(6)?,
                updated_at: row.get(7)?,
            })
        })
        .map_err(|e| format!("Query error: {}", e))?
        .filter_map(|r| r.ok())
        .collect();
    Ok(items)
}

#[tauri::command(rename_all = "camelCase")]
pub fn save_search_history(
    id: String,
    query: String,
    title: String,
    chat_state: Option<String>,
    mode: Option<String>,
) -> Result<HistoryItem, String> {
    let conn = get_conn()?;
    let now = chrono::Utc::now().format("%Y-%m-%dT%H:%M:%SZ").to_string();
    let effective_mode = mode.unwrap_or_else(|| "keyword".to_string());
    conn.execute(
        "INSERT OR REPLACE INTO search_history (id, query, title, chat_state, mode, pinned, updated_at) VALUES (?1, ?2, ?3, ?4, ?5, (SELECT COALESCE((SELECT pinned FROM search_history WHERE id = ?1), 0)), ?6)",
        rusqlite::params![id, query, title, chat_state, effective_mode, now],
    )
    .map_err(|e| format!("Insert error: {}", e))?;

    conn.execute(
        "DELETE FROM search_history WHERE id NOT IN (SELECT id FROM search_history ORDER BY pinned DESC, updated_at DESC LIMIT 30)",
        [],
    )
    .ok();

    Ok(HistoryItem {
        id,
        query: query.clone(),
        title: title.clone(),
        chat_state,
        mode: Some(effective_mode),
        pinned: false,
        created_at: None,
        updated_at: Some(now),
    })
}

#[tauri::command(rename_all = "camelCase")]
pub fn rename_search_history(id: String, title: String) -> Result<(), String> {
    let conn = get_conn()?;
    let now = chrono::Utc::now().format("%Y-%m-%dT%H:%M:%SZ").to_string();
    conn.execute(
        "UPDATE search_history SET title = ?1, updated_at = ?2 WHERE id = ?3",
        rusqlite::params![title, now, id],
    )
    .map_err(|e| format!("Rename error: {}", e))?;
    Ok(())
}

#[tauri::command(rename_all = "camelCase")]
pub fn pin_search_history(id: String, pinned: bool) -> Result<(), String> {
    let conn = get_conn()?;
    let now = chrono::Utc::now().format("%Y-%m-%dT%H:%M:%SZ").to_string();
    conn.execute(
        "UPDATE search_history SET pinned = ?1, updated_at = ?2 WHERE id = ?3",
        rusqlite::params![pinned as i64, now, id],
    )
    .map_err(|e| format!("Pin error: {}", e))?;
    Ok(())
}

#[tauri::command(rename_all = "camelCase")]
pub fn delete_search_history(id: String) -> Result<(), String> {
    let conn = get_conn()?;
    conn.execute("DELETE FROM search_history WHERE id = ?1", [id])
        .map_err(|e| format!("Delete error: {}", e))?;
    Ok(())
}

#[tauri::command(rename_all = "camelCase")]
pub fn get_bookmarks() -> Result<Vec<BookmarkItem>, String> {
    let conn = get_conn()?;
    let mut stmt = conn
        .prepare("SELECT id, label, history_id, created_at FROM bookmarks ORDER BY created_at DESC")
        .map_err(|e| format!("Prepare error: {}", e))?;
    let items = stmt
        .query_map([], |row| {
            Ok(BookmarkItem {
                id: row.get(0)?,
                label: row.get(1)?,
                history_id: row.get(2)?,
                created_at: row.get(3)?,
            })
        })
        .map_err(|e| format!("Query error: {}", e))?
        .filter_map(|r| r.ok())
        .collect();
    Ok(items)
}

#[tauri::command(rename_all = "camelCase")]
pub fn save_bookmark(label: String, history_id: Option<String>) -> Result<BookmarkItem, String> {
    let conn = get_conn()?;
    conn.execute(
        "INSERT INTO bookmarks (label, history_id) VALUES (?1, ?2)",
        rusqlite::params![label, history_id],
    )
    .map_err(|e| format!("Insert error: {}", e))?;
    let id = conn.last_insert_rowid();
    Ok(BookmarkItem {
        id,
        label,
        history_id,
        created_at: None,
    })
}

#[tauri::command(rename_all = "camelCase")]
pub fn delete_bookmark(id: i64) -> Result<(), String> {
    let conn = get_conn()?;
    conn.execute("DELETE FROM bookmarks WHERE id = ?1", [id])
        .map_err(|e| format!("Delete error: {}", e))?;
    Ok(())
}

pub fn verify_user(username: &str, password: &str) -> Result<Option<User>, String> {
    let conn = get_conn()?;
    let mut stmt = conn
        .prepare("SELECT id, username, password_hash, display_name, role FROM app_users WHERE username = ?1")
        .map_err(|e| format!("Prepare error: {}", e))?;
    
    let result = stmt
        .query_row([username], |row| {
            let id: String = row.get(0)?;
            let username: String = row.get(1)?;
            let password_hash: Option<String> = row.get(2)?;
            let display_name: Option<String> = row.get(3)?;
            let role: String = row.get(4)?;
            Ok((id, username, password_hash, display_name, role))
        });

    match result {
        Ok((id, uname, password_hash, display_name, role)) => {
            if let Some(hash) = password_hash {
                if verify_password(&hash, password) {
                    Ok(Some(User { id, username: uname, display_name, role }))
                } else {
                    Ok(None)
                }
            } else {
                Ok(None)
            }
        }
        Err(_) => Ok(None),
    }
}

pub fn get_user_by_username(username: &str) -> Result<Option<User>, String> {
    let conn = get_conn()?;
    let mut stmt = conn
        .prepare("SELECT id, username, display_name, role FROM app_users WHERE username = ?1")
        .map_err(|e| format!("Prepare error: {}", e))?;
    
    let result = stmt
        .query_row([username], |row| {
            Ok(User {
                id: row.get(0)?,
                username: row.get(1)?,
                display_name: row.get(2)?,
                role: row.get(3)?,
            })
        });

    match result {
        Ok(user) => Ok(Some(user)),
        Err(_) => Ok(None),
    }
}