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

/// The users of an htpasswd file (`user:hash` lines; empty lines and `#` comments are
/// skipped). A line that is neither, or a file without users, is an error: every request
/// would be refused, which is better told at startup.
fn load_users(path: &str) -> anyhow::Result<HashMap<String, String>> {
    let file = File::open(path).with_context(|| format!("Opening the htpasswd file {path}"))?;
    let reader = BufReader::new(file);

    let mut users = HashMap::new();

    for (number, line) in reader.lines().enumerate() {
        let line = line.with_context(|| format!("Reading the htpasswd file {path}"))?;
        let line = line.trim();
        if line.is_empty() || line.starts_with('#') {
            continue;
        }
        let Some((user, hash)) = line.split_once(':') else {
            anyhow::bail!("{path}, line {}: not `user:hash`", number + 1);
        };
        users.insert(user.to_owned(), hash.to_owned());
    }
    if users.is_empty() {
        anyhow::bail!("The htpasswd file {path} has no users");
    }

    Ok(users)
}

/// Who may use the HTTP API: the users of an htpasswd file (basic authentication),
/// and the user a reverse proxy in front names (`--user-header`), if it is one of
/// `--trusted-proxies`.
#[derive(Clone)]
pub(crate) struct AuthState {
    idle_timeout: Duration,

    /// None without an htpasswd file: then only the proxy's users.
    users: Option<Arc<HashMap<String, String>>>, // user -> PHC hash string
    proxy: Option<Arc<ProxyAuth>>,
    cache: Arc<DashMap<[u8; 32], Instant>>,
    /// Bounds concurrent Argon2 verifications (CPU and ~19 MiB each).
    verify_slots: Arc<Semaphore>,
}

/// A reverse proxy authenticates, and names the user in `header`.
struct ProxyAuth {
    header: HeaderName,
    trusted: Vec<Range>,
}

/// An address, or the addresses sharing its first `bits` bits.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct Range {
    address: IpAddr,
    bits: u32,
}

impl Range {
    fn parse(text: &str) -> anyhow::Result<Self> {
        let (address, bits) = match text.trim().split_once('/') {
            Some((address, bits)) => (address, Some(bits)),
            None => (text.trim(), None),
        };
        let address: IpAddr = address
            .parse()
            .with_context(|| format!("trusted-proxies: {text:?} is no IP address"))?;
        let all = if address.is_ipv4() { 32 } else { 128 };
        let bits = match bits {
            None => all,
            Some(bits) => bits
                .parse()
                .ok()
                .filter(|bits| *bits <= all)
                .with_context(|| format!("trusted-proxies: {text:?}: not 0 to {all} bits"))?,
        };
        Ok(Self { address, bits })
    }

    fn contains(&self, address: IpAddr) -> bool {
        // An IPv4 address as IPv6 (`::ffff:a.b.c.d`) is that IPv4 address.
        let address = match address {
            IpAddr::V6(v6) => v6.to_ipv4_mapped().map_or(address, IpAddr::V4),
            v4 => v4,
        };
        let leading =
            |a: u128, b: u128, width: u32| self.bits == 0 || (a ^ b) >> (width - self.bits) == 0;
        match (self.address, address) {
            (IpAddr::V4(range), IpAddr::V4(address)) => {
                leading(u32::from(range).into(), u32::from(address).into(), 32)
            }
            (IpAddr::V6(range), IpAddr::V6(address)) => {
                leading(u128::from(range), u128::from(address), 128)
            }
            _ => false,
        }
    }
}

