//! 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, decrypted (None if not set up).
///
/// Secrets are encrypted at rest with the global signing secret; this returns
/// the plaintext base32 seed. A stored value that isn't a valid `enc:v1:`
/// ciphertext (e.g. a pre-encryption plaintext seed, backwards compat was
/// cut) errors out, forcing the user to re-enroll their authenticator.
#[tracing::instrument(skip_all)]
pub(crate) async fn get_totp_secret(
    pool: &PgPool,
    user_id: UserId,
    signing_secret: &str,
) -> Result<Option<String>> {
    let stored: Option<String> = sqlx::query_scalar("SELECT totp_secret FROM users WHERE id = $1")
        .bind(user_id)
        .fetch_one(pool)
        .await?;

    stored
        .map(|s| crate::crypto::decrypt_totp_secret(&s, signing_secret))
        .transpose()
}

/// Store a TOTP secret for a user (does not enable 2FA yet).
///
/// The plaintext base32 seed is encrypted at rest with the signing secret
/// before it touches the database, so a DB read alone cannot recover a usable
/// second factor.
#[tracing::instrument(skip_all)]
pub(crate) async fn set_totp_secret(
    pool: &PgPool,
    user_id: UserId,
    secret: &str,
    signing_secret: &str,
) -> Result<()> {
    let encrypted = crate::crypto::encrypt_totp_secret(secret, signing_secret);

    // 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(encrypted)
        .execute(pool)
        .await?;

    Ok(())
}

/// Enable TOTP 2FA for a user (called after first successful code verification).
#[tracing::instrument(skip_all)]
pub(crate) 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(crate) 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(())
}

/// Atomically advance the last accepted TOTP time step. Returns `true` if this
/// call recorded the step (it was strictly newer than the stored one), `false`
/// if a concurrent verification already consumed this step or a later one.
///
/// The `AND $2 > totp_last_used_step` guard makes the read-then-write in the
/// caller race-free: two concurrent submissions of the same code both read the
/// old step and both match, but only the winner's UPDATE affects a row, so a
/// code cannot be accepted twice within its 30s window. Callers must gate
/// acceptance on the returned bool, not on their own prior step read.
#[tracing::instrument(skip_all)]
pub(crate) async fn set_totp_last_used_step(
    pool: &PgPool,
    user_id: UserId,
    step: i64,
) -> Result<bool> {
    let updated = sqlx::query_scalar::<_, UserId>(
        "UPDATE users SET totp_last_used_step = $2 \
         WHERE id = $1 AND $2 > COALESCE(totp_last_used_step, 0) RETURNING id",
    )
    .bind(user_id)
    .bind(step)
    .fetch_optional(pool)
    .await?;

    Ok(updated.is_some())
}

/// Check if a user has TOTP 2FA enabled.
#[tracing::instrument(skip_all)]
pub(crate) 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(crate) 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(crate) async fn verify_and_consume_backup_code(
    pool: &PgPool,
    user_id: UserId,
    code: &str,
    legacy_hmac: &str,
) -> Result<bool> {
    use argon2::{Argon2, PasswordHash, password_hash::PasswordVerifier};

    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.
    // Run the per-code Argon2 verifies on a blocking thread so a 2FA attempt
    // (up to N verifies) can't occupy a Tokio worker.
    let matched_id: Option<uuid::Uuid> = {
        let code = code.to_string();
        let legacy_hmac = legacy_hmac.to_string();
        tokio::task::spawn_blocking(move || {
            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 {
                    return Some(*id);
                }
            }
            None
        })
        .await
        .map_err(|e| anyhow::anyhow!("backup-code verify task join: {e}"))?
    };

    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)
}