//! Rate limiting (SPEC §8.6, AC-2.12, §8.5), as a tower layer that `router.rs` attaches to every route.
//!
//! * Unauthenticated endpoints (E-1 to E-6 and unknown paths): at most 60 requests per minute per client IP.
//! * Authenticated endpoints: at most 600 requests per minute per device, keyed by the hash of the bearer token
//!   (the token identifies the device; the limiter never touches the database). A request to an authenticated
//!   endpoint without a well-formed bearer token counts against its IP like an unauthenticated one. A response 401
//!   also counts against the IP, and once the IP is over its limit, tokens that have never been accepted are refused
//!   too, so random tokens do not bypass the IP limit.
//! * Invite accepts: after 10 failed accepts (any 4xx other than 429) from one IP within 60 minutes, every accept
//!   from that IP is `rate_limited` until the 60 minutes since the first of those failures have passed.
//!
//! Refusals are 429 `rate_limited` with `Retry-After` in whole seconds. Windows are fixed: they start at the first
//! request (or failure) counted and last one minute (one hour for accept failures). Everything is in memory and resets
//! on restart; there is no table.
//!
//! The client IP is the TCP peer (axum `ConnectInfo<SocketAddr>`), or the last `X-Forwarded-For` entry when the peer
//! is an address of `KVITT_TRUSTED_PROXY` (Caddy, SPEC §8.5). The host name is resolved lazily and re-resolved at most
//! once per [`PROXY_RESOLVE_INTERVAL`]; a name that does not resolve means no peer is trusted (D-022). Without
//! `ConnectInfo` (a server not started with `into_make_service_with_connect_info`) all clients share one IP bucket.

use std::collections::HashMap;
use std::convert::Infallible;
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::sync::{Arc, Mutex, MutexGuard};
use std::task::{Context, Poll};
use std::time::Instant;

use axum::extract::{ConnectInfo, MatchedPath, Request};
use axum::http::StatusCode;
use axum::response::{IntoResponse, Response};
use chrono::{DateTime, Duration, Utc};
use tower::{Layer, Service};

use crate::AppState;
use crate::audit::BoxFuture;
use crate::auth::{is_token_shaped, sha256_hex};
use crate::error::AppError;
use crate::time::Clock;

/// Unauthenticated requests per window per client IP (SPEC §8.6).
pub const UNAUTHENTICATED_LIMIT: u32 = 60;
/// Authenticated requests per window per device (SPEC §8.6).
pub const AUTHENTICATED_LIMIT: u32 = 600;
/// The request window (SPEC §8.6: per minute).
pub const REQUEST_WINDOW: Duration = Duration::minutes(1);
/// Failed invite accepts tolerated per IP within [`ACCEPT_FAILURE_WINDOW`] (AC-2.12).
pub const ACCEPT_FAILURE_LIMIT: u32 = 10;
/// The invite-accept failure window (AC-2.12).
pub const ACCEPT_FAILURE_WINDOW: Duration = Duration::minutes(60);
/// How long a resolution of `KVITT_TRUSTED_PROXY` is reused.
pub const PROXY_RESOLVE_INTERVAL: std::time::Duration = std::time::Duration::from_secs(60);
/// How often expired windows are pruned from memory.
const PRUNE_INTERVAL: Duration = Duration::minutes(10);

/// Route templates of the unauthenticated endpoints (E-1 to E-6).
const PUBLIC_ROUTES: [&str; 6] = [
    "/health",
    "/setup/status",
    "/invites/accept",
    "/i/{code}",
    "/apk/latest",
    "/apk/files/{file_name}",
];
const ACCEPT_ROUTE: &str = "/invites/accept";

/// The layer for `router.rs`. One limiter per router: build the router once and reuse it (as `serve` does).
pub fn layer(state: &AppState) -> RateLimitLayer {
    RateLimitLayer {
        limiter: Arc::new(Limiter::new(
            state.clock.clone(),
            state.config.trusted_proxy.clone(),
        )),
    }
}

/// A fixed counting window.
#[derive(Debug, Clone, Copy)]
struct Window {
    start: DateTime<Utc>,
    count: u32,
}

