GopherGate/src/database/mod.rs

use anyhow::Result;
use sqlx::sqlite::{SqliteConnectOptions, SqlitePool};
use std::str::FromStr;
use tracing::info;

use crate::config::DatabaseConfig;

pub type DbPool = SqlitePool;

pub async fn init(config: &DatabaseConfig) -> Result<DbPool> {
    // Ensure the database directory exists
    if let Some(parent) = config.path.parent()
        && !parent.as_os_str().is_empty()
    {
        tokio::fs::create_dir_all(parent).await?;
    }

    let database_path = config.path.to_string_lossy().to_string();
    info!("Connecting to database at {}", database_path);

    let options = SqliteConnectOptions::from_str(&format!("sqlite:{}", database_path))?.create_if_missing(true);

    let pool = SqlitePool::connect_with(options).await?;

    // Run migrations
    run_migrations(&pool).await?;
    info!("Database migrations completed");

    Ok(pool)
}

async fn run_migrations(pool: &DbPool) -> Result<()> {
    // Create clients table if it doesn't exist
    sqlx::query(
        r#"
        CREATE TABLE IF NOT EXISTS clients (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            client_id TEXT UNIQUE NOT NULL,
            name TEXT,
            description TEXT,
            created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
            updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
            is_active BOOLEAN DEFAULT TRUE,
            rate_limit_per_minute INTEGER DEFAULT 60,
            total_requests INTEGER DEFAULT 0,
            total_tokens INTEGER DEFAULT 0,
            total_cost REAL DEFAULT 0.0
        )
        "#,
    )
    .execute(pool)
    .await?;

    // Create llm_requests table if it doesn't exist
    sqlx::query(
        r#"
        CREATE TABLE IF NOT EXISTS llm_requests (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
            client_id TEXT,
            provider TEXT,
            model TEXT,
            prompt_tokens INTEGER,
            completion_tokens INTEGER,
            total_tokens INTEGER,
            cost REAL,
            has_images BOOLEAN DEFAULT FALSE,
            status TEXT DEFAULT 'success',
            error_message TEXT,
            duration_ms INTEGER,
            request_body TEXT,
            response_body TEXT,
            FOREIGN KEY (client_id) REFERENCES clients(client_id) ON DELETE SET NULL
        )
        "#,
    )
    .execute(pool)
    .await?;

    // Create provider_configs table
    sqlx::query(
        r#"
        CREATE TABLE IF NOT EXISTS provider_configs (
            id TEXT PRIMARY KEY,
            display_name TEXT NOT NULL,
            enabled BOOLEAN DEFAULT TRUE,
            base_url TEXT,
            api_key TEXT,
            credit_balance REAL DEFAULT 0.0,
            low_credit_threshold REAL DEFAULT 5.0,
            updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
        )
        "#,
    )
    .execute(pool)
    .await?;

    // Create model_configs table
    sqlx::query(
        r#"
        CREATE TABLE IF NOT EXISTS model_configs (
            id TEXT PRIMARY KEY,
            provider_id TEXT NOT NULL,
            display_name TEXT,
            enabled BOOLEAN DEFAULT TRUE,
            prompt_cost_per_m REAL,
            completion_cost_per_m REAL,
            mapping TEXT,
            updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
            FOREIGN KEY (provider_id) REFERENCES provider_configs(id) ON DELETE CASCADE
        )
        "#,
    )
    .execute(pool)
    .await?;

    // Create users table for dashboard access
    sqlx::query(
        r#"
        CREATE TABLE IF NOT EXISTS users (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            username TEXT UNIQUE NOT NULL,
            password_hash TEXT NOT NULL,
            display_name TEXT,
            role TEXT DEFAULT 'admin',
            must_change_password BOOLEAN DEFAULT FALSE,
            created_at DATETIME DEFAULT CURRENT_TIMESTAMP
        )
        "#,
    )
    .execute(pool)
    .await?;

    // Create client_tokens table for DB-based token auth
    sqlx::query(
        r#"
        CREATE TABLE IF NOT EXISTS client_tokens (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            client_id TEXT NOT NULL,
            token TEXT NOT NULL UNIQUE,
            name TEXT DEFAULT 'default',
            is_active BOOLEAN DEFAULT TRUE,
            created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
            last_used_at DATETIME,
            FOREIGN KEY (client_id) REFERENCES clients(client_id) ON DELETE CASCADE
        )
        "#,
    )
    .execute(pool)
    .await?;

    // Add must_change_password column if it doesn't exist (migration for existing DBs)
    let _ = sqlx::query("ALTER TABLE users ADD COLUMN must_change_password BOOLEAN DEFAULT FALSE")
        .execute(pool)
        .await;

    // Add display_name column if it doesn't exist (migration for existing DBs)
    let _ = sqlx::query("ALTER TABLE users ADD COLUMN display_name TEXT")
        .execute(pool)
        .await;

    // Add cache token columns if they don't exist (migration for existing DBs)
    let _ = sqlx::query("ALTER TABLE llm_requests ADD COLUMN cache_read_tokens INTEGER DEFAULT 0")
        .execute(pool)
        .await;
    let _ = sqlx::query("ALTER TABLE llm_requests ADD COLUMN cache_write_tokens INTEGER DEFAULT 0")
        .execute(pool)
        .await;

    // Insert default admin user if none exists (default password: admin)
    let user_count: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM users").fetch_one(pool).await?;

    if user_count.0 == 0 {
        // 'admin' hashed with default cost (12)
        let default_admin_hash =
            bcrypt::hash("admin", 12).map_err(|e| anyhow::anyhow!("Failed to hash default password: {}", e))?;
        sqlx::query(
            "INSERT INTO users (username, password_hash, role, must_change_password) VALUES ('admin', ?, 'admin', TRUE)"
        )
        .bind(default_admin_hash)
        .execute(pool)
        .await?;
        info!("Created default admin user with password 'admin' (must change on first login)");
    }

    // Create indices
    sqlx::query("CREATE INDEX IF NOT EXISTS idx_clients_client_id ON clients(client_id)")
        .execute(pool)
        .await?;

    sqlx::query("CREATE INDEX IF NOT EXISTS idx_clients_created_at ON clients(created_at)")
        .execute(pool)
        .await?;

    sqlx::query("CREATE INDEX IF NOT EXISTS idx_llm_requests_timestamp ON llm_requests(timestamp)")
        .execute(pool)
        .await?;

    sqlx::query("CREATE INDEX IF NOT EXISTS idx_llm_requests_client_id ON llm_requests(client_id)")
        .execute(pool)
        .await?;

    sqlx::query("CREATE INDEX IF NOT EXISTS idx_llm_requests_provider ON llm_requests(provider)")
        .execute(pool)
        .await?;

    sqlx::query("CREATE INDEX IF NOT EXISTS idx_llm_requests_status ON llm_requests(status)")
        .execute(pool)
        .await?;

    sqlx::query("CREATE UNIQUE INDEX IF NOT EXISTS idx_client_tokens_token ON client_tokens(token)")
        .execute(pool)
        .await?;

    sqlx::query("CREATE INDEX IF NOT EXISTS idx_client_tokens_client_id ON client_tokens(client_id)")
        .execute(pool)
        .await?;

    // Insert default client if none exists
    sqlx::query(
        r#"
        INSERT OR IGNORE INTO clients (client_id, name, description)
        VALUES ('default', 'Default Client', 'Default client for anonymous requests')
        "#,
    )
    .execute(pool)
    .await?;

    Ok(())
}

pub async fn test_connection(pool: &DbPool) -> Result<()> {
    sqlx::query("SELECT 1").execute(pool).await?;
    Ok(())
}