//! HTTP server for exposing one or more Anda engines.
//!
//! `anda_engine_server` builds an axum router with well-known information
//! endpoints and a signed RPC endpoint for agent/tool calls. Applications can
//! embed the router or run it directly with graceful shutdown.

use anda_core::{BoxError, Json};
use anda_engine::engine::Engine;
use axum::{Router, routing};
use candid::Principal;
use ic_cose_types::cose::ed25519::VerifyingKey;
use std::{collections::BTreeMap, future::Future, net::SocketAddr, sync::Arc};
use structured_logger::unix_ms;
use tokio::signal;
use tokio_util::sync::CancellationToken;

/// HTTP handlers and authentication helpers.
pub mod handler;
/// Pluggable axum middleware helpers.
pub mod middleware;
/// Public response types returned by server endpoints.
pub mod types;

use handler::*;
use middleware::*;

const APP_NAME: &str = env!("CARGO_PKG_NAME");
const APP_VERSION: &str = env!("CARGO_PKG_VERSION");

/// Builder for an Anda engine HTTP server.
pub struct ServerBuilder {
    app_name: String,
    app_version: String,
    addr: String,
    engines: BTreeMap<Principal, Arc<Engine>>,
    default_engine: Option<Principal>,
    middlewares: Vec<Box<dyn HttpMiddleware>>,
    ed25519_pubkeys: Vec<VerifyingKey>,
    extra_info: BTreeMap<String, Json>,
}

impl Default for ServerBuilder {
    fn default() -> Self {
        Self::new()
    }
}

/// Builder for creating a new server.
impl ServerBuilder {
    /// Creates a new ServerBuilder with default values.
    pub fn new() -> Self {
        ServerBuilder {
            app_name: APP_NAME.to_string(),
            app_version: APP_VERSION.to_string(),
            addr: "127.0.0.1:8042".to_string(),
            engines: BTreeMap::new(),
            default_engine: None,
            middlewares: Vec::new(),
            ed25519_pubkeys: Vec::new(),
            extra_info: BTreeMap::new(),
        }
    }

    /// Sets the application name used in [`Self::serve`] startup logs.
    pub fn with_app_name(mut self, app_name: String) -> Self {
        self.app_name = app_name;
        self
    }

    /// Sets the application version used in [`Self::serve`] startup logs.
    pub fn with_app_version(mut self, app_version: String) -> Self {
        self.app_version = app_version;
        self
    }

    /// Sets the TCP bind address, for example `127.0.0.1:8042`.
    pub fn with_addr(mut self, addr: String) -> Self {
        self.addr = addr;
        self
    }

    /// Adds extra metadata to information responses.
    pub fn with_extra_info(mut self, extra_info: BTreeMap<String, Json>) -> Self {
        self.extra_info = extra_info;
        self
    }

    /// Registers engines and optionally selects the default engine principal.
    ///
    /// Each map key must equal its engine's [`Engine::id`]; [`Self::build_router`]
    /// rejects mismatches. Without an explicit default, the first map key is used.
    pub fn with_engines(
        mut self,
        engines: BTreeMap<Principal, Arc<Engine>>,
        default_engine: Option<Principal>,
    ) -> Self {
        self.engines = engines;
        self.default_engine = default_engine;
        self
    }

    /// Sets trusted Ed25519 public keys for bearer CWT authentication.
    pub fn with_ed25519_pubkeys(mut self, pubkeys: Vec<VerifyingKey>) -> Self {
        self.ed25519_pubkeys = pubkeys;
        self
    }

    /// Register a router middleware.
    ///
    /// This is the low-level API. The middleware will be applied to the internal
    /// axum `Router` (typically via `router.layer(...)`). Middlewares are applied
    /// in the order they are added. Each `router.layer(...)` wraps the previous
    /// layers, so adding A then B runs requests as B → A → handler and responses
    /// as handler → A → B.
    ///
    /// More details: <https://docs.rs/axum/latest/axum/middleware/index.html#ordering>
    ///
    /// If you want a middleware that looks like `axum::middleware::from_fn`
    /// (i.e. can operate on `(req, next)`), prefer [`Self::with_request_middleware`].
    ///
    /// Example:
    /// ```
    /// use anda_engine_server::{ServerBuilder, middleware::AppRouter};
    /// use axum::{extract::Request, middleware::Next};
    ///
    /// let server = ServerBuilder::new()
    ///   .with_middleware(|router: AppRouter| {
    ///     router.layer(axum::middleware::from_fn(|req: Request, next: Next| async move {
    ///       // custom auth / param checks here
    ///       next.run(req).await
    ///     }))
    ///   });
    /// ```
    pub fn with_middleware<M>(mut self, middleware: M) -> Self
    where
        M: HttpMiddleware,
    {
        self.middlewares.push(Box::new(middleware));
        self
    }

