//! Passkeys (WebAuthn): registration, sign-in, and the second factor after a password.
//!
//! Sign-in uses discoverable credentials only, so the user types no name. That keeps the
//! unauthenticated part free of anything that could tell whether a username exists.

use std::collections::HashMap;
use std::sync::Mutex;
use std::time::{Duration, Instant};

use api::{Challenge, ChallengeAnswer, LoginResult};
use axum::Json;
use axum::extract::{FromRequestParts, Path as UrlPath, State};
use axum::http::request::Parts;
use axum::http::{HeaderMap, Uri, header};
use axum::response::{IntoResponse, Response};
use rusqlite::{Connection, OptionalExtension, params};
use webauthn_rs::prelude::*;
use webauthn_rs_proto::ResidentKeyRequirement;

use crate::auth::{self, ClientIp, User};
use crate::{AppState, Error, now};

pub const PASSKEY_LIMIT: i64 = 10;
/// How long a browser has to answer a challenge.
const TTL: Duration = Duration::from_secs(300);
/// Anyone can start a passkey sign-in, so their pending challenges need a cap.
/// Past it the oldest goes, so a flood of starts cannot block everyone else's sign-in.
const MAX_ANONYMOUS: usize = 1000;
/// Passkey sign-in starts per address in 15 minutes, so one client cannot push out the others' challenges.
const MAX_STARTS: u32 = 30;

pub enum Pending {
    Register {
        user_id: i64,
        state: Box<PasskeyRegistration>,
    },
    SignIn(Box<DiscoverableAuthentication>),
    /// The password passed. The account also needs a passkey.
    SecondFactor {
        user_id: i64,
        state: Box<PasskeyAuthentication>,
    },
    /// A passkey passed. The account also needs its password.
    NeedsPassword {
        user_id: i64,
    },
}

/// WebAuthn takes two requests. This holds what the second one needs, keyed by a handle the client echoes.
#[derive(Default)]
pub struct Ceremonies(Mutex<HashMap<String, (Pending, Instant)>>);

impl Ceremonies {
    pub fn put(&self, pending: Pending) -> String {
        let mut map = self.0.lock().unwrap();
        map.retain(|_, (_, at)| at.elapsed() < TTL);
        let anonymous = |p: &Pending| matches!(p, Pending::SignIn(_));
        if anonymous(&pending) {
            let mut started: Vec<(String, Instant)> = map
                .iter()
                .filter(|(_, (p, _))| anonymous(p))
                .map(|(id, (_, at))| (id.clone(), *at))
                .collect();
            if started.len() >= MAX_ANONYMOUS {
                started.sort_by_key(|(_, at)| *at);
                for (id, _) in &started[..=started.len() - MAX_ANONYMOUS] {
                    map.remove(id);
                }
            }
        }
        let (id, _) = auth::new_secret();
        map.insert(id.clone(), (pending, Instant::now()));
        id
    }

    /// One handle answers one challenge.
    pub fn take(&self, id: &str) -> Option<Pending> {
        let mut map = self.0.lock().unwrap();
        map.retain(|_, (_, at)| at.elapsed() < TTL);
        map.remove(id).map(|(p, _)| p)
    }
}

pub fn expired() -> Error {
    Error::BadRequest("that took too long, please try again".into())
}

/// The details go to the log. To the user every failure is the same.
fn failed(e: WebauthnError) -> Error {
    eprintln!("webauthn ceremony failed: {e:?}");
    Error::BadRequest("that passkey could not be used".into())
}

fn challenge<T: serde::Serialize>(
    state: &AppState,
    pending: Pending,
    options: &T,
) -> Result<Challenge, Error> {
    let options = serde_json::to_string(options).map_err(|e| Error::Internal(e.to_string()))?;
    Ok(Challenge {
        state_id: state.ceremonies.put(pending),
        options,
    })
}

/// The relying party for the address the browser is on.
pub struct Rp(Webauthn);

impl FromRequestParts<AppState> for Rp {
    type Rejection = Error;

