//! Two-factor authentication verification page (login flow).

use axum::{
    extract::State,
    http::{header::HeaderMap, StatusCode},
    response::{IntoResponse, Redirect, Response},
    Form,
};
use serde::Deserialize;
use sqlx::PgPool;
use tower_sessions::{Expiry, Session};

use crate::{
    auth::{login_user, track_session, SessionUser},
    constants,
    db::{self, UserId},
    error::{AppError, Result, ResultExt},
    helpers::{get_csrf_token, is_htmx_request, spawn_email},
    routes::api::totp::{build_totp, find_matching_step},
    templates::*,
    AppState,
};

/// Session key for the pending 2FA user ID.
const PENDING_2FA_KEY: &str = "pending_2fa_user_id";
const PENDING_2FA_STARTED_AT: &str = "pending_2fa_started_at";
const PENDING_2FA_NOTIFY_ENABLED: &str = "pending_2fa_notify_enabled";
const PENDING_2FA_NOTIFY_EMAIL: &str = "pending_2fa_notify_email";
const PENDING_2FA_NOTIFY_NAME: &str = "pending_2fa_notify_name";
const PENDING_2FA_TRACKING_KEY: &str = "pending_2fa_tracking_id";

/// Clear every pending-2FA session key and delete the pending user_sessions
/// tracking row. Used both on successful login and when the pending state
/// expires, so a stale unattended browser can't sit "one TOTP away from
/// logged in" past `PENDING_2FA_TTL_SECS`.
async fn clear_pending_2fa(session: &Session, pool: &PgPool) {
    if let Ok(Some(tracking_id)) = session
        .get::<crate::db::UserSessionId>(PENDING_2FA_TRACKING_KEY)
        .await
        && let Err(e) = db::sessions::delete_pending_2fa_session(pool, tracking_id).await
    {
        tracing::warn!(error = ?e, "failed to delete pending_2fa tracking row");
    }
    session.remove::<UserId>(PENDING_2FA_KEY).await.ok();
    session.remove::<i64>(PENDING_2FA_STARTED_AT).await.ok();
    session.remove::<bool>(PENDING_2FA_NOTIFY_ENABLED).await.ok();
    session.remove::<String>(PENDING_2FA_NOTIFY_EMAIL).await.ok();
    session.remove::<String>(PENDING_2FA_NOTIFY_NAME).await.ok();
    session.remove::<bool>("pending_2fa_remember_me").await.ok();
    session.remove::<crate::db::UserSessionId>(PENDING_2FA_TRACKING_KEY).await.ok();
}

/// Check whether the pending-2FA state has aged past `PENDING_2FA_TTL_SECS`.
/// Missing `started_at` (older session pre-TTL) is treated as expired.
async fn pending_2fa_expired(session: &Session) -> bool {
    let started_at: Option<i64> = session.get(PENDING_2FA_STARTED_AT).await.ok().flatten();
    match started_at {
        Some(ts) => chrono::Utc::now().timestamp() - ts > constants::PENDING_2FA_TTL_SECS,
        None => true,
    }
}

/// Render the 2FA verification page (GET /auth/2fa).
#[tracing::instrument(skip_all, name = "two_factor::two_factor_page")]
pub(super) async fn two_factor_page(
    State(state): State<AppState>,
    session: Session,
) -> Result<Response> {
    // Verify the user is in a valid 2FA flow
    let user_id: UserId = session
        .get(PENDING_2FA_KEY)
        .await
        .context("session error")?
        .ok_or_else(|| AppError::BadRequest("No pending 2FA session".to_string()))?;

    if pending_2fa_expired(&session).await {
        clear_pending_2fa(&session, &state.db).await;
        return Err(AppError::BadRequest(
            "Your 2FA session expired. Please log in again.".to_string(),
        ));
    }

    // Confirm the pending_2fa tracking row still exists. If it was swept by
    // `delete_all_sessions_for_user` ("log out everywhere"), abort the flow.
    if let Some(tracking_id) = session
        .get::<crate::db::UserSessionId>(PENDING_2FA_TRACKING_KEY)
        .await
        .ok()
        .flatten()
        && !db::sessions::pending_2fa_session_exists(&state.db, tracking_id, user_id).await?
    {
        clear_pending_2fa(&session, &state.db).await;
        return Err(AppError::BadRequest(
            "Your session was revoked. Please log in again.".to_string(),
        ));
    }

    let csrf_token = get_csrf_token(&session).await;

    Ok(TwoFactorTemplate {
        csrf_token,
        session_user: None,
        error: None,
    }
    .into_response())
}

