aro-web 1.0.0

//! Application builder for the Aro web framework.

use axum::Router;
use axum::extract::DefaultBodyLimit;
use std::future::Future;
use std::net::SocketAddr;
use std::pin::Pin;
use tokio::net::TcpListener;
use tower_http::trace::TraceLayer;

use std::sync::Arc;

use crate::fallback::{fallback_404, fallback_405};
use crate::state::{AroState, AroStateBuilder};

/// A boxed startup hook: an async function that receives the resolved `AroState`
/// and may fail with any error.
type StartupHook = Box<
    dyn FnOnce(
            Arc<AroState>,
        ) -> Pin<Box<dyn Future<Output = Result<(), Box<dyn std::error::Error>>>>>
        + Send,
>;

struct UserLayer {
    apply: Box<dyn Fn(Router) -> Router + Send + Sync>,
}

/// Builder for an Aro web application.
///
/// Wraps an Axum `Router` and applies framework defaults (tracing middleware,
/// JSON error fallbacks) while allowing full customization.
///
/// # Example
///
/// ```no_run
/// use aro_web::App;
/// use axum::Router;
/// use axum::routing::get;
///
/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
/// App::new()
///     .routes(Router::new().route("/", get(|| async { "ok" })))
///     .serve("0.0.0.0:3000".parse()?)
///     .await?;
/// # Ok(())
/// # }
/// ```
pub struct App {
    router: Router,
    stateful_router: Option<Router<AroState>>,
    state: Option<AroState>,
    builder: AroStateBuilder,
    default_middleware: bool,
    body_limit: Option<usize>,
    user_layers: Vec<UserLayer>,
    startup_hooks: Vec<StartupHook>,
}

impl App {
    /// Create a new `App` with default settings.
    pub fn new() -> Self {
        Self {
            router: Router::new(),
            stateful_router: None,
            state: None,
            builder: AroState::builder(),
            default_middleware: true,
            body_limit: None,
            user_layers: Vec::new(),
            startup_hooks: Vec::new(),
        }
    }

    /// Merge the given router's routes into the application.
    ///
    /// Can be called multiple times to compose routes from different modules.
    pub fn routes(mut self, router: Router) -> Self {
        self.router = self.router.merge(router);
        self
    }

    /// Set the application state for dependency injection.
    ///
    /// The state is made available to handlers via the
    /// [`Dep<T>`](crate::dep::Dep) extractor.
    ///
    /// Registrations from [`register`](Self::register) take precedence over
    /// entries with the same type when both are present at [`build`](Self::build) time.
    pub fn state(mut self, state: AroState) -> Self {
        self.state = Some(state);
        self
    }

    /// Register a dependency for injection.
    ///
    /// The type parameter `T` should be a trait object type (e.g. `dyn UserRepo`).
    /// Registered dependencies are accumulated internally and merged into the
    /// final [`AroState`] at [`build`](Self::build) time. When combined with
    /// [`state`](Self::state), registrations from this method take precedence
    /// over entries with the same type from external state.
    ///
    /// # Example
    ///
    /// ```ignore
    /// // Domain types (`UserRepo`, `UserRepoImpl`) come from your app.
    /// App::new()
    ///     .register::<dyn UserRepo>(Arc::new(UserRepoImpl::new(pool)))
    ///     .routes_with_state(routes)
    ///     .build()
    /// ```
    pub fn register<T: ?Sized + 'static>(mut self, service: Arc<T>) -> Self
    where
        Arc<T>: Send + Sync + Clone + 'static,
    {
        self.builder = self.builder.register::<T>(service);
        self
    }

    /// Merge routes that require `AroState` into the application.
    ///
    /// Use this for routers containing handlers that use `Dep<T>` extractors.
    /// The state must be set via [`state`](Self::state) before calling
    /// [`build`](Self::build).
    pub fn routes_with_state(mut self, router: Router<AroState>) -> Self {
        self.stateful_router = Some(match self.stateful_router {
            Some(existing) => existing.merge(router),
            None => router,
        });
        self
    }

