use argon2::{password_hash::PasswordVerifier, Argon2};
use axum::extract::State;
use axum::{
    body::Body,
    http::{
        header::{AUTHORIZATION, WWW_AUTHENTICATE},
        HeaderMap, Request, StatusCode,
    },
    response::{IntoResponse, Response},
};
use base64::{engine::general_purpose::STANDARD as BASE64, Engine};
use dashmap::DashMap;
use password_hash::phc::PasswordHash;
use std::{
    collections::HashMap,
    fs::File,
    io::{BufRead, BufReader},
    sync::Arc,
    time::{Duration, Instant},
};
use tokio::sync::Semaphore;
use tokio::time::sleep;

fn load_users(path: &str) -> std::io::Result<HashMap<String, String>> {
    let file = File::open(path)?;
    let reader = BufReader::new(file);

    let mut users = HashMap::new();

    for line in reader.lines() {
        let line = line?;

        if let Some((user, hash)) = line.split_once(':') {
            users.insert(user.to_owned(), hash.to_owned());
        }
    }

    Ok(users)
}

#[derive(Clone)]
pub struct AuthState {
    idle_timeout: Duration,

    users: Arc<HashMap<String, String>>, // user -> PHC hash string
    cache: Arc<DashMap<[u8; 32], Instant>>,
    /// Bounds concurrent Argon2 verifications (CPU and ~19 MiB each).
    verify_slots: Arc<Semaphore>,
}

impl AuthState {
    pub fn new(path: &str) -> Self {
        let users = load_users(path).unwrap_or_default();

        Self {
            idle_timeout: Duration::from_secs(300),
            cache: Arc::new(DashMap::new()),
            users: Arc::new(users),
            verify_slots: Arc::new(Semaphore::new(2)),
        }
    }
}

fn parse_basic(headers: &HeaderMap) -> Option<(String, String)> {
    let encoded = headers
        .get(AUTHORIZATION)?
        .to_str()
        .ok()?
        .strip_prefix("Basic ")?;
    let decoded = String::from_utf8(BASE64.decode(encoded).ok()?).ok()?;
    let (user, pass) = decoded.split_once(':')?;
    Some((user.to_owned(), pass.to_owned()))
}

async fn authenticate(state: &AuthState, headers: &HeaderMap) -> bool {
    let Some((user, pass)) = parse_basic(headers) else {
        return false;
    };
    let Some(hash) = state.users.get(&user).cloned() else {
        return false;
    };

    let key = *blake3::hash(format!("{user}\0{pass}\0{hash}").as_bytes()).as_bytes();

    // Fast path. The shard guard lives only until the end of this `match`;
    // never call into the map while holding it.
    let fresh = match state.cache.get_mut(&key) {
        Some(mut last_seen) if last_seen.elapsed() < state.idle_timeout => {
            *last_seen = Instant::now(); // slide
            true
        }
        _ => false,
    };
    if fresh {
        return true;
    }

    // Slow path: Argon2 is deliberately expensive, so keep it off the runtime
    // threads and bound how many run at once.
    let Ok(_permit) = state.verify_slots.acquire().await else {
        return false;
    };

    let verified = tokio::task::spawn_blocking(move || {
        PasswordHash::new(&hash)
            .map(|ph| {
                Argon2::default()
                    .verify_password(pass.as_bytes(), &ph)
                    .is_ok()
            })
            .unwrap_or(false)
    })
    .await
    .unwrap_or(false);

    if verified {
        // Overwrites an expired entry; no separate remove needed.
        state.cache.insert(key, Instant::now());
    }

    verified
}

pub async fn basic_auth(
    State(state): State<AuthState>,
    headers: HeaderMap,
    request: Request<Body>,
    next: axum::middleware::Next,
) -> Response {
    if !authenticate(&state, &headers).await {
        sleep(Duration::from_millis(1500)).await;
        return (
            StatusCode::UNAUTHORIZED,
            [(WWW_AUTHENTICATE, r#"Basic realm="Restricted""#)],
            "Unauthorized",
        )
            .into_response();
    }

    next.run(request).await
}