/// Form input for 2FA verification.
#[derive(Deserialize)]
pub struct VerifyTwoFactorForm {
    code: String,
}

/// Verify the TOTP or backup code and complete login (POST /auth/verify-2fa).
#[tracing::instrument(skip_all, name = "two_factor::verify_two_factor")]
pub(super) async fn verify_two_factor(
    State(state): State<AppState>,
    headers: HeaderMap,
    session: Session,
    Form(form): Form<VerifyTwoFactorForm>,
) -> Result<Response> {
    let is_htmx = is_htmx_request(&headers);

    let user_id: UserId = session
        .get(PENDING_2FA_KEY)
        .await
        .context("session error")?
        .ok_or_else(|| AppError::BadRequest("No pending 2FA session".to_string()))?;

    if pending_2fa_expired(&session).await {
        clear_pending_2fa(&session, &state.db).await;
        return Err(AppError::BadRequest(
            "Your 2FA session expired. Please log in again.".to_string(),
        ));
    }

    // Confirm the pending_2fa tracking row still exists (see two_factor_page
    // for rationale). Rejecting here closes the "phisher mid-2FA-prompt"
    // window even when the legitimate user's "log out everywhere" landed
    // between page render and code submission.
    if let Some(tracking_id) = session
        .get::<crate::db::UserSessionId>(PENDING_2FA_TRACKING_KEY)
        .await
        .ok()
        .flatten()
        && !db::sessions::pending_2fa_session_exists(&state.db, tracking_id, user_id).await?
    {
        clear_pending_2fa(&session, &state.db).await;
        return Err(AppError::BadRequest(
            "Your session was revoked. Please log in again.".to_string(),
        ));
    }

    let user = db::users::get_user_by_id(&state.db, user_id)
        .await?
        .ok_or(AppError::Unauthorized)?;

    // Re-check lockout status before attempting verification (account may have
    // been locked by a concurrent session since the 2FA page was shown)
    if let Some(locked_until) = user.locked_until
        && locked_until > chrono::Utc::now()
    {
        clear_pending_2fa(&session, &state.db).await;
        let remaining = (locked_until - chrono::Utc::now()).num_minutes() + 1;
        let csrf_token = get_csrf_token(&session).await;
        return Ok(TwoFactorTemplate {
            csrf_token,
            session_user: None,
            error: Some(format!(
                "Account is locked. Try again in {} minute(s).",
                remaining
            )),
        }
        .into_response());
    }

    let code = form.code.trim().to_string();
    let mut verified = false;

    // Try TOTP verification first (6-digit numeric codes)
    if let Some(ref secret) = user.totp_secret {
        let totp = build_totp(secret, &user.email)?;
        // Find the actual step that matched (not just wall-clock step) to
        // prevent replay across the skew boundary.
        let now = chrono::Utc::now().timestamp() as u64;
        let matched_step = find_matching_step(&totp, &code, now);
        if let Some(step) = matched_step {
            let last_step = db::totp::get_totp_last_used_step(&state.db, user_id).await?.unwrap_or(0);
            if step > last_step {
                db::totp::set_totp_last_used_step(&state.db, user_id, step).await?;
                verified = true;
            }
            // If step <= last_step, the code was already used — fall through to backup codes
        }
    }

    // If TOTP didn't match, try backup code. We pass both the raw code (for
    // Argon2 verify of newer rows) and the legacy HMAC (for pre-migration
    // rows that haven't been regenerated yet). Both are evaluated in
    // `verify_and_consume_backup_code` per row.
    if !verified {
        let legacy_hmac = crate::routes::api::totp::legacy_hmac_backup_code(
            &code, &state.config.signing_secret,
        );
        if db::totp::verify_and_consume_backup_code(
            &state.db, user_id, &code, &legacy_hmac,
        ).await? {
            verified = true;
        }
    }

    if !verified {
        // Track failed 2FA attempts toward account lockout (same counter as
        // failed password attempts — prevents brute-forcing 6-digit TOTP codes)
        db::auth::increment_failed_login(
            &state.db,
            user_id,
            constants::MAX_LOGIN_ATTEMPTS,
            constants::LOCKOUT_MINUTES,
        )
        .await?;

        // Check if this attempt triggered a lockout
        let user_after = db::users::get_user_by_id(&state.db, user_id).await?;
        if let Some(ref u) = user_after
            && let Some(locked_until) = u.locked_until
            && locked_until > chrono::Utc::now()
        {
            // Clear the 2FA flow — account is now locked
            clear_pending_2fa(&session, &state.db).await;
            let remaining = (locked_until - chrono::Utc::now()).num_minutes() + 1;
            let csrf_token = get_csrf_token(&session).await;
            return Ok(TwoFactorTemplate {
                csrf_token,
                session_user: None,
                error: Some(format!(
                    "Too many failed attempts. Account locked for {} minute(s).",
                    remaining
                )),
            }
            .into_response());
        }

        let csrf_token = get_csrf_token(&session).await;
        return Ok(TwoFactorTemplate {
            csrf_token,
            session_user: None,
            error: Some("Invalid code. Please try again.".to_string()),
        }
        .into_response());
    }

    // Successful 2FA — reset failed login counter
    db::auth::reset_failed_login(&state.db, user_id).await?;

    // Retrieve stored notification info
    let notify_enabled: bool = session
        .get(PENDING_2FA_NOTIFY_ENABLED)
        .await
        .ok()
        .flatten()
        .unwrap_or(false);
    let notify_email: Option<String> = session
        .get(PENDING_2FA_NOTIFY_EMAIL)
        .await
        .ok()
        .flatten();
    let notify_name: Option<String> = session
        .get(PENDING_2FA_NOTIFY_NAME)
        .await
        .ok()
        .flatten();

    // Retrieve remember-me preference
    let remember: bool = session
        .get("pending_2fa_remember_me")
        .await
        .ok()
        .flatten()
        .unwrap_or(false);

    // Clear pending 2FA state
    clear_pending_2fa(&session, &state.db).await;

    // Complete login
    let session_user = SessionUser::from_db_user(user, &state.db, state.config.admin_user_id).await;

    login_user(&session, session_user).await?;
    if !remember {
        session.set_expiry(Some(Expiry::OnSessionEnd));
    }
    track_session(&session, &state.db, user_id, &headers).await?;
    tracing::info!(user_id = %user_id, event = "login_2fa_success", "User completed 2FA login");

    // Send login notification (same as in auth.rs)
    if notify_enabled
        && let Some(email_addr) = notify_email
    {
        let session_count = match db::sessions::count_user_sessions(&state.db, user_id).await {
            Ok(n) => n,
            Err(e) => { tracing::warn!("Failed to count sessions for login notification: {e}"); 0 }
        };
        if session_count > 1 {
            let user_agent = headers
                .get("user-agent")
                .and_then(|v| v.to_str().ok())
                .map(|s| s.chars().take(constants::USER_AGENT_MAX_LENGTH).collect::<String>());
            let ip = crate::helpers::extract_client_ip(&headers);
            let unsub_url = crate::email::generate_unsubscribe_url(
                &state.config.host_url, user_id, crate::email::UnsubscribeAction::Login, &user_id.to_string(), &state.config.signing_secret,
            );
            spawn_email!(state, "login notification", |email| {
                email.send_new_login_notification(
                    &email_addr,
                    notify_name.as_deref(),
                    user_agent.as_deref(),
                    ip.as_deref(),
                    Some(&unsub_url),
                )
            });
        }
    }

    if is_htmx {
        return Ok((
            StatusCode::OK,
            [("HX-Redirect", "/dashboard")],
            "",
        )
            .into_response());
    }

    Ok(Redirect::to("/dashboard").into_response())
}