    /// Add a tower `Layer` to the application.
    ///
    /// User layers wrap the default middleware stack (outermost).
    /// They execute in the order they are added and apply to all routes,
    /// including those registered via [`routes_with_state`](Self::routes_with_state).
    /// Each call wraps the stack, so the last layer added is outermost on inbound
    /// requests.
    pub fn layer<L>(mut self, layer: L) -> Self
    where
        L: tower::Layer<axum::routing::Route> + Clone + Send + Sync + 'static,
        L::Service:
            tower::Service<axum::http::Request<axum::body::Body>> + Clone + Send + Sync + 'static,
        <L::Service as tower::Service<axum::http::Request<axum::body::Body>>>::Response:
            axum::response::IntoResponse + 'static,
        <L::Service as tower::Service<axum::http::Request<axum::body::Body>>>::Error:
            Into<std::convert::Infallible> + 'static,
        <L::Service as tower::Service<axum::http::Request<axum::body::Body>>>::Future:
            Send + 'static,
    {
        self.user_layers.push(UserLayer {
            apply: Box::new(move |router| router.layer(layer.clone())),
        });
        self
    }

    /// Set the maximum allowed request body size in bytes.
    ///
    /// Applies [`DefaultBodyLimit::max`] as a global layer, overriding
    /// Axum's default 2 MB limit for extractors like `Json`, `Form`, and
    /// `Bytes`. Individual routes can still override this with a
    /// per-route `DefaultBodyLimit` layer.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use aro_web::App;
    /// use axum::Router;
    ///
    /// let _app = App::new()
    ///     .body_limit(10 * 1024 * 1024) // 10 MB
    ///     .routes(Router::new())
    ///     .build();
    /// ```
    pub fn body_limit(mut self, max_bytes: usize) -> Self {
        self.body_limit = Some(max_bytes);
        self
    }

