//! Two-factor authentication queries: TOTP secrets, backup codes.

use sqlx::PgPool;

use super::UserId;
use crate::error::Result;

/// Get the stored TOTP secret for a user (None if not set up).
#[tracing::instrument(skip_all)]
pub async fn get_totp_secret(pool: &PgPool, user_id: UserId) -> Result<Option<String>> {
    let secret: Option<String> =
        sqlx::query_scalar("SELECT totp_secret FROM users WHERE id = $1")
            .bind(user_id)
            .fetch_one(pool)
            .await?;

    Ok(secret)
}

/// Store a TOTP secret for a user (does not enable 2FA yet).
#[tracing::instrument(skip_all)]
pub async fn set_totp_secret(pool: &PgPool, user_id: UserId, secret: &str) -> Result<()> {
    // Clear the replay step alongside the secret. Without this, a user who
    // disables and re-enables TOTP (potentially with a new secret) inherits a
    // stale `totp_last_used_step` and any first-attempt code in a lower step
    // window is false-rejected as a replay.
    sqlx::query("UPDATE users SET totp_secret = $2, totp_last_used_step = NULL WHERE id = $1")
        .bind(user_id)
        .bind(secret)
        .execute(pool)
        .await?;

    Ok(())
}

/// Enable TOTP 2FA for a user (called after first successful code verification).
#[tracing::instrument(skip_all)]
pub async fn enable_totp(pool: &PgPool, user_id: UserId) -> Result<()> {
    sqlx::query("UPDATE users SET totp_enabled = true WHERE id = $1")
        .bind(user_id)
        .execute(pool)
        .await?;

    Ok(())
}

/// Disable TOTP 2FA: clear the secret, set enabled to false, delete backup codes.
#[tracing::instrument(skip_all)]
pub async fn disable_totp(pool: &PgPool, user_id: UserId) -> Result<()> {
    sqlx::query("UPDATE users SET totp_secret = NULL, totp_enabled = false, totp_last_used_step = NULL WHERE id = $1")
        .bind(user_id)
        .execute(pool)
        .await?;

    sqlx::query("DELETE FROM backup_codes WHERE user_id = $1")
        .bind(user_id)
        .execute(pool)
        .await?;

    Ok(())
}

/// Get the last accepted TOTP time step for a user (for replay prevention).
#[tracing::instrument(skip_all)]
pub async fn get_totp_last_used_step(pool: &PgPool, user_id: UserId) -> Result<Option<i64>> {
    let step: Option<i64> =
        sqlx::query_scalar("SELECT totp_last_used_step FROM users WHERE id = $1")
            .bind(user_id)
            .fetch_one(pool)
            .await?;

    Ok(step)
}

/// Update the last accepted TOTP time step (set after successful verification).
#[tracing::instrument(skip_all)]
pub async fn set_totp_last_used_step(pool: &PgPool, user_id: UserId, step: i64) -> Result<()> {
    sqlx::query("UPDATE users SET totp_last_used_step = $2 WHERE id = $1")
        .bind(user_id)
        .bind(step)
        .execute(pool)
        .await?;

    Ok(())
}

/// Check if a user has TOTP 2FA enabled.
#[tracing::instrument(skip_all)]
pub async fn is_totp_enabled(pool: &PgPool, user_id: UserId) -> Result<bool> {
    let enabled: bool =
        sqlx::query_scalar("SELECT totp_enabled FROM users WHERE id = $1")
            .bind(user_id)
            .fetch_one(pool)
            .await?;

    Ok(enabled)
}

/// Delete existing backup codes and insert new ones (atomic replacement).
#[tracing::instrument(skip_all)]
pub async fn create_backup_codes(
    pool: &PgPool,
    user_id: UserId,
    code_hashes: &[String],
) -> Result<()> {
    let mut tx = pool.begin().await?;

    // Delete any existing codes
    sqlx::query("DELETE FROM backup_codes WHERE user_id = $1")
        .bind(user_id)
        .execute(&mut *tx)
        .await?;

    // Batch insert all codes in a single query
    sqlx::query(
        "INSERT INTO backup_codes (user_id, code_hash) SELECT $1, UNNEST($2::text[])",
    )
    .bind(user_id)
    .bind(code_hashes)
    .execute(&mut *tx)
    .await?;

    tx.commit().await?;
    Ok(())
}

/// Verify a backup code and mark it as used if found.
///
/// `code` is the raw 8-char token the user typed; `legacy_hmac` is the
/// HMAC-SHA256 of the same code (passed in pre-computed by the caller so the
/// secret stays in route-layer scope). Returns `Ok(true)` when a matching
/// unused code is consumed.
///
/// Dual-read window: rows hashed under the old HMAC scheme remain valid
/// until the user regenerates their backup codes (each regeneration writes
/// fresh Argon2 hashes). Argon2 PHC strings begin with `$argon2`; anything
/// else is treated as a legacy 64-char hex HMAC.
#[tracing::instrument(skip_all)]
pub async fn verify_and_consume_backup_code(
    pool: &PgPool,
    user_id: UserId,
    code: &str,
    legacy_hmac: &str,
) -> Result<bool> {
    use argon2::{password_hash::PasswordVerifier, Argon2, PasswordHash};

    let rows: Vec<(uuid::Uuid, String)> = sqlx::query_as(
        "SELECT id, code_hash FROM backup_codes WHERE user_id = $1 AND used_at IS NULL",
    )
    .bind(user_id)
    .fetch_all(pool)
    .await?;

    // Timing note: the loop `break`s on the first match, which is NOT a usable
    // timing oracle. A wrong guess (the attacker's case) matches nothing, so the
    // loop always runs to completion and scans every row in constant time —
    // independent of code ordering. The early exit fires only on a *successful*
    // verify, by which point the caller already supplied a valid code and has
    // nothing left to learn. Retaining the break also avoids forcing N Argon2
    // verifications (each ~46 MiB) on every attempt, which would hand an attacker
    // a memory-amplification lever on the 2FA endpoint. Brute force is bounded
    // separately by the shared failed-attempt lockout.
    let mut matched_id: Option<uuid::Uuid> = None;
    for (id, stored) in &rows {
        let is_match = if stored.starts_with("$argon2") {
            match PasswordHash::new(stored) {
                Ok(parsed) => Argon2::default()
                    .verify_password(code.as_bytes(), &parsed)
                    .is_ok(),
                Err(e) => {
                    tracing::warn!(error = %e, "malformed argon2 backup code hash in DB; skipping");
                    false
                }
            }
        } else {
            // Legacy HMAC-SHA256 hex. Length-equality short-circuits before
            // the constant-time compare, matching the existing behavior of
            // `crypto::constant_time_compare`.
            crate::crypto::constant_time_compare(stored, legacy_hmac)
        };
        if is_match {
            matched_id = Some(*id);
            break;
        }
    }

    let Some(id) = matched_id else {
        return Ok(false);
    };

    let result = sqlx::query(
        "UPDATE backup_codes SET used_at = NOW() WHERE id = $1 AND used_at IS NULL",
    )
    .bind(id)
    .execute(pool)
    .await?;

    Ok(result.rows_affected() > 0)
}