//! Session cookie, current user, CSRF checks and login redirect helpers.
//!
//! The browser holds the raw session token in `__Host-forge_session`. The database
//! stores only `hash_token(token)`, so a leaked DB row cannot be replayed as a cookie.

use std::time::Duration;

use forge_core::password::{hash_password, verify_password};
use forge_core::session::{new_csrf_token, new_session_token};
use forge_core::time::stamp_after;
use forge_core::token::{constant_time_eq, hash_token};
use forge_core::{Session, Store, User, secret, totp};
use topcoat::Result;
use topcoat::context::{Cx, app_context};
use topcoat::cookie::{Cookie, Cookies, SameSite, cookies, time};
use topcoat::router::error::{forbidden, see_other};
use topcoat::router::request::{headers, uri};

use crate::state::AppState;

/// Logical cookie name. `override_prefix_host()` stores it as `__Host-forge_session`.
pub const SESSION_COOKIE: &str = "forge_session";
/// Lifetime of a web session, both in the database and on the cookie.
pub const SESSION_TTL: Duration = Duration::from_secs(30 * 24 * 60 * 60);

/// Logical name of the pre-session CSRF cookie, used by `/login` and `/setup` before any
/// session exists. Stored as `__Host-forge_precsrf`.
pub const PRE_CSRF_COOKIE: &str = "forge_precsrf";
/// Lifetime of the pre-session CSRF cookie.
const PRE_CSRF_TTL: Duration = Duration::from_secs(60 * 60);

/// Logical name of the pending-login cookie set between the password and the second factor.
/// Stored as `__Host-forge_2fa`. Its value is sealed with the instance key.
pub const PENDING_2FA_COOKIE: &str = "forge_2fa";
/// How long a password-verified login may wait for its second factor.
pub const PENDING_2FA_TTL: Duration = Duration::from_secs(5 * 60);

/// Maximum accepted length of a `next` redirect target.
const MAX_NEXT_LEN: usize = 2048;

/// Cookie jar that forces the `__Host-` prefix (Secure, Path=/, no Domain).
fn session_jar(cx: &Cx) -> impl Cookies + '_ {
    cookies(cx).override_prefix_host()
}

/// Resolves the signed-in user from the session cookie.
///
/// Returns `None` when there is no cookie, the session is unknown or expired, or its
/// user no longer exists. Every successful lookup refreshes `last_seen`.
pub async fn current_user(cx: &Cx) -> forge_core::Result<Option<(User, Session)>> {
    let Some(cookie) = session_jar(cx).get(SESSION_COOKIE) else {
        return Ok(None);
    };
    let store = &app_context::<AppState>(cx).store;
    let Some(session) = store
        .session_by_token_hash(&hash_token(cookie.value()))
        .await?
    else {
        return Ok(None);
    };
    // An organization never signs in, even if a session row names one.
    let Some(user) = store
        .user_by_id(&session.user_id)
        .await?
        .filter(|user| !user.is_org())
    else {
        return Ok(None);
    };
    store.touch_session(&session.id).await?;
    Ok(Some((user, session)))
}

/// Creates a session for `user` and sets the session cookie on the response.
pub async fn login(cx: &Cx, user: &User) -> forge_core::Result<Session> {
    let store = &app_context::<AppState>(cx).store;
    let token = new_session_token();
    let session = store
        .create_session(
            &user.id,
            &hash_token(&token),
            &new_csrf_token(),
            &stamp_after(SESSION_TTL),
            user_agent(cx).as_deref(),
        )
        .await?;
    let max_age = time::Duration::seconds(i64::try_from(SESSION_TTL.as_secs()).unwrap_or(i64::MAX));
    session_jar(cx).add(
        Cookie::build((SESSION_COOKIE, token))
            .path("/")
            .secure(true)
            .http_only(true)
            .same_site(SameSite::Lax)
            .max_age(max_age)
            .build(),
    );
    Ok(session)
}

/// Deletes the session row and clears the cookie.
pub async fn logout(cx: &Cx, session: &Session) -> forge_core::Result<()> {
    app_context::<AppState>(cx)
        .store
        .delete_session(&session.id)
        .await?;
    session_jar(cx).remove(Cookie::build((SESSION_COOKIE, "")).path("/"));
    Ok(())
}

