//! TOTP 2FA management API: setup, confirm, disable, backup codes, status.

use axum::{
    extract::State,
    response::{Html, IntoResponse, Response},
    Form,
};
use serde::Deserialize;

use crate::{
    auth::{verify_password, AuthUser},
    constants::{BACKUP_CODE_COUNT, BACKUP_CODE_LENGTH, TOTP_DIGITS, TOTP_SKEW, TOTP_STEP},
    db,
    error::{AppError, Result, ResultExt},
    helpers::hx_toast,
    templates::{TotpSetupTemplate, TotpStatusTemplate},
    AppState,
};

/// Generate a TOTP secret, QR code, and backup codes (does not enable 2FA yet).
#[tracing::instrument(skip_all, name = "totp::setup")]
pub(super) async fn setup(
    State(state): State<AppState>,
    AuthUser(user): AuthUser,
) -> Result<Response> {
    user.check_not_sandbox()?;
    // Generate a 20-byte (160-bit) random secret
    use rand::Rng;
    let secret_bytes: Vec<u8> = (0..20).map(|_| rand::rng().random()).collect();

    let totp = totp_rs::TOTP::new(
        totp_rs::Algorithm::SHA1,
        TOTP_DIGITS,
        TOTP_SKEW,
        TOTP_STEP,
        secret_bytes,
        Some("Makenotwork".to_string()),
        user.email.clone(),
    )
    .context("totp generation")?;

    let secret_base32 = totp.get_secret_base32();

    // Store the secret (not yet enabled)
    db::totp::set_totp_secret(&state.db, user.id, &secret_base32).await?;

    // Generate QR code as base64 PNG
    let qr_base64 = totp
        .get_qr_base64()
        .map_err(|e| AppError::Internal(anyhow::anyhow!("qr code generation: {e}")))?;

    // Generate backup codes
    let backup_codes = generate_backup_codes();
    let code_hashes: Vec<String> = backup_codes
        .iter()
        .map(|code| hash_backup_code(code, &state.config.signing_secret))
        .collect();

    db::totp::create_backup_codes(&state.db, user.id, &code_hashes).await?;

    Ok(TotpSetupTemplate {
        qr_base64,
        secret_base32,
        backup_codes,
    }
    .into_response())
}

/// Verify the user's first TOTP code and enable 2FA.
#[derive(Deserialize)]
pub struct ConfirmForm {
    code: String,
}

#[tracing::instrument(skip_all, name = "totp::confirm")]
pub(super) async fn confirm(
    State(state): State<AppState>,
    AuthUser(user): AuthUser,
    Form(form): Form<ConfirmForm>,
) -> Result<Response> {
    let secret = db::totp::get_totp_secret(&state.db, user.id)
        .await?
        .ok_or_else(|| AppError::BadRequest("2FA setup not started".to_string()))?;

    let totp = build_totp(&secret, &user.email)?;

    if !totp.check_current(&form.code).map_err(|e| {
        AppError::Internal(anyhow::anyhow!("TOTP check failed: {}", e))
    })? {
        return Ok((
            [("HX-Retarget", "#totp-confirm-status"), ("HX-Reswap", "innerHTML")],
            Html("<span class=\"save-error\">Invalid code. Please try again.</span>"),
        )
            .into_response());
    }

    // Record the matched step to prevent replay on the very first login
    let now = chrono::Utc::now().timestamp() as u64;
    if let Some(step) = find_matching_step(&totp, &form.code, now) {
        db::totp::set_totp_last_used_step(&state.db, user.id, step).await?;
    }

    db::totp::enable_totp(&state.db, user.id).await?;

    Ok((
        [("HX-Trigger", hx_toast("Two-factor authentication enabled", "success"))],
        TotpStatusTemplate { enabled: true },
    )
        .into_response())
}

/// Disable 2FA (requires password confirmation).
#[derive(Deserialize)]
pub struct DisableForm {
    password: String,
}

#[tracing::instrument(skip_all, name = "totp::disable")]
pub(super) async fn disable(
    State(state): State<AppState>,
    AuthUser(user): AuthUser,
    Form(form): Form<DisableForm>,
) -> Result<Response> {
    // Verify password
    let db_user = db::users::get_user_by_id(&state.db, user.id)
        .await?
        .ok_or(AppError::Unauthorized)?;

    if !verify_password(&form.password, &db_user.password_hash)? {
        return Ok((
            [("HX-Retarget", "#totp-disable-status"), ("HX-Reswap", "innerHTML")],
            Html("<span class=\"save-error\">Incorrect password.</span>"),
        )
            .into_response());
    }

    db::totp::disable_totp(&state.db, user.id).await?;

    Ok((
        [("HX-Trigger", hx_toast("Two-factor authentication disabled", "success"))],
        TotpStatusTemplate { enabled: false },
    )
        .into_response())
}

