use std::net::{IpAddr, SocketAddr};
use std::sync::Arc;
use std::time::Duration;

use axum::Router;
use axum::extract::ConnectInfo;
use axum::http::Request;
use jiff::Timestamp;
use tokio::net::TcpListener;
use tokio::signal::unix::{SignalKind, signal};
use tower_governor::GovernorLayer;
use tower_governor::errors::GovernorError;
use tower_governor::governor::GovernorConfigBuilder;
use tower_governor::key_extractor::KeyExtractor;
use tower_http::limit::RequestBodyLimitLayer;
use tower_http::timeout::TimeoutLayer;

use crate::atproto::{AtprotoLogin, JacquardLogin, LoginError};
use crate::authstore::SqliteAuthStore;
use crate::clients::ClientStore;
use crate::config::{RateLimit, ServerConfig};
use crate::database::{self, Database, DatabaseError};
use crate::state::StateKeys;
use crate::web::{AppState, AppStateError, router};

pub const SWEEP_INTERVAL: Duration = Duration::from_secs(60);
pub const REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
pub const BODY_LIMIT_BYTES: usize = 64 * 1024;

#[derive(Debug, thiserror::Error)]
pub enum ServeError {
    #[error(transparent)]
    Login(#[from] LoginError),
    #[error("cannot listen on {listen}: {source}")]
    Bind {
        listen: String,
        source: std::io::Error,
    },
    #[error("server failed: {0}")]
    Io(#[from] std::io::Error),
    #[error(transparent)]
    Database(#[from] DatabaseError),
    #[error(transparent)]
    State(#[from] AppStateError),
    #[error("invalid rate limit configuration")]
    RateLimit,
}

pub fn init_tracing() {
    let filter = tracing_subscriber::EnvFilter::try_from_default_env()
        .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info"));
    let _ = tracing_subscriber::fmt()
        .with_env_filter(filter)
        .with_writer(std::io::stderr)
        .try_init();
}

pub async fn serve(config: ServerConfig, keys: StateKeys) -> Result<(), ServeError> {
    let db = Database::open(&config.state_dir.join(database::FILE_NAME))?;
    let login = JacquardLogin::new(
        &config,
        SqliteAuthStore::new(db.clone()),
        keys.client.keyset(),
    )?;
    let shutdown = shutdown_signal()?;
    let listener = TcpListener::bind(&config.listen)
        .await
        .map_err(|source| ServeError::Bind {
            listen: config.listen.clone(),
            source,
        })?;
    run(listener, config, keys, login, db, shutdown).await
}

pub async fn run<L: AtprotoLogin + 'static>(
    listener: TcpListener,
    config: ServerConfig,
    keys: StateKeys,
    login: L,
    db: Database,
    shutdown: impl Future<Output = ()> + Send + 'static,
) -> Result<(), ServeError> {
    warn_about_plaintext_redirects(&config, &db);
    let listen = listener.local_addr()?;
    let issuer = config.issuer().to_owned();
    let trust_proxy_headers = config.trust_proxy_headers;
    let rate_limit = config.rate_limit;
    let state = Arc::new(AppState::new(
        config,
        keys.signing,
        &keys.client,
        login,
        db,
    )?);

    let sweeper = Arc::clone(&state);
    tokio::spawn(async move {
        let mut interval = tokio::time::interval(SWEEP_INTERVAL);
        loop {
            interval.tick().await;
            if let Err(error) = sweeper.flows.sweep(Timestamp::now()) {
                tracing::error!(%error, "sweep failed");
            }
        }
    });

    let app = limited(router(state), trust_proxy_headers, rate_limit)?;
    tracing::info!(
        version = env!("CARGO_PKG_VERSION"),
        %listen,
        %issuer,
        %trust_proxy_headers,
        "listening"
    );
    if !trust_proxy_headers && listen.ip().is_loopback() {
        tracing::warn!(
            "listening on loopback without trust_proxy_headers, every request shares the proxy's one rate-limit key"
        );
    }
    axum::serve(
        listener,
        app.into_make_service_with_connect_info::<SocketAddr>(),
    )
    .with_graceful_shutdown(shutdown)
    .await?;
    tracing::info!("stopped");
    Ok(())
}

fn warn_about_plaintext_redirects(config: &ServerConfig, db: &Database) {
    let stored = ClientStore::new(db.clone())
        .active()
        .unwrap_or_else(|error| {
            tracing::error!(%error, "cannot list stored clients");
            Vec::new()
        });
    for client in config.clients.clients().chain(&stored) {
        let uris = client.plaintext_redirect_uris();
        if !uris.is_empty() {
            tracing::warn!(
                client = client.id().as_str(),
                redirect_uris = uris.join(" "),
                "client accepts redirect URIs that are not HTTPS"
            );
        }
    }
}

fn limited(
    app: Router,
    trust_proxy_headers: bool,
    rate_limit: RateLimit,
) -> Result<Router, ServeError> {
    let app = app
        .layer(TimeoutLayer::with_status_code(
            http::StatusCode::REQUEST_TIMEOUT,
            REQUEST_TIMEOUT,
        ))
        .layer(RequestBodyLimitLayer::new(BODY_LIMIT_BYTES));
    let config = GovernorConfigBuilder::default()
        .per_second(rate_limit.period_seconds)
        .burst_size(rate_limit.burst)
        .key_extractor(RateLimitKey {
            trust_proxy_headers,
        })
        .finish()
        .ok_or(ServeError::RateLimit)?;
    let limiter = config.limiter().clone();
    tokio::spawn(async move {
        let mut interval = tokio::time::interval(SWEEP_INTERVAL);
        loop {
            interval.tick().await;
            limiter.retain_recent();
        }
    });
    Ok(app.layer(GovernorLayer::new(config)))
}

fn shutdown_signal() -> std::io::Result<impl Future<Output = ()> + Send + 'static> {
    let mut terminate = signal(SignalKind::terminate())?;
    let mut interrupt = signal(SignalKind::interrupt())?;
    Ok(async move {
        tokio::select! {
            _ = terminate.recv() => {}
            _ = interrupt.recv() => {}
        }
    })
}

#[derive(Clone)]
struct RateLimitKey {
    trust_proxy_headers: bool,
}

impl KeyExtractor for RateLimitKey {
    type Key = IpAddr;

    fn extract<T>(&self, request: &Request<T>) -> Result<IpAddr, GovernorError> {
        let forwarded = request
            .headers()
            .get_all("x-forwarded-for")
            .iter()
            .filter_map(|value| value.to_str().ok())
            .flat_map(|value| value.split(','))
            .filter_map(|entry| entry.trim().parse::<IpAddr>().ok())
            .next_back()
            .filter(|_| self.trust_proxy_headers);
        forwarded
            .or_else(|| {
                request
                    .extensions()
                    .get::<ConnectInfo<SocketAddr>>()
                    .map(|peer| peer.ip())
            })
            .map(network)
            .ok_or(GovernorError::UnableToExtractKey)
    }
}

fn network(address: IpAddr) -> IpAddr {
    match address {
        IpAddr::V4(_) => address,
        IpAddr::V6(v6) => IpAddr::V6((u128::from(v6) & !((1u128 << 64) - 1)).into()),
    }
}

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

    fn keyed(forwarded: &[&str], peer: &str) -> IpAddr {
        keyed_by(true, forwarded, peer)
    }

    fn keyed_by(trust_proxy_headers: bool, forwarded: &[&str], peer: &str) -> IpAddr {
        let mut request = Request::new(());
        for value in forwarded {
            request
                .headers_mut()
                .append("x-forwarded-for", value.parse().expect("header"));
        }
        request
            .extensions_mut()
            .insert(ConnectInfo(peer.parse::<SocketAddr>().expect("peer")));
        RateLimitKey {
            trust_proxy_headers,
        }
        .extract(&request)
        .expect("key")
    }

    #[test]
    fn the_rate_limit_key_is_the_address_the_proxy_appended() {
        assert_eq!(
            keyed(&[], "10.0.0.1:4000"),
            "10.0.0.1".parse::<IpAddr>().unwrap()
        );
        assert_eq!(
            keyed(&["203.0.113.9, 198.51.100.7"], "10.0.0.1:4000"),
            "198.51.100.7".parse::<IpAddr>().unwrap()
        );
        assert_eq!(
            keyed(&["203.0.113.9", "198.51.100.7"], "10.0.0.1:4000"),
            "198.51.100.7".parse::<IpAddr>().unwrap()
        );
        assert_eq!(
            keyed(&["garbage"], "10.0.0.1:4000"),
            "10.0.0.1".parse::<IpAddr>().unwrap()
        );
        assert_eq!(
            keyed_by(false, &["203.0.113.9"], "10.0.0.1:4000"),
            "10.0.0.1".parse::<IpAddr>().unwrap(),
            "forwarded headers are ignored without a trusted proxy"
        );
    }

    #[test]
    fn ipv6_visitors_share_one_key_per_64() {
        let first = keyed(&[], "[2001:db8:1:2:aaaa::1]:4000");
        let second = keyed(&[], "[2001:db8:1:2:ffff:ffff:ffff:ffff]:4000");
        let other = keyed(&[], "[2001:db8:1:3::1]:4000");
        assert_eq!(first, "2001:db8:1:2::".parse::<IpAddr>().unwrap());
        assert_eq!(first, second);
        assert_ne!(first, other);
        assert_eq!(
            keyed(&["2001:db8:9:9::5"], "10.0.0.1:4000"),
            "2001:db8:9:9::".parse::<IpAddr>().unwrap()
        );
    }
}