/// Constant-time comparison of the submitted CSRF token with the session's token.
pub fn check_csrf(session: &Session, submitted: &str) -> bool {
    // An empty stored token must never match, even an empty submission.
    !session.csrf_token.is_empty()
        && constant_time_eq(session.csrf_token.as_bytes(), submitted.as_bytes())
}

/// Returns the pre-session CSRF token for a form, issuing one when the browser has none.
///
/// The existing cookie is reused so that two open login tabs keep working. The same value
/// goes into a hidden `precsrf` field; the POST must echo it back.
pub fn pre_csrf_token(cx: &Cx) -> String {
    if let Some(cookie) = session_jar(cx).get(PRE_CSRF_COOKIE) {
        let value = cookie.value();
        if !value.is_empty() {
            return value.to_owned();
        }
    }
    let token = new_csrf_token();
    session_jar(cx).add(pre_csrf_cookie(token.clone()));
    token
}

/// True only when the submitted `precsrf` equals the cookie, compared in constant time.
pub fn check_pre_csrf(cx: &Cx, submitted: &str) -> bool {
    let Some(cookie) = session_jar(cx).get(PRE_CSRF_COOKIE) else {
        return false;
    };
    !cookie.value().is_empty() && constant_time_eq(cookie.value().as_bytes(), submitted.as_bytes())
}

/// Removes the pre-session CSRF cookie. Called after a successful login or setup.
pub fn clear_pre_csrf(cx: &Cx) {
    session_jar(cx).remove(Cookie::build((PRE_CSRF_COOKIE, "")).path("/"));
}

fn pre_csrf_cookie(token: String) -> Cookie<'static> {
    Cookie::build((PRE_CSRF_COOKIE, token))
        .path("/")
        .secure(true)
        .http_only(true)
        .same_site(SameSite::Strict)
        .max_age(time::Duration::seconds(
            i64::try_from(PRE_CSRF_TTL.as_secs()).unwrap_or(i64::MAX),
        ))
        .build()
}

/// Hashes a password with argon2 on the blocking pool, keeping the async runtime free.
pub async fn hash_password_async(password: String) -> forge_core::Result<String> {
    tokio::task::spawn_blocking(move || hash_password(&password))
        .await
        .map_err(|_| forge_core::Error::Crypto("password hashing task failed".into()))?
}

/// Verifies a password with argon2 on the blocking pool.
pub async fn verify_password_async(password: String, hash: String) -> forge_core::Result<bool> {
    tokio::task::spawn_blocking(move || verify_password(&password, &hash))
        .await
        .map_err(|_| forge_core::Error::Crypto("password verification task failed".into()))
}

/// Returns the signed-in user, or a 303 redirect to `/login?next=<current path>`.
pub async fn require_user(cx: &Cx) -> Result<(User, Session)> {
    match current_user(cx).await? {
        Some(viewer) => Ok(viewer),
        None => {
            let path = uri(cx).path_and_query().map_or("/", |p| p.as_str());
            let target = format!("/login?next={}", percent_encode(path));
            Err(see_other(target).into())
        }
    }
}

/// Requires a signed-in user and a valid `csrf` form field. Returns 403 on a bad token.
pub async fn require_csrf(cx: &Cx, submitted: &str) -> Result<(User, Session)> {
    let (user, session) = require_user(cx).await?;
    if check_csrf(&session, submitted) {
        Ok((user, session))
    } else {
        Err(forbidden().into())
    }
}

/// A login that has passed its password and waits for the second factor.
pub struct PendingLogin {
    pub user_id: String,
    pub next: String,
}

/// Seconds since the Unix epoch, or 0 if the clock is before it.
pub fn unix_now() -> u64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map_or(0, |d| d.as_secs())
}