    /// Load layered configuration and register it as a dependency.
    ///
    /// Uses [`load_config`](crate::config::load_config) to build `T`
    /// from environment variables and config files, then registers the result
    /// so it can be extracted with [`Dep<T>`](crate::dep::Dep).
    ///
    /// # Errors
    ///
    /// Returns an error if configuration loading or deserialization fails.
    pub fn config<T: serde::de::DeserializeOwned + Send + Sync + 'static>(
        self,
    ) -> Result<Self, crate::config::ConfigError> {
        let cfg: T = crate::config::load_config()?;
        Ok(self.register::<T>(Arc::new(cfg)))
    }

    /// Register an async startup hook.
    ///
    /// Hooks run in the order they are registered during [`serve()`](Self::serve)
    /// or [`build_with_hooks()`](Self::build_with_hooks), **before** the TCP
    /// listener is bound. Each hook receives a shared reference to the resolved
    /// [`AroState`]. If any hook returns an error, the server will not start
    /// and `serve()` will return the error.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use aro_web::App;
    /// use axum::Router;
    ///
    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
    /// App::new()
    ///     .on_startup(|_state| Box::pin(async move { Ok(()) }))
    ///     .routes(Router::new())
    ///     .serve("127.0.0.1:0".parse()?)
    ///     .await?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn on_startup<F>(mut self, hook: F) -> Self
    where
        F: FnOnce(
                Arc<AroState>,
            )
                -> Pin<Box<dyn Future<Output = Result<(), Box<dyn std::error::Error>>>>>
            + Send
            + 'static,
    {
        self.startup_hooks.push(Box::new(hook));
        self
    }

    /// Apply a request timeout to all routes.
    ///
    /// Requests that are not completed within the given duration will receive
    /// a `408 Request Timeout` response.
    ///
    /// This is the canonical builder entry point: the timeout applies to all
    /// routes, including stateful DI routes added via
    /// [`routes_with_state`](Self::routes_with_state).
    ///
    /// For manual `Router::layer(...)` composition, see
    /// [`crate::middleware::timeout`].
    ///
    /// # Example
    ///
    /// ```no_run
    /// use std::time::Duration;
    ///
    /// use aro_web::App;
    /// use axum::Router;
    ///
    /// let _app = App::new()
    ///     .routes(Router::new())
    ///     .timeout(Duration::from_secs(30))
    ///     .build();
    /// ```
    pub fn timeout(self, duration: std::time::Duration) -> Self {
        self.layer(crate::middleware::timeout(duration))
    }

    /// Disable the default middleware stack (`TraceLayer`).
    ///
    /// Call this if you want full control over the middleware pipeline.
    /// The JSON fallback handler is always applied.
    pub fn without_defaults(mut self) -> Self {
        self.default_middleware = false;
        self
    }

    /// Build the final `Router`, applying default middleware.
    ///
    /// The default middleware stack includes:
    /// - JSON 404/405 fallback handlers for unmatched routes and unsupported methods
    /// - `TraceLayer` from tower-http for request/response logging
    ///
    /// **Note:** Startup hooks registered via [`on_startup`](Self::on_startup)
    /// are *not* executed by `build()`. Use [`build_with_hooks()`](Self::build_with_hooks)
    /// or [`serve()`](Self::serve) to run them.
    pub fn build(self) -> Router {
        let (router, _state, _hooks) = self.into_parts();
        router
    }

    /// Build the final `Router`, running all startup hooks first.
    ///
    /// This is the same as [`build()`](Self::build) but runs all registered
    /// startup hooks in order before returning the router. Use this in tests
    /// when your app relies on startup hooks (e.g., database migrations).
    ///
    /// # Errors
    ///
    /// Returns an error if any startup hook fails.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use aro_web::App;
    /// use axum::Router;
    ///
    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
    /// let _router = App::new()
    ///     .on_startup(|_state| Box::pin(async move { Ok(()) }))
    ///     .routes(Router::new())
    ///     .build_with_hooks()
    ///     .await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn build_with_hooks(self) -> Result<Router, Box<dyn std::error::Error>> {
        let (router, state, hooks) = self.into_parts();

        let shared_state = Arc::new(state);
        for hook in hooks {
            hook(Arc::clone(&shared_state)).await?;
        }

        Ok(router)
    }

    /// Build and start serving the application.
    ///
    /// Runs all registered startup hooks in order, then binds to the given
    /// address and serves with graceful shutdown (handles SIGINT/Ctrl+C and
    /// SIGTERM).
    ///
    /// **Note:** Like [`build_with_hooks`](Self::build_with_hooks), this runs
    /// startup hooks. [`build`](Self::build) alone does not.
    ///
    /// If any startup hook returns an error, the server is **not** started
    /// and the error is returned.
    pub async fn serve(self, addr: SocketAddr) -> Result<(), Box<dyn std::error::Error>> {
        let app = self.build_with_hooks().await?;

        let listener = TcpListener::bind(addr).await?;
        tracing::info!("listening on {}", addr);
        axum::serve(listener, app)
            .with_graceful_shutdown(shutdown_signal())
            .await?;
        Ok(())
    }

    /// Consume the App and produce the final router, resolved state, and startup hooks.
    fn into_parts(self) -> (Router, AroState, Vec<StartupHook>) {
        let mut router = self.router;

        // Build the final AroState by merging the internal builder with any
        // externally provided state. Builder registrations (`.register()`)
        // take precedence over entries from `.state()`.
        let state = match self.state {
            Some(external) => external.merge(self.builder.build()),
            None => self.builder.build(),
        };

        // Merge stateful routes with the resolved state
        if let Some(stateful) = self.stateful_router {
            router = router.merge(stateful.with_state(state.clone()));
        }

        // Apply JSON fallbacks for unmatched routes and unsupported methods
        router = router
            .fallback(fallback_404)
            .method_not_allowed_fallback(fallback_405);

        // Apply body limit
        if let Some(max_bytes) = self.body_limit {
            router = router.layer(DefaultBodyLimit::max(max_bytes));
        }

        // Apply default middleware
        if self.default_middleware {
            router = router.layer(TraceLayer::new_for_http());
        }

        // Apply user layers outermost so they wrap the default stack
        for user_layer in self.user_layers {
            router = (user_layer.apply)(router);
        }

        (router, state, self.startup_hooks)
    }
}

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

/// Wait for a shutdown signal (SIGINT or SIGTERM).
#[expect(
    clippy::expect_used,
    reason = "signal handler installation failure is unrecoverable"
)]
async fn shutdown_signal() {
    let ctrl_c = async {
        tokio::signal::ctrl_c()
            .await
            .expect("failed to install Ctrl+C handler");
    };

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

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

    tokio::select! {
        () = ctrl_c => {},
        () = terminate => {},
    }

    tracing::info!("shutdown signal received, starting graceful shutdown");
}