//! TOTP 2FA workflow tests: setup, confirm, login with TOTP, backup codes, disable.

use crate::harness::TestHarness;

// ── Helpers ──

/// Extract the TOTP secret from setup HTML (inside `<details>` > `<code>`).
fn extract_totp_secret(html: &str) -> String {
    let details_start = html.find("<details").expect("No <details> in TOTP setup HTML");
    let details_html = &html[details_start..];
    let code_start = details_html.find("<code").expect("No <code> in details");
    let after_tag = &details_html[code_start..];
    let content_start = after_tag.find('>').expect("No > after <code") + 1;
    let content_end = after_tag[content_start..]
        .find("</code>")
        .expect("No </code>");
    after_tag[content_start..content_start + content_end].to_string()
}

/// Extract backup codes from setup HTML (inside `<div class="backup-codes-grid">`).
fn extract_backup_codes(html: &str) -> Vec<String> {
    let marker = "backup-codes-grid";
    let grid_start = html.find(marker).expect("No backup-codes-grid in HTML");
    let grid_html = &html[grid_start..];
    let grid_end = grid_html.find("</div>").expect("No </div> for backup-codes-grid");
    let grid_content = &grid_html[..grid_end];

    let mut codes = Vec::new();
    let mut search = grid_content;
    while let Some(code_start) = search.find("<code>") {
        let content_start = code_start + "<code>".len();
        let content_end = search[content_start..]
            .find("</code>")
            .expect("Unclosed <code> in backup codes");
        codes.push(search[content_start..content_start + content_end].to_string());
        search = &search[content_start + content_end..];
    }

    codes
}

/// Generate a valid TOTP code from a base32 secret.
fn generate_totp_code(secret_base32: &str, email: &str) -> String {
    let bytes = totp_rs::Secret::Encoded(secret_base32.to_string())
        .to_bytes()
        .expect("Invalid TOTP secret");
    let totp = totp_rs::TOTP::new(
        totp_rs::Algorithm::SHA1,
        6,
        1,
        30,
        bytes,
        Some("Makenotwork".into()),
        email.into(),
    )
    .expect("TOTP creation failed");
    totp.generate_current().expect("TOTP generation failed")
}

/// Set up TOTP for the currently logged-in user and enable it.
/// Returns (secret_base32, backup_codes).
async fn setup_and_enable_totp(h: &mut TestHarness, email: &str) -> (String, Vec<String>) {
    let resp = h.client.post_form("/api/users/me/totp/setup", "").await;
    assert_eq!(resp.status.as_u16(), 200, "TOTP setup failed: {}", resp.text);

    let secret = extract_totp_secret(&resp.text);
    let codes = extract_backup_codes(&resp.text);
    assert!(!secret.is_empty(), "TOTP secret should not be empty");
    assert!(!codes.is_empty(), "Should have backup codes");

    let code = generate_totp_code(&secret, email);
    let resp = h
        .client
        .post_form("/api/users/me/totp/confirm", &format!("code={}", code))
        .await;
    assert_eq!(resp.status.as_u16(), 200, "TOTP confirm failed: {}", resp.text);

    (secret, codes)
}

/// Login flow when TOTP is enabled: POST /login -> 303 to /auth/2fa -> POST /auth/verify-2fa.
async fn login_with_2fa(h: &mut TestHarness, username: &str, password: &str, code: &str) {
    h.client.fetch_csrf_token().await;
    let resp = h
        .client
        .post_form(
            "/login",
            &format!(
                "login={}&password={}",
                urlencoding::encode(username),
                urlencoding::encode(password)
            ),
        )
        .await;
    // Login should redirect to 2FA page
    assert!(
        resp.status.is_redirection()
            || resp.text.contains("/auth/2fa")
            || resp.text.contains("HX-Redirect"),
        "Expected redirect to /auth/2fa, got {} — {}",
        resp.status,
        resp.text
    );

    // Load the 2FA page to get CSRF token
    let resp = h.client.get("/auth/2fa").await;
    assert_eq!(resp.status.as_u16(), 200, "2FA page failed: {}", resp.text);

    // Submit the code
    let resp = h
        .client
        .post_form("/auth/verify-2fa", &format!("code={}", code))
        .await;
    assert!(
        resp.status.is_redirection() || resp.status.is_success(),
        "2FA verification failed with {}: {}",
        resp.status,
        resp.text
    );
}

// ── Tests ──

#[tokio::test]
async fn totp_setup_and_confirm() {
    let mut h = TestHarness::new().await;
    h.signup("totp1", "totp1@test.com", "Password1!").await;

    let (_secret, codes) = setup_and_enable_totp(&mut h, "totp1@test.com").await;
    assert_eq!(codes.len(), 10, "Should have 10 backup codes");

    // Check status
    let resp = h.client.get("/api/users/me/totp/status").await;
    assert_eq!(resp.status.as_u16(), 200);
    assert!(
        resp.text.contains("true") || resp.text.contains("enabled") || resp.text.contains("Enabled"),
        "TOTP should be enabled: {}",
        resp.text
    );
}