impl AuthState {
    /// From the options; None if there is neither an htpasswd file nor a user header
    /// (no users: no HTTP API, for a relay).
    pub(crate) fn from_settings(
        settings: &crate::config::Settings,
    ) -> anyhow::Result<Option<Self>> {
        let users = settings
            .htpasswd
            .as_deref()
            .map(load_users)
            .transpose()?
            .map(Arc::new);
        let proxy = settings
            .user_header
            .as_deref()
            .map(|header| {
                let header = HeaderName::try_from(header.trim())
                    .with_context(|| format!("user-header: {header:?} is no header name"))?;
                let trusted = settings
                    .trusted_proxies
                    .iter()
                    .map(|range| Range::parse(range))
                    .collect::<anyhow::Result<Vec<_>>>()?;
                Ok::<_, anyhow::Error>(Arc::new(ProxyAuth { header, trusted }))
            })
            .transpose()?;
        if users.is_none() && proxy.is_none() {
            return Ok(None);
        }
        Ok(Some(Self {
            idle_timeout: Duration::from_secs(300),
            cache: Arc::new(DashMap::new()),
            users,
            proxy,
            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()))
}

/// The user a request authenticated as (a request extension, set by [`basic_auth`]).
#[derive(Clone, Debug)]
pub(crate) struct AuthenticatedUser(pub(crate) String);

/// The user the trusted proxy the request came from names, if it does.
fn proxied_user(state: &AuthState, headers: &HeaderMap, from: Option<IpAddr>) -> Option<String> {
    let proxy = state.proxy.as_ref()?;
    let user = headers.get(&proxy.header)?.to_str().ok()?.trim();
    if user.is_empty() {
        return None;
    }
    let trusted = from.is_some_and(|from| proxy.trusted.iter().any(|range| range.contains(from)));
    if !trusted {
        tracing::debug!(from = ?from, header = %proxy.header, "User header from an address not in trusted-proxies; ignored");
        return None;
    }
    Some(user.to_owned())
}

/// The user, if the credentials are valid.
async fn authenticate(state: &AuthState, headers: &HeaderMap) -> Option<String> {
    let (user, pass) = parse_basic(headers)?;
    let hash = state.users.as_ref()?.get(&user).cloned()?;
    verify(state, &user, pass, hash).await.then_some(user)
}

async fn verify(state: &AuthState, user: &str, pass: String, hash: String) -> bool {
    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
}

/// The user of a request: as the trusted proxy it came from names it, else by basic
/// authentication; `401` if neither.
pub(crate) async fn basic_auth(
    State(state): State<AuthState>,
    headers: HeaderMap,
    mut request: Request<Body>,
    next: axum::middleware::Next,
) -> Response {
    let from = request
        .extensions()
        .get::<ConnectInfo<SocketAddr>>()
        .map(|ConnectInfo(address)| address.ip());
    let user = match proxied_user(&state, &headers, from) {
        Some(user) => Some(user),
        None => authenticate(&state, &headers).await,
    };
    let Some(user) = user else {
        sleep(Duration::from_millis(1500)).await;
        // A challenge only where basic authentication can answer it.
        return match state.users {
            Some(_) => (
                StatusCode::UNAUTHORIZED,
                [(WWW_AUTHENTICATE, r#"Basic realm="Restricted""#)],
                "Unauthorized",
            )
                .into_response(),
            None => (StatusCode::UNAUTHORIZED, "Unauthorized").into_response(),
        };
    };

    request.extensions_mut().insert(AuthenticatedUser(user));
    next.run(request).await
}

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

    #[test]
    fn ranges_hold_their_addresses() {
        let range = |text: &str| Range::parse(text).unwrap();
        let ip = |text: &str| text.parse().unwrap();
        assert!(range("127.0.0.1").contains(ip("127.0.0.1")));
        assert!(!range("127.0.0.1").contains(ip("127.0.0.2")));
        assert!(range("10.0.0.0/8").contains(ip("10.1.2.3")));
        assert!(!range("10.0.0.0/8").contains(ip("11.0.0.1")));
        assert!(range("::1").contains(ip("::1")));
        assert!(range("fd00::/8").contains(ip("fd12::1")));
        assert!(range("0.0.0.0/0").contains(ip("192.0.2.1")));
        // IPv4 as IPv6, as a dual-stack listener sees it.
        assert!(range("127.0.0.1").contains(ip("::ffff:127.0.0.1")));
        assert!(!range("127.0.0.1").contains(ip("::1")));
        for bad in ["localhost", "10.0.0.0/33", "::/129", "10.0.0.0/x"] {
            assert!(Range::parse(bad).is_err(), "{bad}");
        }
    }
}