impl Window {
    fn expired(&self, now: DateTime<Utc>, len: Duration) -> bool {
        now >= self.start + len
    }

    fn retry_after(&self, now: DateTime<Utc>, len: Duration) -> u64 {
        let remaining = (self.start + len - now).num_milliseconds().max(0);
        u64::try_from((remaining + 999) / 1000).unwrap_or(1).max(1)
    }
}

/// Counts one request in `slot` unless the window is full; `Err(retry_after_secs)` when it is.
fn hit(
    slot: &mut Option<Window>,
    now: DateTime<Utc>,
    len: Duration,
    limit: u32,
) -> Result<(), u64> {
    let window = match slot {
        Some(w) if !w.expired(now, len) => w,
        _ => slot.insert(Window {
            start: now,
            count: 0,
        }),
    };
    if window.count >= limit {
        return Err(window.retry_after(now, len));
    }
    window.count += 1;
    Ok(())
}

/// Whether `slot` is full (without counting).
fn full(slot: Option<&Window>, now: DateTime<Utc>, len: Duration, limit: u32) -> Option<u64> {
    slot.filter(|w| !w.expired(now, len) && w.count >= limit)
        .map(|w| w.retry_after(now, len))
}

#[derive(Debug, Default)]
struct TokenState {
    window: Option<Window>,
    /// The token has had a response other than 401: it belongs to a device.
    accepted: bool,
}

#[derive(Debug, Default)]
struct Buckets {
    ip: HashMap<IpAddr, Window>,
    tokens: HashMap<String, TokenState>,
    accept_failures: HashMap<IpAddr, Window>,
    last_prune: Option<DateTime<Utc>>,
}

impl Buckets {
    fn prune(&mut self, now: DateTime<Utc>) {
        if self.last_prune.is_some_and(|t| now - t < PRUNE_INTERVAL) {
            return;
        }
        self.last_prune = Some(now);
        self.ip.retain(|_, w| !w.expired(now, REQUEST_WINDOW));
        self.tokens
            .retain(|_, t| t.accepted || t.window.is_some_and(|w| !w.expired(now, REQUEST_WINDOW)));
        self.accept_failures
            .retain(|_, w| !w.expired(now, ACCEPT_FAILURE_WINDOW));
    }
}

/// How a request is counted.
enum Key {
    Ip(IpAddr),
    Token { ip: IpAddr, hash: String },
}

/// The shared limiter state.
pub struct Limiter {
    clock: Arc<dyn Clock>,
    trusted_proxy: String,
    proxy: Mutex<Option<(Instant, Vec<IpAddr>)>>,
    buckets: Mutex<Buckets>,
}

fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
    match mutex.lock() {
        Ok(guard) => guard,
        Err(poisoned) => poisoned.into_inner(),
    }
}

impl Limiter {
    pub fn new(clock: Arc<dyn Clock>, trusted_proxy: String) -> Self {
        Self {
            clock,
            trusted_proxy,
            proxy: Mutex::new(None),
            buckets: Mutex::new(Buckets::default()),
        }
    }

    /// The addresses of `KVITT_TRUSTED_PROXY`, resolved at most once per [`PROXY_RESOLVE_INTERVAL`].
    async fn proxy_addrs(&self) -> Vec<IpAddr> {
        if let Some((at, addrs)) = lock(&self.proxy).as_ref()
            && at.elapsed() < PROXY_RESOLVE_INTERVAL
        {
            return addrs.clone();
        }
        let name = self.trusted_proxy.trim();
        let addrs: Vec<IpAddr> = if name.is_empty() {
            Vec::new()
        } else if let Ok(ip) = name.parse::<IpAddr>() {
            vec![ip]
        } else {
            match tokio::net::lookup_host((name, 0)).await {
                Ok(found) => found.map(|a| a.ip()).collect(),
                Err(_) => Vec::new(),
            }
        };
        *lock(&self.proxy) = Some((Instant::now(), addrs.clone()));
        addrs
    }

