use crate::control::{ApiBind, ApiBinds, Hydrant};
use crate::state::AppState;
use axum::{Router, routing::get};
use futures::FutureExt;
use futures::future::BoxFuture;
use std::collections::HashMap;
use std::os::fd::{FromRawFd, IntoRawFd, RawFd};
use std::path::{Path, PathBuf};
use std::sync::{LazyLock, Mutex};
use std::{net::SocketAddr, sync::Arc};
use tower_http::cors::CorsLayer;
use tower_http::trace::TraceLayer;

const LISTEN_BACKLOG: i32 = 1024;

#[cfg(any(feature = "jetstream", feature = "backlinks"))]
pub struct MultiQuery<T>(pub T);

#[cfg(any(feature = "jetstream", feature = "backlinks"))]
impl<T, S> axum::extract::FromRequestParts<S> for MultiQuery<T>
where
    T: serde::de::DeserializeOwned,
    S: Send + Sync,
{
    type Rejection = (axum::http::StatusCode, String);

    async fn from_request_parts(
        parts: &mut axum::http::request::Parts,
        _state: &S,
    ) -> Result<Self, Self::Rejection> {
        serde_html_form::from_str(parts.uri.query().unwrap_or_default())
            .map(Self)
            .map_err(|error| {
                (
                    axum::http::StatusCode::BAD_REQUEST,
                    format!("failed to deserialize query string: {error}"),
                )
            })
    }
}

/// why a repo's records aren't served, none while its account is active. the spec says a
/// service shouldn't hand out an inactive account's content, and a deleted account's records
/// stay readable in the db until its wipe has moved every one of them
#[cfg(feature = "indexer")]
pub(crate) async fn unavailable(
    repo: &crate::control::RepoHandle,
) -> miette::Result<Option<&'static str>> {
    use crate::types::RepoStatus;
    Ok(repo.state().await?.and_then(|state| match state.status {
        RepoStatus::Takendown => Some("repo is takendown"),
        RepoStatus::Suspended => Some("repo is suspended"),
        RepoStatus::Deactivated => Some("repo is deactivated"),
        RepoStatus::Deleted => Some("repo is deleted"),
        _ => None,
    }))
}

#[cfg(feature = "indexer")]
mod crawler;
mod db;
mod debug;
mod filter;
mod firehose;
mod ingestion;
#[cfg(feature = "jetstream")]
mod jetstream;
mod pds;
mod redactions;
mod repos;
mod stats;
#[cfg(feature = "indexer_stream")]
mod stream;
#[cfg(any(feature = "relay", feature = "indexer_stream"))]
mod ws;
mod xrpc;

pub async fn serve(hydrant: Hydrant, binds: ApiBinds) -> miette::Result<()> {
    let record_bodies_available = hydrant.state.stores_record_bodies();
    let app = Router::new()
        .route(
            "/",
            get(async || {
                let kind = cfg!(feature = "indexer")
                    .then_some("indexer")
                    .unwrap_or("relay");
                let subscribe = cfg!(feature = "indexer")
                    .then_some("/stream")
                    .unwrap_or("/xrpc/com.atproto.sync.subscribeRepos");
                include_str!("index.txt")
                    .replace("%type%", kind)
                    .replace("%subscribe_url%", subscribe)
            }),
        )
        .route("/health", get(version))
        .route("/_health", get(version))
        .route("/version", get(version))
        .route("/_version", get(version))
        .route("/stats", get(stats::get_stats));
    #[cfg(feature = "indexer_stream")]
    let app = app.nest("/stream", stream::router());
    #[cfg(feature = "jetstream")]
    let app = app.route("/subscribe", get(jetstream::handle_subscribe));
    let app = app.merge(xrpc::router(record_bodies_available));
    #[cfg(feature = "backlinks")]
    let app = app.merge(crate::backlinks::api::router());

    let management = Router::new()
        .merge(filter::router())
        .merge(pds::router())
        .merge(repos::router())
        .merge(redactions::router())
        .merge(ingestion::router())
        .merge(firehose::router())
        .merge(db::router());
    #[cfg(feature = "indexer")]
    let management = management.merge(crawler::router());

    if !binds.serves_management() {
        tracing::warn!(
            "management api is only served on unix binds and none are configured, add a unix: \
             bind to HYDRANT_API_BIND or set HYDRANT_API_TCP_MANAGEMENT=true"
        );
    }
    let finish = |app: Router<Hydrant>| {
        app.with_state(hydrant.clone())
            .layer(TraceLayer::new_for_http())
            .layer(CorsLayer::permissive())
    };
    let full = finish(app.clone().merge(management));
    let tcp_app = if binds.tcp_management {
        full.clone()
    } else {
        finish(app)
    };

    // bind everything before serving anything, so a bad bind fails startup instead of
    // leaving the api half up
    let tcp_binds: Vec<SocketAddr> = binds.iter().filter_map(ApiBind::tcp).collect();
    let services = binds
        .iter()
        .map(|bind| {
            let service: BoxFuture<'static, std::io::Result<()>> = match bind {
                ApiBind::Tcp(addr) => {
                    let listener = bind_listener(*addr, v6only_for(*addr, &tcp_binds))
                        .map_err(|e| miette::miette!("failed to bind to {bind}: {e}"))?;
                    tracing::info!("API server listening on {}", listener.local_addr().unwrap());
                    axum::serve(
                        listener,
                        tcp_app
                            .clone()
                            .into_make_service_with_connect_info::<SocketAddr>(),
                    )
                    .into_future()
                    .boxed()
                }
                ApiBind::Unix(path) => {
                    let listener = bind_unix(path, binds.unix_mode)
                        .map_err(|e| miette::miette!("failed to bind to {bind}: {e}"))?;
                    tracing::info!("API server listening on {bind}");
                    axum::serve(listener, full.clone().into_make_service())
                        .into_future()
                        .boxed()
                }
            };
            Ok::<_, miette::Report>(
                service.map(|r| r.map_err(|e| miette::miette!("axum server error: {e}"))),
            )
        })
        .collect::<miette::Result<Vec<_>>>()?;

    futures::future::try_join_all(services).await?;

    Ok(())
}