    /// Register a request middleware like `axum::middleware::from_fn`.
    ///
    /// The middleware function runs for every incoming request, and can decide
    /// to short-circuit with a response or call `next.run(req)`.
    ///
    /// Example:
    /// ```
    /// use anda_engine_server::ServerBuilder;
    /// use axum::http::StatusCode;
    /// use axum::response::IntoResponse;
    ///
    /// let server = ServerBuilder::new()
    ///   .with_request_middleware(|req, next| async move {
    ///     // custom auth / param checks here
    ///     if req.headers().get("x-allow").is_none() {
    ///       return (StatusCode::UNAUTHORIZED, "missing x-allow").into_response();
    ///     }
    ///
    ///     next.run(req).await
    ///   });
    /// ```
    pub fn with_request_middleware<F, Fut>(self, f: F) -> Self
    where
        F: Fn(axum::extract::Request, axum::middleware::Next) -> Fut
            + Clone
            + Send
            + Sync
            + 'static,
        Fut: Future<Output = axum::response::Response> + Send + 'static,
    {
        self.with_middleware(middleware::RequestFnMiddleware::new(f))
    }

    /// Builds the axum [`Router`] with all routes, registered middlewares, and
    /// application state applied.
    ///
    /// Useful for embedding the engine server into a larger axum application
    /// or driving it with a custom listener; [`Self::serve`] uses it internally.
    pub fn build_router(self) -> Result<Router, BoxError> {
        let default_engine = self
            .default_engine
            .or_else(|| self.engines.keys().next().copied())
            .ok_or("no engines registered")?;
        if !self.engines.contains_key(&default_engine) {
            return Err("default engine not found".into());
        }
        for (id, engine) in &self.engines {
            if *id != engine.id() {
                return Err(format!(
                    "registered engine ID {id} does not match engine ID {}",
                    engine.id()
                )
                .into());
            }
        }

        let state = AppState {
            engines: Arc::new(self.engines),
            default_engine,
            start_time_ms: unix_ms(),
            extra_info: Arc::new(self.extra_info),
            ed25519_pubkeys: Arc::new(self.ed25519_pubkeys),
        };

        // Build a router that is still "missing" an `AppState`.
        // We'll provide the state at the end (after applying middlewares) so we
        // end up with a `Router<()>` that can be passed to `axum::serve`.
        let mut app: Router<AppState> = Router::new()
            .route("/", routing::get(get_information))
            .route("/.well-known/information", routing::get(get_information))
            .route("/.well-known/agents", routing::get(get_information))
            .route(
                "/.well-known/agents/{id}",
                routing::get(get_engine_information),
            )
            .route("/{*id}", routing::post(anda_engine));

        for middleware in &self.middlewares {
            app = middleware.apply(app);
        }

        Ok(app.with_state(state))
    }

    /// Serves the router on the configured address until `signal` resolves.
    pub async fn serve(
        self,
        signal: impl Future<Output = ()> + Send + 'static,
    ) -> Result<(), BoxError> {
        let addr: SocketAddr = self
            .addr
            .parse()
            .map_err(|err| format!("invalid server address {:?}: {err}", self.addr))?;
        let app_name = self.app_name.clone();
        let app_version = self.app_version.clone();
        let app = self.build_router()?;

        let listener = create_reuse_port_listener(addr).await?;
        log::warn!("{}@{} listening on {:?}", app_name, app_version, addr);

        axum::serve(listener, app)
            .with_graceful_shutdown(signal)
            .await?;

        Ok(())
    }
}

/// Waits for Ctrl+C or SIGTERM, then cancels the provided token.
pub async fn shutdown_signal(cancel_token: CancellationToken) {
    let ctrl_c = async {
        signal::ctrl_c()
            .await
            .expect("failed to install Ctrl+C handler");
    };

    #[cfg(unix)]
    let terminate = async {
        signal::unix::signal(signal::unix::SignalKind::terminate())
            .expect("failed to install signal handler")
            .recv()
            .await;
    };

    #[cfg(not(unix))]
    let terminate = std::future::pending::<()>();

    tokio::select! {
        _ = ctrl_c => {},
        _ = terminate => {},
    }

    log::warn!("received termination signal, starting graceful shutdown");
    cancel_token.cancel();
}

/// Creates a TCP listener with `SO_REUSEPORT` on supported Unix targets.
///
/// Other Unix targets fall back to `SO_REUSEADDR`. Windows sets neither: there
/// `SO_REUSEADDR` would let another socket bind the same port and take over
/// its connections.
pub async fn create_reuse_port_listener(
    addr: SocketAddr,
) -> Result<tokio::net::TcpListener, BoxError> {
    let socket = match &addr {
        SocketAddr::V4(_) => tokio::net::TcpSocket::new_v4()?,
        SocketAddr::V6(_) => tokio::net::TcpSocket::new_v6()?,
    };

    #[cfg(all(
        unix,
        not(target_os = "solaris"),
        not(target_os = "illumos"),
        not(target_os = "cygwin"),
    ))]
    socket.set_reuseport(true)?;

    #[cfg(all(
        unix,
        any(target_os = "solaris", target_os = "illumos", target_os = "cygwin",),
    ))]
    socket.set_reuseaddr(true)?;

    socket.bind(addr)?;
    let listener = socket.listen(1024)?;
    Ok(listener)
}