    /// The TCP peer and the last `X-Forwarded-For` entry of `request`.
    fn peer_and_forwarded(request: &Request) -> (Option<IpAddr>, Option<IpAddr>) {
        let peer = request
            .extensions()
            .get::<ConnectInfo<SocketAddr>>()
            .map(|c| c.0.ip());
        let forwarded = request
            .headers()
            .get("x-forwarded-for")
            .and_then(|v| v.to_str().ok())
            .and_then(|v| v.rsplit(',').next())
            .and_then(|v| v.trim().parse::<IpAddr>().ok());
        (peer, forwarded)
    }

    /// The client IP (SPEC §8.5): the forwarded address when the peer is the trusted proxy, else the peer.
    async fn client_ip(&self, peer: Option<IpAddr>, forwarded: Option<IpAddr>) -> IpAddr {
        let Some(peer) = peer else {
            return IpAddr::V4(Ipv4Addr::UNSPECIFIED);
        };
        match forwarded {
            Some(client) if self.proxy_addrs().await.contains(&peer) => client,
            _ => peer,
        }
    }

    fn key(request: &Request, route: &str, ip: IpAddr) -> Key {
        if PUBLIC_ROUTES.contains(&route) || route.is_empty() {
            return Key::Ip(ip);
        }
        let token = request
            .headers()
            .get(axum::http::header::AUTHORIZATION)
            .and_then(|v| v.to_str().ok())
            .and_then(|v| v.split_once(' '))
            .filter(|(scheme, _)| scheme.eq_ignore_ascii_case("bearer"))
            .map(|(_, token)| token.trim())
            .filter(|token| is_token_shaped(token));
        match token {
            Some(token) => Key::Token {
                ip,
                hash: sha256_hex(token),
            },
            None => Key::Ip(ip),
        }
    }

    /// Counts the request before it runs; `Err(retry_after_secs)` refuses it.
    fn admit(&self, key: &Key, accept: bool, now: DateTime<Utc>) -> Result<(), u64> {
        let mut buckets = lock(&self.buckets);
        buckets.prune(now);
        match key {
            Key::Ip(ip) => {
                if accept
                    && let Some(retry) = full(
                        buckets.accept_failures.get(ip),
                        now,
                        ACCEPT_FAILURE_WINDOW,
                        ACCEPT_FAILURE_LIMIT,
                    )
                {
                    return Err(retry);
                }
                let mut slot = buckets.ip.get(ip).copied();
                let result = hit(&mut slot, now, REQUEST_WINDOW, UNAUTHENTICATED_LIMIT);
                if let Some(w) = slot {
                    buckets.ip.insert(*ip, w);
                }
                result
            }
            Key::Token { ip, hash } => {
                let accepted = buckets.tokens.get(hash).is_some_and(|t| t.accepted);
                if !accepted
                    && let Some(retry) = full(
                        buckets.ip.get(ip),
                        now,
                        REQUEST_WINDOW,
                        UNAUTHENTICATED_LIMIT,
                    )
                {
                    return Err(retry);
                }
                let state = buckets.tokens.entry(hash.clone()).or_default();
                hit(&mut state.window, now, REQUEST_WINDOW, AUTHENTICATED_LIMIT)
            }
        }
    }

    /// Records what the response says about the request (401s, accept failures, accepted tokens).
    fn record(&self, key: &Key, accept: bool, status: StatusCode, now: DateTime<Utc>) {
        let mut buckets = lock(&self.buckets);
        match key {
            Key::Ip(ip) => {
                if accept && status.is_client_error() && status != StatusCode::TOO_MANY_REQUESTS {
                    let mut slot = buckets.accept_failures.get(ip).copied();
                    // The window is refused once full (admit), so this count never exceeds the limit.
                    let _ = hit(&mut slot, now, ACCEPT_FAILURE_WINDOW, ACCEPT_FAILURE_LIMIT);
                    if let Some(w) = slot {
                        buckets.accept_failures.insert(*ip, w);
                    }
                }
            }
            Key::Token { ip, hash } => {
                if status == StatusCode::UNAUTHORIZED {
                    let mut slot = buckets.ip.get(ip).copied();
                    let _ = hit(&mut slot, now, REQUEST_WINDOW, UNAUTHENTICATED_LIMIT);
                    if let Some(w) = slot {
                        buckets.ip.insert(*ip, w);
                    }
                } else if status != StatusCode::TOO_MANY_REQUESTS {
                    buckets.tokens.entry(hash.clone()).or_default().accepted = true;
                }
            }
        }
    }