async fn version() -> axum::Json<serde_json::Value> {
    axum::Json(serde_json::json!({
        "name": env!("CARGO_PKG_NAME"),
        "version": env!("CARGO_PKG_VERSION"),
        "mode": cfg!(feature = "indexer").then_some("indexer").unwrap_or("relay"),
    }))
}

fn v6only_for(addr: SocketAddr, all: &[SocketAddr]) -> Option<bool> {
    let SocketAddr::V6(v6) = addr else {
        return None;
    };
    let has_v4_sibling = all
        .iter()
        .any(|a| matches!(a, SocketAddr::V4(v4) if v4.port() == v6.port()));
    Some(has_v4_sibling)
}

const LISTEN_FDS_START: RawFd = 3;

#[derive(Default)]
struct Inherited {
    tcp: HashMap<SocketAddr, std::net::TcpListener>,
    unix: HashMap<PathBuf, std::os::unix::net::UnixListener>,
}

/// systemd holds these sockets across restarts and connects queue until hydrant is back
fn inherited_listeners(
    lookup: impl Fn(&str) -> Option<String>,
    pid: u32,
    first: RawFd,
) -> Inherited {
    let ours = lookup("LISTEN_PID").and_then(|p| p.parse::<u32>().ok()) == Some(pid);
    let count = lookup("LISTEN_FDS")
        .and_then(|n| n.parse::<RawFd>().ok())
        .filter(|_| ours)
        .unwrap_or(0);
    let mut inherited = Inherited::default();
    for fd in first..first.saturating_add(count) {
        // safety: systemd passed this fd to this pid, and it's taken only once
        let socket = unsafe { socket2::Socket::from_raw_fd(fd) };
        let stream = socket.r#type().ok() == Some(socket2::Type::STREAM);
        let local = socket.local_addr().ok().filter(|_| stream);
        if let Some(addr) = local.as_ref().and_then(|addr| addr.as_socket()) {
            inherited.tcp.insert(addr, socket.into());
        } else if let Some(path) = local.as_ref().and_then(|addr| addr.as_pathname()) {
            inherited.unix.insert(path.to_path_buf(), socket.into());
        } else {
            // dropping the socket would close an fd that isn't ours
            let _ = socket.into_raw_fd();
        }
    }
    inherited
}

static INHERITED_LISTENERS: LazyLock<Mutex<Inherited>> = LazyLock::new(|| {
    let env = |key: &str| std::env::var(key).ok();
    Mutex::new(inherited_listeners(
        env,
        std::process::id(),
        LISTEN_FDS_START,
    ))
});

fn bind_listener(
    addr: SocketAddr,
    v6only: Option<bool>,
) -> std::io::Result<tokio::net::TcpListener> {
    let inherited = INHERITED_LISTENERS.lock().unwrap().tcp.remove(&addr);
    if let Some(listener) = inherited {
        tracing::info!("using the listener systemd passed in for {addr}");
        listener.set_nonblocking(true)?;
        return tokio::net::TcpListener::from_std(listener);
    }
    let domain = match addr {
        SocketAddr::V4(_) => socket2::Domain::IPV4,
        SocketAddr::V6(_) => socket2::Domain::IPV6,
    };
    let socket = socket2::Socket::new(domain, socket2::Type::STREAM, Some(socket2::Protocol::TCP))?;
    if let Some(only) = v6only {
        socket.set_only_v6(only)?;
    }
    socket.set_reuse_address(true)?;
    socket.set_nonblocking(true)?;
    socket.bind(&addr.into())?;
    socket.listen(LISTEN_BACKLOG)?;
    let std_listener: std::net::TcpListener = socket.into();
    tokio::net::TcpListener::from_std(std_listener)
}