/// Sets the pending-login cookie. Its payload is `expires|user_id|next`, sealed so the browser
/// cannot change the user or the expiry.
pub fn set_pending_login(
    cx: &Cx,
    key: &[u8; 32],
    user_id: &str,
    next: &str,
) -> forge_core::Result<()> {
    let expires = unix_now() + PENDING_2FA_TTL.as_secs();
    let sealed = secret::seal(key, format!("{expires}|{user_id}|{next}").as_bytes())?;
    let max_age = time::Duration::seconds(i64::try_from(PENDING_2FA_TTL.as_secs()).unwrap_or(0));
    session_jar(cx).add(
        Cookie::build((PENDING_2FA_COOKIE, sealed))
            .path("/")
            .secure(true)
            .http_only(true)
            .same_site(SameSite::Strict)
            .max_age(max_age)
            .build(),
    );
    Ok(())
}

/// The pending login from the cookie, or `None` when the cookie is missing, tampered with,
/// malformed or past its expiry.
pub fn pending_login(cx: &Cx, key: &[u8; 32]) -> Option<PendingLogin> {
    let cookie = session_jar(cx).get(PENDING_2FA_COOKIE)?;
    let plain = secret::open(key, cookie.value()).ok()?;
    let text = String::from_utf8(plain).ok()?;
    let mut parts = text.splitn(3, '|');
    let expires: u64 = parts.next()?.parse().ok()?;
    let user_id = parts.next()?;
    let next = parts.next()?;
    if unix_now() >= expires {
        return None;
    }
    Some(PendingLogin {
        user_id: user_id.to_owned(),
        next: next.to_owned(),
    })
}

/// Removes the pending-login cookie.
pub fn clear_pending_login(cx: &Cx) {
    session_jar(cx).remove(Cookie::build((PENDING_2FA_COOKIE, "")).path("/"));
}

/// True when the user has confirmed two-factor authentication.
pub async fn two_factor_enabled(store: &Store, user_id: &str) -> forge_core::Result<bool> {
    Ok(store
        .totp_state(user_id)
        .await?
        .is_some_and(|state| state.enabled))
}

/// Checks a second-factor input for a user with two-factor on. A six-digit input is a TOTP
/// code, checked with replay protection. Anything else is a recovery code, which is used up
/// on success. Returns false for any wrong or spent input.
pub async fn verify_second_factor(
    store: &Store,
    key: &[u8; 32],
    user_id: &str,
    input: &str,
) -> forge_core::Result<bool> {
    let input = input.trim();
    // An organization has no second factor and never passes one.
    if store
        .user_by_id(user_id)
        .await?
        .is_none_or(|user| user.is_org())
    {
        return Ok(false);
    }
    let Some(state) = store.totp_state(user_id).await? else {
        return Ok(false);
    };
    let Some(sealed) = state.sealed_secret.as_deref() else {
        return Ok(false);
    };
    if !state.enabled {
        return Ok(false);
    }
    let shared = secret::open(key, sealed)?;
    let is_totp = input.len() == totp::DIGITS && input.bytes().all(|b| b.is_ascii_digit());
    if is_totp {
        return match totp::verify(&shared, input, totp::step_for(unix_now()), state.last_step) {
            Some(step) => store.accept_totp_step(user_id, step).await,
            None => Ok(false),
        };
    }
    store
        .consume_recovery_code(user_id, &totp::recovery_code_hash(input))
        .await
}

/// Returns `next` only when it is a same-site relative path; otherwise `/`.
///
/// Accepts a single leading `/`, rejects `//` and `/\` (protocol-relative or browser
/// backslash tricks), rejects control characters and caps the length.
pub fn safe_next(next: Option<&str>) -> &str {
    match next {
        Some(n) if is_local_path(n) => n,
        _ => "/",
    }
}

fn is_local_path(path: &str) -> bool {
    path.len() <= MAX_NEXT_LEN
        && path.starts_with('/')
        && !path.starts_with("//")
        && !path.starts_with("/\\")
        && path.chars().all(|c| !c.is_control())
}

/// Percent-encodes everything except unreserved characters and `/`.
pub fn percent_encode(input: &str) -> String {
    let mut out = String::with_capacity(input.len());
    for b in input.bytes() {
        match b {
            b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' | b'/' => {
                out.push(char::from(b));
            }
            _ => out.push_str(&format!("%{b:02X}")),
        }
    }
    out
}