/// Regenerate backup codes (requires password confirmation).
#[derive(Deserialize)]
pub struct RegenerateForm {
    password: String,
}

#[tracing::instrument(skip_all, name = "totp::regenerate_backup_codes")]
pub(super) async fn regenerate_backup_codes(
    State(state): State<AppState>,
    AuthUser(user): AuthUser,
    Form(form): Form<RegenerateForm>,
) -> Result<Response> {
    // Verify password
    let db_user = db::users::get_user_by_id(&state.db, user.id)
        .await?
        .ok_or(AppError::Unauthorized)?;

    if !verify_password(&form.password, &db_user.password_hash)? {
        return Ok((
            [("HX-Retarget", "#backup-regen-status"), ("HX-Reswap", "innerHTML")],
            Html("<span class=\"save-error\">Incorrect password.</span>"),
        )
            .into_response());
    }

    let backup_codes = generate_backup_codes();
    let code_hashes: Vec<String> = backup_codes
        .iter()
        .map(|code| hash_backup_code(code, &state.config.signing_secret))
        .collect();

    db::totp::create_backup_codes(&state.db, user.id, &code_hashes).await?;

    // Return the new codes as an HTML partial
    let codes_html: String = backup_codes
        .iter()
        .map(|c| format!("<code>{}</code>", c))
        .collect::<Vec<_>>()
        .join("\n");

    Ok((
        [("HX-Trigger", hx_toast("Backup codes regenerated", "success"))],
        Html(format!(
            "<div class=\"backup-codes-grid\">\n{}\n</div>\n<p style=\"opacity: 0.7; font-size: 0.85rem; margin-top: 0.75rem;\">Save these codes somewhere safe. Each code can only be used once.</p>",
            codes_html
        )),
    )
        .into_response())
}

/// Return the current 2FA status as an HTMX partial for the dashboard.
#[tracing::instrument(skip_all, name = "totp::status")]
pub(super) async fn status(
    State(state): State<AppState>,
    AuthUser(user): AuthUser,
) -> Result<Response> {
    let enabled = db::totp::is_totp_enabled(&state.db, user.id).await?;

    Ok(TotpStatusTemplate { enabled }.into_response())
}

// ── Helpers ─────────────────────────────────────────────────────────────────

/// Build a TOTP instance from a stored base32 secret.
pub(crate) fn build_totp(secret_base32: &str, account_name: &str) -> Result<totp_rs::TOTP> {
    let secret_bytes = totp_rs::Secret::Encoded(secret_base32.to_string())
        .to_bytes()
        .context("parse totp secret")?;

    totp_rs::TOTP::new(
        totp_rs::Algorithm::SHA1,
        TOTP_DIGITS,
        TOTP_SKEW,
        TOTP_STEP,
        secret_bytes,
        Some("Makenotwork".to_string()),
        account_name.to_string(),
    )
    .context("totp creation")
}

/// Generate random alphanumeric backup codes.
fn generate_backup_codes() -> Vec<String> {
    use rand::Rng;
    let mut rng = rand::rng();

    (0..BACKUP_CODE_COUNT)
        .map(|_| {
            (0..BACKUP_CODE_LENGTH)
                .map(|_| {
                    let idx: u8 = rng.random_range(0..36);
                    if idx < 10 {
                        (b'0' + idx) as char
                    } else {
                        (b'a' + idx - 10) as char
                    }
                })
                .collect()
        })
        .collect()
}

/// Find which TOTP time step a code matches, returning the step number.
///
/// This is used instead of `totp.check_current()` so we can store the
/// *matched* step for replay prevention, not just the wall-clock step.
/// Without this, a code valid for step N can be replayed at step N+1
/// within the skew window.
pub(crate) fn find_matching_step(totp: &totp_rs::TOTP, code: &str, time_secs: u64) -> Option<i64> {
    let base_step = time_secs / TOTP_STEP;
    let skew = TOTP_SKEW as u64;
    let start = base_step.saturating_sub(skew);
    for i in 0..=(skew * 2) {
        let step = start + i;
        let step_time = step * TOTP_STEP;
        let expected = totp.generate(step_time);
        if crate::crypto::constant_time_compare(&expected, code) {
            return Some(step as i64);
        }
    }
    None
}