#[tokio::test]
async fn totp_login_requires_2fa() {
    let mut h = TestHarness::new().await;
    let uid = h.signup("totp2", "totp2@test.com", "Password1!").await;
    let (secret, _) = setup_and_enable_totp(&mut h, "totp2@test.com").await;

    // Reset the anti-replay counter so the login TOTP code (same 30s window)
    // is not rejected as a replay of the confirm step.
    sqlx::query("UPDATE users SET totp_last_used_step = 0 WHERE id = $1")
        .bind(uid)
        .execute(&h.db)
        .await
        .expect("reset totp_last_used_step");

    // Logout
    h.client.post_form("/logout", "").await;

    // Login with TOTP
    let code = generate_totp_code(&secret, "totp2@test.com");
    login_with_2fa(&mut h, "totp2", "Password1!", &code).await;

    // Verify we're logged in
    let resp = h.client.get("/dashboard").await;
    assert_eq!(resp.status.as_u16(), 200, "Should be on dashboard after 2FA login");
}

#[tokio::test]
async fn totp_backup_code_login() {
    let mut h = TestHarness::new().await;
    h.signup("totp3", "totp3@test.com", "Password1!").await;
    let (_, codes) = setup_and_enable_totp(&mut h, "totp3@test.com").await;

    h.client.post_form("/logout", "").await;

    // Login using a backup code
    login_with_2fa(&mut h, "totp3", "Password1!", &codes[0]).await;

    let resp = h.client.get("/dashboard").await;
    assert_eq!(resp.status.as_u16(), 200, "Should be on dashboard after backup code login");
}

#[tokio::test]
async fn totp_backup_code_single_use() {
    let mut h = TestHarness::new().await;
    h.signup("totp4", "totp4@test.com", "Password1!").await;
    let (_, codes) = setup_and_enable_totp(&mut h, "totp4@test.com").await;
    let used_code = codes[0].clone();

    // Use backup code once
    h.client.post_form("/logout", "").await;
    login_with_2fa(&mut h, "totp4", "Password1!", &used_code).await;

    // Logout and try the same code again
    h.client.post_form("/logout", "").await;
    h.client.fetch_csrf_token().await;
    let resp = h
        .client
        .post_form(
            "/login",
            &format!("login=totp4&password={}", urlencoding::encode("Password1!")),
        )
        .await;
    assert!(
        resp.status.is_redirection()
            || resp.text.contains("/auth/2fa")
            || resp.text.contains("HX-Redirect"),
        "Expected redirect to 2FA"
    );

    let _resp = h.client.get("/auth/2fa").await;
    let resp = h
        .client
        .post_form("/auth/verify-2fa", &format!("code={}", used_code))
        .await;

    // Should fail — backup code already consumed
    assert!(
        resp.text.contains("Invalid") || resp.text.contains("error") || resp.text.contains("invalid"),
        "Used backup code should be rejected: {}",
        resp.text
    );
}

#[tokio::test]
async fn totp_disable_with_password() {
    let mut h = TestHarness::new().await;
    h.signup("totp5", "totp5@test.com", "Password1!").await;
    setup_and_enable_totp(&mut h, "totp5@test.com").await;

    // Disable TOTP
    let resp = h
        .client
        .post_form(
            "/api/users/me/totp/disable",
            &format!("password={}", urlencoding::encode("Password1!")),
        )
        .await;
    assert_eq!(resp.status.as_u16(), 200, "TOTP disable failed: {}", resp.text);

    // Verify status is disabled
    let resp = h.client.get("/api/users/me/totp/status").await;
    assert_eq!(resp.status.as_u16(), 200);
    assert!(
        resp.text.contains("false") || resp.text.contains("disabled") || resp.text.contains("Disabled")
            || !resp.text.contains("Enabled"),
        "TOTP should be disabled: {}",
        resp.text
    );

    // Logout and login — should not require 2FA
    h.client.post_form("/logout", "").await;
    h.login("totp5", "Password1!").await;

    let resp = h.client.get("/dashboard").await;
    assert_eq!(resp.status.as_u16(), 200, "Login after TOTP disable should not require 2FA");
}

#[tokio::test]
async fn totp_invalid_code_rejected() {
    let mut h = TestHarness::new().await;
    h.signup("totp6", "totp6@test.com", "Password1!").await;
    setup_and_enable_totp(&mut h, "totp6@test.com").await;

    h.client.post_form("/logout", "").await;

    // Start login
    h.client.fetch_csrf_token().await;
    let resp = h
        .client
        .post_form(
            "/login",
            &format!("login=totp6&password={}", urlencoding::encode("Password1!")),
        )
        .await;
    assert!(
        resp.status.is_redirection()
            || resp.text.contains("/auth/2fa")
            || resp.text.contains("HX-Redirect"),
        "Expected redirect to 2FA"
    );

    let _resp = h.client.get("/auth/2fa").await;

    // Submit an invalid code
    let resp = h
        .client
        .post_form("/auth/verify-2fa", "code=000000")
        .await;

    // Should show error, not redirect to dashboard
    assert!(
        resp.text.contains("Invalid") || resp.text.contains("invalid") || resp.text.contains("error"),
        "Invalid TOTP code should be rejected: {}",
        resp.text
    );
    assert!(
        !resp.text.contains("/dashboard") || resp.status.as_u16() == 200,
        "Should not redirect to dashboard with invalid code"
    );
}