    async fn handle<S>(self: Arc<Self>, request: Request, mut inner: S) -> Response
    where
        S: Service<Request, Response = Response, Error = Infallible>,
    {
        let route = request
            .extensions()
            .get::<MatchedPath>()
            .map(|p| p.as_str().to_owned())
            .unwrap_or_default();
        let accept = route == ACCEPT_ROUTE;
        let (peer, forwarded) = Self::peer_and_forwarded(&request);
        let ip = self.client_ip(peer, forwarded).await;
        let key = Self::key(&request, &route, ip);
        if let Err(retry_after) = self.admit(&key, accept, self.clock.now()) {
            return AppError::rate_limited(retry_after).into_response();
        }
        let response = match inner.call(request).await {
            Ok(response) => response,
            Err(never) => match never {},
        };
        self.record(&key, accept, response.status(), self.clock.now());
        response
    }
}

/// The tower layer of [`layer`].
#[derive(Clone)]
pub struct RateLimitLayer {
    limiter: Arc<Limiter>,
}

impl<S> Layer<S> for RateLimitLayer {
    type Service = RateLimit<S>;

    fn layer(&self, inner: S) -> Self::Service {
        RateLimit {
            inner,
            limiter: self.limiter.clone(),
        }
    }
}

/// The rate-limiting service wrapping one route.
#[derive(Clone)]
pub struct RateLimit<S> {
    inner: S,
    limiter: Arc<Limiter>,
}

impl<S> Service<Request> for RateLimit<S>
where
    S: Service<Request, Response = Response, Error = Infallible> + Clone + Send + 'static,
    S::Future: Send + 'static,
{
    type Response = Response;
    type Error = Infallible;
    type Future = BoxFuture<'static, Result<Response, Infallible>>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, request: Request) -> Self::Future {
        // The ready clone goes into the future; a fresh clone stays for the next call (tower's usual pattern).
        let clone = self.inner.clone();
        let inner = std::mem::replace(&mut self.inner, clone);
        let limiter = self.limiter.clone();
        Box::pin(async move { Ok(limiter.handle(request, inner).await) })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::time::{FixedClock, parse_ts};

    #[test]
    fn fixed_window_counts_and_reports_retry_after() {
        let now = parse_ts("2026-10-17T12:00:00.000Z").unwrap();
        let mut slot = None;
        for _ in 0..3 {
            assert!(hit(&mut slot, now, REQUEST_WINDOW, 3).is_ok());
        }
        assert_eq!(
            hit(&mut slot, now + Duration::seconds(20), REQUEST_WINDOW, 3),
            Err(40)
        );
        assert!(hit(&mut slot, now + Duration::seconds(60), REQUEST_WINDOW, 3).is_ok());
    }

    #[test]
    fn retry_after_rounds_up_and_is_at_least_one() {
        let now = parse_ts("2026-10-17T12:00:00.000Z").unwrap();
        let w = Window {
            start: now,
            count: 1,
        };
        assert_eq!(
            w.retry_after(now + Duration::milliseconds(59_500), REQUEST_WINDOW),
            1
        );
        assert_eq!(
            w.retry_after(now + Duration::milliseconds(100), REQUEST_WINDOW),
            60
        );
    }

    #[test]
    fn pruning_drops_expired_windows() {
        let now = parse_ts("2026-10-17T12:00:00.000Z").unwrap();
        let limiter = Limiter::new(Arc::new(FixedClock::new(now)), String::new());
        let ip = IpAddr::V4(Ipv4Addr::LOCALHOST);
        assert!(limiter.admit(&Key::Ip(ip), false, now).is_ok());
        assert_eq!(lock(&limiter.buckets).ip.len(), 1);
        assert!(
            limiter
                .admit(
                    &Key::Ip(IpAddr::V4(Ipv4Addr::BROADCAST)),
                    false,
                    now + Duration::minutes(11)
                )
                .is_ok()
        );
        assert_eq!(lock(&limiter.buckets).ip.len(), 1);
    }
}