/// a socket systemd passed in keeps the mode its unit gave it. `mode` is for sockets hydrant binds
fn bind_unix(path: &Path, mode: u32) -> std::io::Result<tokio::net::UnixListener> {
    let inherited = INHERITED_LISTENERS.lock().unwrap().unix.remove(path);
    if let Some(listener) = inherited {
        tracing::info!(
            "using the listener systemd passed in for {}",
            path.display()
        );
        listener.set_nonblocking(true)?;
        return tokio::net::UnixListener::from_std(listener);
    }
    bind_unix_listener(path, mode)
}

/// replaces a socket left behind by a hydrant that died, but never one something is still
/// listening on, and never a file that isn't a socket.
fn bind_unix_listener(path: &Path, mode: u32) -> std::io::Result<tokio::net::UnixListener> {
    use std::io::{Error, ErrorKind};
    use std::os::unix::fs::{DirBuilderExt, FileTypeExt, PermissionsExt};

    match std::fs::symlink_metadata(path) {
        Ok(meta) if meta.file_type().is_socket() => {
            if std::os::unix::net::UnixStream::connect(path).is_ok() {
                return Err(Error::new(
                    ErrorKind::AddrInUse,
                    "another process is listening on this socket",
                ));
            }
        }
        Ok(_) => {
            return Err(Error::new(
                ErrorKind::AlreadyExists,
                "path exists and is not a socket",
            ));
        }
        Err(e) if e.kind() == ErrorKind::NotFound => {}
        Err(e) => return Err(e),
    }

    // binding straight at `path` and chmodding after would let the umask decide who can
    // connect in between, so bind in a dir only we can enter and rename the socket out once
    // it has its mode. the rename also swaps out a stale socket in one step
    let parent = path.parent().unwrap_or(Path::new(""));
    let staging = parent.join(format!(".hydrant-{:08x}", rand::random::<u32>()));
    std::fs::DirBuilder::new().mode(0o700).create(&staging)?;
    let staged = staging.join("s");
    let bound = std::os::unix::net::UnixListener::bind(&staged).and_then(|listener| {
        std::fs::set_permissions(&staged, std::fs::Permissions::from_mode(mode))?;
        std::fs::rename(&staged, path)?;
        listener.set_nonblocking(true)?;
        tokio::net::UnixListener::from_std(listener)
    });
    let _ = std::fs::remove_dir_all(&staging);
    bound
}

