use rusqlite::Connection;

pub fn init_db(conn: &Connection) {
    conn.execute(
        "CREATE TABLE IF NOT EXISTS users (
            username TEXT PRIMARY KEY,
            hashed_password TEXT NOT NULL,
            account_id INTEGER UNIQUE,
            balance INTEGER DEFAULT 1000
        )",
        [],
    )
    .unwrap();
    // Add column if not exists (for migration)
    let _ = conn.execute("ALTER TABLE users ADD COLUMN account_id INTEGER UNIQUE", []);
    let _ = conn.execute(
        "ALTER TABLE users ADD COLUMN balance INTEGER DEFAULT 1000",
        [],
    );
}

pub fn add_user(
    conn: &Connection,
    username: &str,
    hashed_password: &str,
) -> Result<(), Box<dyn std::error::Error>> {
    loop {
        let account_id: u32 = rand::random::<u32>();
        match conn.execute(
            "INSERT INTO users (username, hashed_password, account_id, balance) VALUES (?1, ?2, ?3, ?4)",
            [username, hashed_password, &account_id.to_string(), "1000"],
        ) {
            Ok(_) => break,
            Err(_) => continue, // Retry if account_id conflict
        }
    }
    Ok(())
}

pub fn get_hashed_password(conn: &Connection, username: &str) -> Option<String> {
    let mut stmt = conn
        .prepare("SELECT hashed_password FROM users WHERE username = ?1")
        .unwrap();
    let mut rows = stmt.query([username]).unwrap();
    if let Ok(Some(row)) = rows.next() {
        row.get(0).ok()
    } else {
        None
    }
}

pub fn get_account_id(conn: &Connection, username: &str) -> Option<u32> {
    let mut stmt = conn
        .prepare("SELECT account_id FROM users WHERE username = ?1")
        .unwrap();
    let mut rows = stmt.query([username]).unwrap();
    if let Ok(Some(row)) = rows.next() {
        Some(row.get::<_, u32>(0).ok()?)
    } else {
        None
    }
}

pub fn get_balance(conn: &Connection, username: &str) -> Option<i64> {
    let mut stmt = conn
        .prepare("SELECT balance FROM users WHERE username = ?1")
        .unwrap();
    let mut rows = stmt.query([username]).unwrap();
    if let Ok(Some(row)) = rows.next() {
        Some(row.get::<_, i64>(0).ok()?)
    } else {
        None
    }
}