    async fn from_request_parts(parts: &mut Parts, state: &AppState) -> Result<Self, Error> {
        relying_party(state, &parts.uri, &parts.headers).map(Rp)
    }
}

pub fn relying_party(state: &AppState, uri: &Uri, headers: &HeaderMap) -> Result<Webauthn, Error> {
    let origin = match &state.public_url {
        Some(url) => url.clone(),
        None => {
            // HTTP/2 carries the host in the URI, HTTP/1.1 in the Host header.
            let host = match uri.authority() {
                Some(a) => a.as_str().to_owned(),
                None => headers
                    .get(header::HOST)
                    .and_then(|v| v.to_str().ok())
                    .ok_or_else(|| Error::BadRequest("no Host header".into()))?
                    .to_owned(),
            };
            Url::parse(&format!("http://{host}")).map_err(|e| Error::BadRequest(e.to_string()))?
        }
    };
    // The browser signs its own origin. A mismatch would only fail later, with no useful message.
    let expected = origin.origin().ascii_serialization();
    if let Some(browser) = headers.get(header::ORIGIN).and_then(|v| v.to_str().ok())
        && browser != expected
    {
        return Err(Error::BadRequest(format!(
            "passkeys are set up for {expected}, but this page is {browser}. Set --public-url to the address you use."
        )));
    }
    let rp_id = origin.domain().ok_or_else(|| {
        Error::BadRequest("passkeys need a domain name, not an IP address".into())
    })?;
    WebauthnBuilder::new(rp_id, &origin)
        .and_then(|b| b.rp_name("opentracker").build())
        .map_err(failed)
}

/// The stored passkeys of a user. An unreadable row is skipped, so it cannot lock the user out of the others.
pub fn load(db: &Connection, user_id: i64) -> Result<Vec<(i64, Passkey)>, Error> {
    let rows: Vec<(i64, String)> = db
        .prepare_cached("SELECT id, passkey FROM passkeys WHERE user_id = ?1")?
        .query_map([user_id], |r| Ok((r.get(0)?, r.get(1)?)))?
        .collect::<rusqlite::Result<_>>()?;
    Ok(rows
        .into_iter()
        .filter_map(|(id, json)| match serde_json::from_str(&json) {
            Ok(key) => Some((id, key)),
            Err(e) => {
                eprintln!("passkey {id} is unreadable: {e}");
                None
            }
        })
        .collect())
}

pub fn count(db: &Connection, user_id: i64) -> Result<i64, Error> {
    Ok(db.query_row(
        "SELECT COUNT(*) FROM passkeys WHERE user_id = ?1",
        [user_id],
        |r| r.get(0),
    )?)
}

/// Stores the new signature counter and the time of use.
fn record_use(db: &Connection, user_id: i64, result: &AuthenticationResult) -> Result<(), Error> {
    for (id, mut key) in load(db, user_id)? {
        if key.cred_id() == result.cred_id() {
            key.update_credential(result);
            let json = serde_json::to_string(&key).map_err(|e| Error::Internal(e.to_string()))?;
            db.execute(
                "UPDATE passkeys SET passkey = ?1, last_used_at = ?2 WHERE id = ?3",
                params![json, now(), id],
            )?;
        }
    }
    Ok(())
}

pub fn sign_in(state: &AppState, user_id: i64) -> Result<Response, Error> {
    let cookie = auth::create_session(state, user_id)?;
    Ok((
        [(header::SET_COOKIE, cookie)],
        Json(LoginResult {
            ok: true,
            ..Default::default()
        }),
    )
        .into_response())
}

/// The password passed. Asks for one of the user's passkeys next.
pub fn second_factor(
    state: &AppState,
    uri: &Uri,
    headers: &HeaderMap,
    user_id: i64,
) -> Result<Response, Error> {
    let rp = relying_party(state, uri, headers)?;
    let keys: Vec<Passkey> = load(&state.db(), user_id)?
        .into_iter()
        .map(|(_, k)| k)
        .collect();
    if keys.is_empty() {
        return Err(Error::Internal(format!(
            "user {user_id} needs a passkey but has none"
        )));
    }
    let (options, auth_state) = rp.start_passkey_authentication(&keys).map_err(failed)?;
    let ch = challenge(
        state,
        Pending::SecondFactor {
            user_id,
            state: Box::new(auth_state),
        },
        &options,
    )?;
    Ok(Json(LoginResult {
        ok: false,
        passkey_challenge: Some(ch),
        ..Default::default()
    })
    .into_response())
}