/// The request's `User-Agent`, truncated to 512 characters.
pub fn user_agent(cx: &Cx) -> Option<String> {
    headers(cx)
        .get(http::header::USER_AGENT)
        .and_then(|v| v.to_str().ok())
        .map(|s| s.chars().take(512).collect())
}

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

    fn session_with_csrf(csrf: &str) -> Session {
        Session {
            id: "s1".into(),
            user_id: "u1".into(),
            csrf_token: csrf.into(),
            created_at: String::new(),
            expires_at: String::new(),
            last_seen: String::new(),
            user_agent: None,
        }
    }

    #[test]
    fn csrf_accepts_only_the_exact_token() {
        let s = session_with_csrf("abc123");
        assert!(check_csrf(&s, "abc123"));
        assert!(!check_csrf(&s, "abc124"));
        assert!(!check_csrf(&s, "abc12"));
        assert!(!check_csrf(&s, ""));
    }

    #[test]
    fn csrf_rejects_empty_session_token() {
        let s = session_with_csrf("");
        assert!(!check_csrf(&s, ""));
        assert!(!check_csrf(&s, "x"));
    }

    #[test]
    fn next_allows_local_paths() {
        assert_eq!(safe_next(Some("/")), "/");
        assert_eq!(safe_next(Some("/alice/repo")), "/alice/repo");
        assert_eq!(safe_next(Some("/settings/keys?x=1")), "/settings/keys?x=1");
    }

    #[test]
    fn next_rejects_open_redirects() {
        for bad in [
            "",
            "https://evil.example/",
            "//evil.example/",
            "/\\evil.example/",
            "evil.example",
            "\\\\evil",
            "/\n/evil",
            "/\t",
        ] {
            assert_eq!(safe_next(Some(bad)), "/", "input {bad:?}");
        }
        assert_eq!(safe_next(None), "/");
    }

    #[test]
    fn next_rejects_overlong_paths() {
        let long = format!("/{}", "a".repeat(MAX_NEXT_LEN));
        assert_eq!(safe_next(Some(&long)), "/");
    }

    #[test]
    fn percent_encode_keeps_slashes_and_encodes_the_rest() {
        assert_eq!(percent_encode("/a/b-c_d.e~f"), "/a/b-c_d.e~f");
        assert_eq!(percent_encode("/s?x=1&y=2"), "/s%3Fx%3D1%26y%3D2");
        assert_eq!(percent_encode("/a b"), "/a%20b");
    }

    /// An organization has no second factor and never passes one, even with a valid TOTP code or
    /// a recovery code on its row. A person with the same setup does pass.
    #[tokio::test]
    async fn an_org_never_passes_a_second_factor() {
        let dir = tempfile::tempdir().unwrap();
        let store = Store::open(&dir.path().join("forge.db")).await.unwrap();
        let key = [7u8; 32];
        let alice = store.create_user("alice", None, "h", false).await.unwrap();
        let org = store.create_org("acme", &alice.id).await.unwrap();
        let shared = forge_core::totp::new_secret();
        let sealed = forge_core::secret::seal(&key, &shared).unwrap();
        let code = forge_core::totp::code_at(&shared, forge_core::totp::step_for(unix_now()));
        let recovery = forge_core::totp::recovery_code_hash("aaaa-bbbb");
        for id in [&org.id, &alice.id] {
            store.set_totp_secret(id, &sealed).await.unwrap();
            store
                .enable_totp(id, 0, std::slice::from_ref(&recovery))
                .await
                .unwrap();
        }

        assert!(
            !verify_second_factor(&store, &key, &org.id, &code)
                .await
                .unwrap()
        );
        assert!(
            !verify_second_factor(&store, &key, &org.id, "aaaa-bbbb")
                .await
                .unwrap()
        );
        // The same setup on a person works, so the refusal is the org rule.
        assert!(
            verify_second_factor(&store, &key, &alice.id, "aaaa-bbbb")
                .await
                .unwrap()
        );
    }
}