pub async fn serve_debug(state: Arc<AppState>, port: u16) -> miette::Result<()> {
    let app = debug::router()
        .with_state(state)
        .layer(TraceLayer::new_for_http());

    let listener = tokio::net::TcpListener::bind(format!("127.0.0.1:{port}"))
        .await
        .map_err(|e| miette::miette!("failed to bind debug server to port {port}: {e}"))?;

    tracing::info!(
        "debug server listening on {}",
        listener.local_addr().unwrap()
    );

    axum::serve(
        listener,
        app.into_make_service_with_connect_info::<SocketAddr>(),
    )
    .await
    .map_err(|e| miette::miette!("debug server error: {e}"))?;

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::os::unix::fs::PermissionsExt;

    fn listen_env(pid: u32, fds: usize) -> impl Fn(&str) -> Option<String> {
        move |key| match key {
            "LISTEN_PID" => Some(pid.to_string()),
            "LISTEN_FDS" => Some(fds.to_string()),
            _ => None,
        }
    }

    #[test]
    fn inherited_listeners_are_found_by_address() {
        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
        let addr = listener.local_addr().unwrap();
        let fd = listener.into_raw_fd();

        let mut found =
            inherited_listeners(listen_env(std::process::id(), 1), std::process::id(), fd);
        let listener = found
            .tcp
            .remove(&addr)
            .expect("the passed listener is keyed by its address");
        assert!(found.tcp.is_empty() && found.unix.is_empty());
        let client = std::net::TcpStream::connect(addr).unwrap();
        let (accepted, _) = listener.accept().unwrap();
        assert_eq!(accepted.peer_addr().unwrap(), client.local_addr().unwrap());
    }

    #[test]
    fn listen_fds_meant_for_another_pid_are_left_alone() {
        use std::os::fd::AsRawFd;
        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
        let fd = listener.as_raw_fd();
        let pid = std::process::id();
        for found in [
            inherited_listeners(listen_env(pid + 1, 1), pid, fd),
            inherited_listeners(|_| None, pid, fd),
        ] {
            assert!(found.tcp.is_empty() && found.unix.is_empty());
        }
        listener.local_addr().unwrap();
    }

    #[test]
    fn inherited_unix_listeners_are_found_by_path() {
        let tmp = tempfile::tempdir().unwrap();
        let path = tmp.path().join("s");
        let listener = std::os::unix::net::UnixListener::bind(&path).unwrap();
        let fd = listener.into_raw_fd();

        let pid = std::process::id();
        let mut found = inherited_listeners(listen_env(pid, 1), pid, fd);
        let listener = found
            .unix
            .remove(&path)
            .expect("the passed listener is keyed by its path");
        assert!(found.tcp.is_empty() && found.unix.is_empty());
        std::os::unix::net::UnixStream::connect(&path).unwrap();
        listener.accept().unwrap();
    }

    #[test]
    fn inherited_sockets_that_are_not_streams_stay_open() {
        use std::os::fd::AsRawFd;
        let tmp = tempfile::tempdir().unwrap();
        let path = tmp.path().join("s");
        let socket = std::os::unix::net::UnixDatagram::bind(&path).unwrap();
        let pid = std::process::id();
        let found = inherited_listeners(listen_env(pid, 1), pid, socket.as_raw_fd());
        assert!(found.tcp.is_empty() && found.unix.is_empty());
        let client = std::os::unix::net::UnixDatagram::unbound().unwrap();
        client.send_to(b"x", &path).unwrap();
        socket.recv(&mut [0; 1]).unwrap();
    }

    #[tokio::test]
    async fn bind_unix_takes_the_listener_systemd_passed_in() {
        use std::os::unix::fs::PermissionsExt;
        let tmp = tempfile::tempdir().unwrap();
        let path = tmp.path().join("api.sock");
        let passed = std::os::unix::net::UnixListener::bind(&path).unwrap();
        std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o660)).unwrap();
        INHERITED_LISTENERS
            .lock()
            .unwrap()
            .unix
            .insert(path.clone(), passed);

        // a live socket would make a fresh bind fail, and a different mode would show a rebind
        let listener = bind_unix(&path, 0o600).unwrap();
        assert_eq!(
            std::fs::metadata(&path).unwrap().permissions().mode() & 0o777,
            0o660
        );
        let client = tokio::net::UnixStream::connect(&path);
        let (accepted, connected) = tokio::join!(listener.accept(), client);
        accepted.unwrap();
        connected.unwrap();
    }

    #[tokio::test]
    async fn unix_bind_uses_the_configured_mode() {
        let tmp = tempfile::tempdir().unwrap();
        let path = tmp.path().join("api.sock");
        let listener = bind_unix_listener(&path, 0o660).unwrap();
        let mode = std::fs::metadata(&path).unwrap().permissions().mode();
        assert_eq!(mode & 0o777, 0o660);

        // the socket still answers after being renamed out of the staging dir, which is gone
        let client = tokio::net::UnixStream::connect(&path);
        let (accepted, connected) = tokio::join!(listener.accept(), client);
        accepted.unwrap();
        connected.unwrap();
        let left: Vec<_> = std::fs::read_dir(tmp.path()).unwrap().collect();
        assert_eq!(left.len(), 1);
    }

    #[tokio::test]
    async fn unix_bind_replaces_a_stale_socket() {
        let tmp = tempfile::tempdir().unwrap();
        let path = tmp.path().join("api.sock");
        // a dropped listener leaves its socket file behind, like a hydrant that crashed
        drop(std::os::unix::net::UnixListener::bind(&path).unwrap());
        assert!(path.exists());

        let listener = bind_unix_listener(&path, 0o600).unwrap();
        let client = tokio::net::UnixStream::connect(&path);
        let (accepted, connected) = tokio::join!(listener.accept(), client);
        accepted.unwrap();
        connected.unwrap();
    }

    #[tokio::test]
    async fn unix_bind_refuses_a_live_socket() {
        let tmp = tempfile::tempdir().unwrap();
        let path = tmp.path().join("api.sock");
        let _live = bind_unix_listener(&path, 0o600).unwrap();

        let err = bind_unix_listener(&path, 0o600).unwrap_err();
        assert_eq!(err.kind(), std::io::ErrorKind::AddrInUse);
        assert!(path.exists());
    }

    #[tokio::test]
    async fn unix_bind_never_deletes_a_regular_file() {
        let tmp = tempfile::tempdir().unwrap();
        let path = tmp.path().join("api.sock");
        std::fs::write(&path, "not a socket").unwrap();

        let err = bind_unix_listener(&path, 0o600).unwrap_err();
        assert_eq!(err.kind(), std::io::ErrorKind::AlreadyExists);
        assert_eq!(std::fs::read_to_string(&path).unwrap(), "not a socket");
    }
}