pub async fn login_begin(
    State(s): State<AppState>,
    ClientIp(ip): ClientIp,
    Rp(rp): Rp,
) -> Result<Json<Challenge>, Error> {
    let key = format!("passkey {}", auth::ip_group(ip));
    s.limiter.attempt(&[(&key, MAX_STARTS)])?;
    let (mut options, auth_state) = rp.start_discoverable_authentication().map_err(failed)?;
    // Without this the browser waits for the autofill dropdown instead of showing its dialog.
    options.mediation = None;
    Ok(Json(challenge(
        &s,
        Pending::SignIn(Box::new(auth_state)),
        &options,
    )?))
}

pub async fn login_finish(
    State(s): State<AppState>,
    Rp(rp): Rp,
    Json(b): Json<ChallengeAnswer>,
) -> Result<Response, Error> {
    let pending = s.ceremonies.take(&b.state_id).ok_or_else(expired)?;
    let cred: PublicKeyCredential = serde_json::from_str(&b.credential)
        .map_err(|_| Error::BadRequest("unreadable credential".into()))?;
    let db = s.db();
    let (user_id, password_done) = match pending {
        Pending::SignIn(auth_state) => {
            // The user handle is only a claim until the signature checks out against that user's keys.
            let (handle, _) = rp
                .identify_discoverable_authentication(&cred)
                .map_err(failed)?;
            let user_id: i64 = db
                .query_row(
                    "SELECT id FROM users WHERE webauthn_id = ?1",
                    [handle.to_string()],
                    |r| r.get(0),
                )
                .optional()?
                .ok_or_else(|| failed(WebauthnError::CredentialNotFound))?;
            let keys: Vec<DiscoverableKey> =
                load(&db, user_id)?.iter().map(|(_, k)| k.into()).collect();
            let result = rp
                .finish_discoverable_authentication(&cred, *auth_state, &keys)
                .map_err(failed)?;
            record_use(&db, user_id, &result)?;
            (user_id, false)
        }
        Pending::SecondFactor {
            user_id,
            state: auth_state,
        } => {
            let result = rp
                .finish_passkey_authentication(&cred, &auth_state)
                .map_err(failed)?;
            record_use(&db, user_id, &result)?;
            (user_id, true)
        }
        _ => return Err(expired()),
    };
    let two_factor: bool = db.query_row(
        "SELECT two_factor FROM users WHERE id = ?1",
        [user_id],
        |r| r.get(0),
    )?;
    drop(db);
    if two_factor && !password_done {
        let state_id = s.ceremonies.put(Pending::NeedsPassword { user_id });
        return Ok(Json(LoginResult {
            ok: false,
            password_required: Some(state_id),
            ..Default::default()
        })
        .into_response());
    }
    sign_in(&s, user_id)
}

pub async fn list(State(s): State<AppState>, user: User) -> Result<Json<Vec<api::Passkey>>, Error> {
    let keys = s
        .db()
        .prepare_cached("SELECT id, name, created_at, last_used_at FROM passkeys WHERE user_id = ?1 ORDER BY id")?
        .query_map([user.id], |r| {
            Ok(api::Passkey { id: r.get(0)?, name: r.get(1)?, created_at: r.get(2)?, last_used_at: r.get(3)? })
        })?
        .collect::<rusqlite::Result<_>>()?;
    Ok(Json(keys))
}

fn too_many() -> Error {
    Error::BadRequest(format!("you can have at most {PASSKEY_LIMIT} passkeys"))
}

