use axum::Router;

use hyper_util::rt::{TokioExecutor, TokioIo};

use hyper_util::server::conn::auto;

use hyper_util::service::TowerToHyperService;

use lyquor_api::anyhow;

use lyquor_image_store::OciDistributionStore;

use lyquor_jsonrpc::client::ClientHandle;

use std::sync::Arc;

use tokio::net::TcpListener;

use tokio_util::sync::CancellationToken;

use tokio_util::task::TaskTracker;

pub mod httptls;

mod jsonrpc;

mod lyquid;

mod node;

mod oci;

pub use self::lyquid::LyquidServiceContext;

pub use self::node::NodeServiceContext;

pub struct HttpServer {

    local_addr: String,

    shutdown: CancellationToken,

    task_tracker: TaskTracker,

}

impl std::fmt::Debug for HttpServer {

    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {

        f.debug_struct("HttpServer")

            .field("local_addr", &self.local_addr)

            .finish()

    }

}

impl HttpServer {

    pub async fn new(

        bind_addr: &str, client: ClientHandle, fco: lyquor_seq::fco::FCOHandle,

        process_manager: lyquor_hosting::ProcessManagerHandle, registry: lyquor_hosting::LyquidRegistry,

        node_service_context: NodeServiceContext, oci_store: Arc<dyn OciDistributionStore>,

        shutdown: CancellationToken, http_tls: Option<httptls::HttpsTlsRuntime>,

    ) -> anyhow::Result<Self> {

        let listener = TcpListener::bind(bind_addr).await?;

        let local_addr = listener.local_addr()?.to_string();

        let router = build_router(

            jsonrpc::JsonRpcRouteState::new(

                jsonrpc::JsonRpcContext::new(client.clone(), process_manager.clone(), registry.clone(), fco),

                shutdown.child_token(),

            ),

            Arc::new(LyquidServiceContext::new(process_manager.clone(), registry)),

            Arc::new(node_service_context.clone()),

            oci::OciRouteState::new(oci_store.clone()),

        );

        let cert_state_rx = http_tls.as_ref().map(httptls::HttpsTlsRuntime::cert_state_receiver);

        let task_tracker = TaskTracker::new();

        task_tracker.spawn(run_http_server(listener, router, cert_state_rx, shutdown.clone()));

        // Don't change this line as it is parsed by tests.

        tracing::info!("API listening at {local_addr}");

        Ok(Self {

            local_addr,

            shutdown,

            task_tracker,

        })

    }

    pub fn stop(&self) {

        // Single shutdown signal for all HTTP/JSON-RPC tasks; cleanup is handled

        // by the token-driven tasks spawned in `start` and `JsonRpcRouteState::new`.

        self.shutdown.cancel();

        self.task_tracker.close();

    }

    pub async fn wait(&self) {

        self.task_tracker.wait().await;

    }

}

fn build_router(

    jsonrpc_state: jsonrpc::JsonRpcRouteState, lyquid_service_context: Arc<lyquid::LyquidServiceContext>,

    node_service_context: Arc<node::NodeServiceContext>, oci_state: oci::OciRouteState,

) -> Router {

    // Keep top-level router composition in mod.rs so adding REST APIs later only

    // extends this function.

    jsonrpc::build_jsonrpc_router(jsonrpc_state)

        .merge(oci::build_oci_router(oci_state))

        .merge(node::build_node_router(node_service_context))

        .merge(lyquid::build_lyquid_router(lyquid_service_context))

}

async fn serve_connection<I>(i: I, app: Router, shutdown: CancellationToken)

where

    I: tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + Unpin + 'static,

{

    let service = TowerToHyperService::new(app);

    let builder = auto::Builder::new(TokioExecutor::new());

    let mut connection = std::pin::pin!(builder.serve_connection_with_upgrades(TokioIo::new(i), service));

    tokio::select! {

        biased;

        _ = shutdown.cancelled() => {

            // Shutdown signal received, stop accepting new requests but allow in-flight requests to complete.

            tracing::info!("Shutting down HTTP connection.");

            connection.as_mut().graceful_shutdown();

        }

        res = connection.as_mut() => {

            if let Err(e) = res {

                tracing::trace!("HTTP API connection error: {e:?}.");

            }

        }

    }

}

async fn run_http_server(

    listener: TcpListener, router: Router, cert_state_rx: Option<httptls::HttpsCertReceiver>,

    shutdown: CancellationToken,

) {

    loop {

        let (socket, _remote) = tokio::select! {

            biased;

            _ = shutdown.cancelled() => break,

            res = listener.accept() => match res {

                Ok(pair) => pair,

                Err(e) => {

                    tracing::error!("Failed to accept HTTP API connection: {e:?}.");

                    continue;

                }

            },

        };

        let tls = cert_state_rx.as_ref().map(|rx| rx.acceptor());

        let app = router.clone();

        let connection_shutdown = shutdown.child_token();

        tokio::spawn(async move {

            if let Some(tls) = tls {

                match tls.accept(socket).await {

                    Ok(tls_stream) => serve_connection(tls_stream, app, connection_shutdown).await,

                    Err(e) => {

                        tracing::trace!("HTTPS handshake error: {e:?}.");

                        return;

                    }

                }

            } else {

                serve_connection(socket, app, connection_shutdown).await;

            }

        });

    }

}