/// Argon2id hash of a backup code.
///
/// Backup codes have only ~41 bits of entropy (8 alphanumeric chars), so the
/// previous HMAC-SHA256 scheme was brute-forceable in minutes if the DB and
/// the server's signing secret both leaked. Argon2id with even modest
/// parameters multiplies that work by ~10^5, putting offline attack in the
/// "needs a real GPU farm and time" range.
///
/// Uses lower-than-password params (8 MiB, 1 iteration) — backup codes are
/// random tokens, not user-chosen passwords, so the security floor is set by
/// the wordlist, not the hash function. Per-code unique salt.
pub(crate) fn hash_backup_code(code: &str, _secret: &str) -> String {
    use argon2::password_hash::{rand_core::OsRng, PasswordHasher, SaltString};
    use argon2::{Algorithm, Argon2, Params, Version};

    let salt = SaltString::generate(&mut OsRng);
    let params = Params::new(8 * 1024, 1, 1, None)
        .expect("argon2 backup-code params are valid");
    let argon2 = Argon2::new(Algorithm::Argon2id, Version::V0x13, params);
    let hash = argon2
        .hash_password(code.as_bytes(), &salt)
        .expect("argon2 backup-code hashing");
    hash.to_string()
}

/// Legacy HMAC-SHA256 hash kept for verifying pre-migration backup codes.
///
/// Used only as a fallback inside `verify_and_consume_backup_code` when the
/// stored hash isn't an Argon2 PHC string. Do not call from new code paths —
/// any newly issued backup code goes through `hash_backup_code` (Argon2).
pub(crate) fn legacy_hmac_backup_code(code: &str, secret: &str) -> String {
    use hmac::{Hmac, Mac};
    use sha2::Sha256;

    let mut mac =
        Hmac::<Sha256>::new_from_slice(secret.as_bytes()).expect("HMAC accepts any key length");
    mac.update(code.as_bytes());
    hex::encode(mac.finalize().into_bytes())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn backup_code_generation_produces_correct_count() {
        let codes = generate_backup_codes();
        assert_eq!(codes.len(), BACKUP_CODE_COUNT);
    }

    #[test]
    fn backup_codes_are_correct_length() {
        let codes = generate_backup_codes();
        for code in &codes {
            assert_eq!(code.len(), BACKUP_CODE_LENGTH);
        }
    }

    #[test]
    fn backup_codes_are_alphanumeric() {
        let codes = generate_backup_codes();
        for code in &codes {
            assert!(
                code.chars().all(|c| c.is_ascii_alphanumeric()),
                "Code should be alphanumeric: {}",
                code
            );
        }
    }

    #[test]
    fn backup_codes_are_unique() {
        let codes = generate_backup_codes();
        let unique: std::collections::HashSet<&String> = codes.iter().collect();
        assert_eq!(unique.len(), codes.len());
    }

    #[test]
    fn hash_backup_code_is_argon2_phc() {
        // Argon2 hashes are non-deterministic (unique salt per call) and use
        // the PHC string format that starts with `$argon2`.
        let h = hash_backup_code("abc12345", "secret");
        assert!(h.starts_with("$argon2"), "got {h}");
    }

    #[test]
    fn hash_backup_code_non_deterministic() {
        // Two hashes of the same code must differ — distinct salts.
        let h1 = hash_backup_code("abc12345", "secret");
        let h2 = hash_backup_code("abc12345", "secret");
        assert_ne!(h1, h2);
    }

    #[test]
    fn hash_backup_code_verifies_against_itself() {
        use argon2::{password_hash::PasswordVerifier, Argon2, PasswordHash};
        let h = hash_backup_code("abc12345", "ignored");
        let parsed = PasswordHash::new(&h).unwrap();
        assert!(Argon2::default().verify_password(b"abc12345", &parsed).is_ok());
        assert!(Argon2::default().verify_password(b"wrong", &parsed).is_err());
    }

    #[test]
    fn legacy_hmac_is_deterministic_and_secret_keyed() {
        let h1 = legacy_hmac_backup_code("abc12345", "secret");
        let h2 = legacy_hmac_backup_code("abc12345", "secret");
        assert_eq!(h1, h2);
        let h3 = legacy_hmac_backup_code("abc12345", "different-secret");
        assert_ne!(h1, h3);
    }
}