pub async fn register_begin(
    State(s): State<AppState>,
    user: User,
    Rp(rp): Rp,
) -> Result<Json<Challenge>, Error> {
    user.check_recent()?;
    let db = s.db();
    if count(&db, user.id)? >= PASSKEY_LIMIT {
        return Err(too_many());
    }
    let handle: String = db.query_row(
        "SELECT webauthn_id FROM users WHERE id = ?1",
        [user.id],
        |r| r.get(0),
    )?;
    let handle = Uuid::parse_str(&handle).map_err(|e| Error::Internal(e.to_string()))?;
    // The authenticator then refuses a second credential for the same account.
    let existing: Vec<CredentialID> = load(&db, user.id)?
        .iter()
        .map(|(_, k)| k.cred_id().clone())
        .collect();
    drop(db);
    let (mut options, reg) = rp
        .start_passkey_registration(handle, &user.username, &user.username, Some(existing))
        .map_err(failed)?;
    // webauthn-rs asks for a non-discoverable credential, but sign-in without a username needs a discoverable one.
    if let Some(sel) = options.public_key.authenticator_selection.as_mut() {
        sel.resident_key = Some(ResidentKeyRequirement::Required);
    }
    Ok(Json(challenge(
        &s,
        Pending::Register {
            user_id: user.id,
            state: Box::new(reg),
        },
        &options,
    )?))
}

pub async fn register_finish(
    State(s): State<AppState>,
    user: User,
    Rp(rp): Rp,
    Json(b): Json<ChallengeAnswer>,
) -> Result<Json<api::Passkey>, Error> {
    let Some(Pending::Register {
        user_id,
        state: reg,
    }) = s.ceremonies.take(&b.state_id)
    else {
        return Err(expired());
    };
    if user_id != user.id {
        return Err(expired());
    }
    let cred: RegisterPublicKeyCredential = serde_json::from_str(&b.credential)
        .map_err(|_| Error::BadRequest("unreadable credential".into()))?;
    let key = rp
        .finish_passkey_registration(&cred, &reg)
        .map_err(failed)?;
    let json = serde_json::to_string(&key).map_err(|e| Error::Internal(e.to_string()))?;
    let name = match b.name.trim() {
        "" => "Passkey",
        n => n,
    };
    let name: String = name.chars().take(100).collect();

    let db = s.db();
    if count(&db, user.id)? >= PASSKEY_LIMIT {
        return Err(too_many());
    }
    let created_at = now();
    db.execute(
        "INSERT INTO passkeys (user_id, cred_id, passkey, name, created_at) VALUES (?1, ?2, ?3, ?4, ?5)",
        params![user.id, key.cred_id().as_ref(), json, name, created_at],
    )
    .map_err(crate::taken("that passkey is already registered"))?;
    let id = db.last_insert_rowid();
    drop(db);
    auth::end_other_sessions(&s, &user)?;
    Ok(Json(api::Passkey {
        id,
        name,
        created_at,
        last_used_at: None,
    }))
}

pub async fn delete(
    State(s): State<AppState>,
    user: User,
    UrlPath(id): UrlPath<i64>,
) -> Result<Json<()>, Error> {
    user.check_recent()?;
    let db = s.db();
    let (has_password, two_factor): (bool, bool) = db.query_row(
        "SELECT pw_hash IS NOT NULL, two_factor FROM users WHERE id = ?1",
        [user.id],
        |r| Ok((r.get(0)?, r.get(1)?)),
    )?;
    if count(&db, user.id)? == 1 {
        if two_factor {
            return Err(Error::BadRequest(
                "turn off two-factor sign-in before removing your last passkey".into(),
            ));
        }
        if !has_password {
            return Err(Error::BadRequest(
                "set a password before removing your last passkey".into(),
            ));
        }
    }
    if db.execute(
        "DELETE FROM passkeys WHERE id = ?1 AND user_id = ?2",
        [id, user.id],
    )? == 0
    {
        return Err(Error::NotFound);
    }
    drop(db);
    auth::end_other_sessions(&s, &user)?;
    Ok(Json(()))
}