product_os_router/

lib.rs

//! # Product OS Router
//!
//! A fully-featured HTTP router built on top of [Axum](https://github.com/tokio-rs/axum) and [Tower](https://github.com/tower-rs/tower),
//! providing convenient helper methods for creating HTTP, HTTPS, WebSocket, and Server-Sent Events (SSE) servers.
//!
//! ## Features
//!
//! - **Easy Routing**: Simple, chainable API for defining routes
//! - **HTTP Methods**: Support for GET, POST, PUT, PATCH, DELETE, and more
//! - **State Management**: Both stateful and stateless routing options
//! - **WebSockets**: Built-in WebSocket support (with `ws` feature)
//! - **Server-Sent Events**: SSE support (with `sse` feature)
//! - **CORS**: Cross-Origin Resource Sharing support (with `cors` feature)
//! - **Middleware**: Custom middleware support (with `middleware` feature)
//! - **Default Headers**: Apply default headers to all responses
//! - **Protocol Upgrade**: HTTP to HTTPS automatic redirection
//! - **Path Parameters**: Express-style path parameter syntax (`:param`, `*wildcard`)
//! - **no_std Support**: Can be used in no_std environments with alloc
//!
//! ## Quick Start
//!
//! ```rust
//! use product_os_router::{ProductOSRouter, IntoResponse};
//!
//! async fn hello_world() -> &'static str {
//!     "Hello, World!"
//! }
//!
//! # async fn example() {
//! let mut router = ProductOSRouter::new();
//! router.add_get_no_state("/", hello_world);
//!
//! let app = router.get_router();
//! // Use app with your server (e.g., axum::Server)
//! # }
//! ```
//!
//! ## Path Parameter Syntax
//!
//! The router automatically converts Express-style path parameters to Axum format:
//! - `:param` becomes `{param}` (single segment)
//! - `*wildcard` becomes `{wildcard}` (catch-all)
//!
//! ```rust
//! # use product_os_router::{ProductOSRouter, Path};
//! async fn user_handler(Path(id): Path<u32>) -> String {
//!     format!("User ID: {}", id)
//! }
//!
//! # async fn example() {
//! let mut router = ProductOSRouter::new();
//! router.add_get_no_state("/users/:id", user_handler);  // Converted to /users/{id}
//! # }
//! ```
//!
//! ## With State
//!
//! ```rust
//! # use product_os_router::{ProductOSRouter, State};
//! #[derive(Clone)]
//! struct AppState {
//!     counter: u32,
//! }
//!
//! async fn handler(State(state): State<AppState>) -> String {
//!     format!("Counter: {}", state.counter)
//! }
//!
//! # async fn example() {
//! let state = AppState { counter: 0 };
//! let mut router = ProductOSRouter::new_with_state(state);
//! router.add_get("/", handler);
//! # }
//! ```
//!
//! ## Feature Flags
//!
//! - `default`: Includes `framework_axum` and `std` features
//! - `framework_axum`: Axum framework (existing, stable)
//! - `framework_flow`: Flow framework (new, high-performance)
//! - `std`: Standard library support
//! - `core`: Core routing functionality
//! - `ws`: WebSocket support
//! - `sse`: Server-Sent Events support
//! - `cors`: CORS middleware support
//! - `sessions`: Session management support
//! - `middleware`: Custom middleware traits
//! - `debug`: Debug handler macros

#![no_std]
extern crate no_std_compat as std;

use std::prelude::v1::*;

// ============================================================================
// Compile-Time Feature Validation
// ============================================================================

// Ensure at least one framework is enabled
#[cfg(not(any(feature = "framework_axum", feature = "framework_flow")))]
compile_error!(
    "At least one framework feature must be enabled: 'framework_axum' or 'framework_flow'"
);

// Prevent using both frameworks simultaneously
#[cfg(all(feature = "framework_axum", feature = "framework_flow"))]
compile_error!(
    "Both 'framework_axum' and 'framework_flow' features cannot be enabled simultaneously. \
     \n\nChoose ONE framework:\
     \n  1. Axum (default):  --features framework_axum\
     \n  2. Flow (new):      --no-default-features --features framework_flow,std"
);

// ============================================================================
// Module Imports
// ============================================================================

// Conditional module imports based on feature flags
#[cfg(feature = "framework_axum")]
mod axum_impl;

#[cfg(feature = "framework_flow")]
pub mod flow_impl;

// BTreeMap needed by both frameworks
#[cfg(any(feature = "framework_axum", feature = "framework_flow"))]
use std::collections::BTreeMap;

#[cfg(any(feature = "framework_axum", feature = "core"))]
pub use product_os_http::{
    header::{HeaderMap, HeaderName, HeaderValue},
    request::Parts as RequestParts,
    response::Parts as ResponseParts,
    Request, Response, StatusCode,
};

#[cfg(feature = "core")]
pub use product_os_http_body::{BodyBytes, Bytes};

#[cfg(feature = "middleware")]
pub use product_os_http_body::*;

#[cfg(feature = "framework_axum")]
pub use axum::{
    body::{Body, HttpBody},
    extract::{
        ws::{Message, WebSocket, WebSocketUpgrade},
        ConnectInfo, FromRef, FromRequest, FromRequestParts, Path, Query, State,
    },
    handler::Handler,
    http::Uri,
    middleware::{from_fn, from_fn_with_state, Next},
    response::{IntoResponse, Redirect},
    routing::{any, delete, get, head, options, patch, post, put, trace, MethodRouter},
    BoxError, Form, Json, Router,
};

#[cfg(all(feature = "sse", feature = "framework_axum"))]
pub use axum::response::sse::{Event, Sse};

#[cfg(feature = "framework_axum")]
pub use crate::extractors::RequestMethod;

// When only Flow is enabled (without axum compat), re-export Flow-native types.
#[cfg(all(
    feature = "framework_flow",
    not(feature = "framework_axum"),
    not(feature = "flow_axum_compat")
))]
pub use flow_impl::core::body::Body;
#[cfg(all(
    feature = "framework_flow",
    not(feature = "framework_axum"),
    not(feature = "flow_axum_compat")
))]
pub use flow_impl::core::Handler;
#[cfg(all(
    feature = "framework_flow",
    not(feature = "framework_axum"),
    not(feature = "flow_axum_compat")
))]
pub use flow_impl::core::{FromRequest, FromRequestParts, IntoResponse};
#[cfg(all(
    feature = "framework_flow",
    not(feature = "framework_axum"),
    not(feature = "flow_axum_compat")
))]
pub use flow_impl::extractors::form::StreamingForm as Form;
#[cfg(all(
    feature = "framework_flow",
    not(feature = "framework_axum"),
    not(feature = "flow_axum_compat")
))]
pub use flow_impl::extractors::json::StreamingJson as Json;
#[cfg(all(
    feature = "framework_flow",
    not(feature = "framework_axum"),
    not(feature = "flow_axum_compat")
))]
pub use flow_impl::extractors::path::Path;
#[cfg(all(
    feature = "framework_flow",
    not(feature = "framework_axum"),
    not(feature = "flow_axum_compat")
))]
pub use flow_impl::extractors::query::Query;
#[cfg(all(
    feature = "framework_flow",
    not(feature = "framework_axum"),
    not(feature = "flow_axum_compat")
))]
pub use flow_impl::extractors::state::State;
#[cfg(all(
    feature = "framework_flow",
    not(feature = "framework_axum"),
    not(feature = "flow_axum_compat")
))]
pub use flow_impl::router::NextService;
#[cfg(all(
    feature = "framework_flow",
    not(feature = "framework_axum"),
    not(feature = "flow_axum_compat")
))]
pub use product_os_http::uri::Uri;
#[cfg(all(
    feature = "framework_flow",
    not(feature = "framework_axum"),
    not(feature = "flow_axum_compat")
))]
pub use product_os_http_body::Body as HttpBody;

#[cfg(all(
    feature = "framework_flow",
    not(feature = "framework_axum"),
    not(feature = "flow_axum_compat")
))]
pub type BoxError = std::boxed::Box<dyn core::error::Error + Send + Sync>;

#[cfg(all(
    feature = "sse",
    feature = "framework_flow",
    not(feature = "framework_axum"),
    not(feature = "flow_axum_compat")
))]
pub use flow_impl::sse::{Event, KeepAlive, Sse};

#[cfg(all(
    feature = "ws",
    feature = "framework_flow",
    not(feature = "framework_axum"),
    not(feature = "flow_axum_compat")
))]
pub use flow_impl::ws::{Message, WebSocket};

#[cfg(any(feature = "framework_axum", feature = "core"))]
pub use product_os_http::uri::Scheme;

// When flow_axum_compat is enabled, re-export from Axum for full compatibility.
// Path, Query, and State use Flow-native extractors since the Flow router
// populates extensions differently than Axum's internal machinery.
#[cfg(all(feature = "flow_axum_compat", not(feature = "framework_axum")))]
pub use axum::{
    body::{Body, HttpBody},
    extract::{
        ws::{Message, WebSocket, WebSocketUpgrade},
        ConnectInfo, FromRef, FromRequest, FromRequestParts, State,
    },
    http::Uri,
    response::{IntoResponse, Redirect},
    BoxError, Form, Json,
};

// Path and Query use Flow-native extractors under flow_axum_compat because
// the Flow router populates PathParams in extensions, not Axum's internal UrlParams.
#[cfg(all(feature = "flow_axum_compat", not(feature = "framework_axum")))]
pub use flow_impl::extractors::path::Path;
#[cfg(all(feature = "flow_axum_compat", not(feature = "framework_axum")))]
pub use flow_impl::extractors::query::Query;

#[cfg(all(
    feature = "sse",
    feature = "flow_axum_compat",
    not(feature = "framework_axum")
))]
pub use axum::response::sse::{Event, Sse};

#[cfg(all(feature = "flow_axum_compat", not(feature = "framework_axum")))]
pub use flow_impl::core::Handler;

#[cfg(all(feature = "flow_axum_compat", not(feature = "framework_axum")))]
pub use flow_impl::router::NextService;

#[cfg(all(feature = "flow_axum_compat", not(feature = "framework_axum")))]
pub use axum::routing::MethodRouter;

#[cfg(all(
    feature = "debug",
    any(feature = "framework_axum", feature = "flow_axum_compat")
))]
pub use axum_macros::debug_handler;

#[cfg(feature = "core")]
pub use tower::{
    make::future::SharedFuture, make::AsService, make::IntoService, make::MakeConnection,
    make::Shared, util::service_fn, util::ServiceFn, Layer, MakeService, Service, ServiceBuilder,
    ServiceExt,
};

#[cfg(feature = "core")]
pub use product_os_request::Method;

#[cfg(feature = "middleware")]
pub mod middleware;

#[cfg(feature = "framework_axum")]
mod default_headers;
#[cfg(feature = "core")]
pub mod dual_protocol;
#[cfg(feature = "framework_axum")]
mod extractors;
#[cfg(feature = "framework_axum")]
mod service_wrapper;

#[cfg(feature = "framework_axum")]
pub use crate::default_headers::DefaultHeadersLayer;
#[cfg(feature = "core")]
pub use crate::dual_protocol::{Protocol, UpgradeHttpLayer};

#[cfg(feature = "middleware")]
pub use crate::middleware::*;

#[cfg(feature = "framework_flow")]
pub use flow_impl::extractors::headers::Headers;

#[cfg(all(feature = "framework_axum", not(feature = "framework_flow")))]
mod axum_request_id;
#[cfg(all(feature = "framework_axum", not(feature = "framework_flow")))]
mod axum_trace_context;

#[cfg(all(feature = "framework_axum", not(feature = "framework_flow")))]
pub use axum_request_id::RequestId;
#[cfg(all(feature = "framework_axum", not(feature = "framework_flow")))]
pub use axum_trace_context::TraceContext;

#[cfg(feature = "framework_flow")]
pub use flow_impl::extractors::trace_context::TraceContext;

#[cfg(all(
    feature = "framework_flow",
    not(feature = "framework_axum"),
    not(feature = "flow_axum_compat")
))]
pub use flow_impl::extractors::context::{ClientAddr, ConnectInfo, RequestId};

#[cfg(all(feature = "framework_flow", feature = "flow_axum_compat"))]
pub use flow_impl::extractors::context::{ClientAddr, RequestId};

/// Error types for router operations
///
/// Provides semantic error types for different HTTP error scenarios.
/// Each variant maps to an appropriate HTTP status code.
///
/// # Examples
///
/// ```rust
/// use product_os_router::{ProductOSRouterError, StatusCode};
///
/// let error = ProductOSRouterError::Authentication("Invalid token".to_string());
/// let response = ProductOSRouterError::handle_error(&error);
/// assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
/// ```
use core::fmt;

#[derive(Debug)]
pub enum ProductOSRouterError {
    /// Invalid header error (400 Bad Request)
    Headers(String),
    /// Invalid query parameter error (400 Bad Request)
    Query(String),
    /// Invalid request body error (400 Bad Request)
    Body(String),
    /// Authentication failure (401 Unauthorized)
    Authentication(String),
    /// Authorization failure (403 Forbidden)
    Authorization(String),
    /// Processing timeout or failure (504 Gateway Timeout)
    Process(String),
    /// Service unavailable (503 Service Unavailable)
    Unavailable(String),
}

impl fmt::Display for ProductOSRouterError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            ProductOSRouterError::Headers(m) => write!(f, "header error: {m}"),
            ProductOSRouterError::Query(m) => write!(f, "query error: {m}"),
            ProductOSRouterError::Body(m) => write!(f, "body error: {m}"),
            ProductOSRouterError::Authentication(m) => write!(f, "authentication error: {m}"),
            ProductOSRouterError::Authorization(m) => write!(f, "authorization error: {m}"),
            ProductOSRouterError::Process(m) => write!(f, "process error: {m}"),
            ProductOSRouterError::Unavailable(m) => write!(f, "unavailable: {m}"),
        }
    }
}

#[cfg(feature = "std")]
extern crate std as _real_std;

#[cfg(feature = "std")]
impl _real_std::error::Error for ProductOSRouterError {}

impl ProductOSRouterError {
    /// Extract the error message from any variant.
    pub fn message(&self) -> &str {
        match self {
            ProductOSRouterError::Headers(m) => m,
            ProductOSRouterError::Query(m) => m,
            ProductOSRouterError::Body(m) => m,
            ProductOSRouterError::Authentication(m) => m,
            ProductOSRouterError::Authorization(m) => m,
            ProductOSRouterError::Process(m) => m,
            ProductOSRouterError::Unavailable(m) => m,
        }
    }
}

#[cfg(feature = "core")]
impl ProductOSRouterError {
    /// Return the appropriate HTTP status code for this error variant.
    pub fn status_code(&self) -> StatusCode {
        match self {
            ProductOSRouterError::Headers(_) => StatusCode::BAD_REQUEST,
            ProductOSRouterError::Query(_) => StatusCode::BAD_REQUEST,
            ProductOSRouterError::Body(_) => StatusCode::BAD_REQUEST,
            ProductOSRouterError::Authentication(_) => StatusCode::UNAUTHORIZED,
            ProductOSRouterError::Authorization(_) => StatusCode::FORBIDDEN,
            ProductOSRouterError::Process(_) => StatusCode::GATEWAY_TIMEOUT,
            ProductOSRouterError::Unavailable(_) => StatusCode::SERVICE_UNAVAILABLE,
        }
    }
}

#[cfg(feature = "std")]
impl ProductOSRouterError {
    /// Build a JSON error response body from a message and status code.
    ///
    /// Quotes in the error message are replaced with single quotes to
    /// ensure valid JSON output.
    fn build_error_response(message: &str, status_code: &StatusCode) -> Response<Body> {
        let safe_message = message.replace("\"", "'");
        let status_code_u16 = status_code.as_u16();

        Response::builder()
            .status(status_code_u16)
            .header("content-type", "application/json")
            .body(Body::from(format!("{{\"error\": \"{safe_message}\"}}")))
            .unwrap_or_else(|_| {
                Response::builder()
                    .status(StatusCode::INTERNAL_SERVER_ERROR.as_u16())
                    .body(Body::from("{\"error\": \"internal error\"}"))
                    .expect("fallback response must be valid")
            })
    }

    /// Create an HTTP response from an error with a specific status code.
    ///
    /// Consumes the error and converts it into a JSON response with the error
    /// message. Quotes in the error message are replaced with single quotes.
    ///
    /// # Arguments
    ///
    /// * `error` - The error to convert (consumed)
    /// * `status_code` - The HTTP status code to use
    ///
    /// # Returns
    ///
    /// An HTTP response with JSON body containing the error message
    ///
    /// # Examples
    ///
    /// ```rust
    /// use product_os_router::{ProductOSRouterError, StatusCode};
    ///
    /// let error = ProductOSRouterError::Headers("Invalid header".to_string());
    /// let response = ProductOSRouterError::error_response(error, &StatusCode::BAD_REQUEST);
    /// ```
    pub fn error_response(error: ProductOSRouterError, status_code: &StatusCode) -> Response<Body> {
        Self::build_error_response(error.message(), status_code)
    }

    /// Handle an error and return an appropriate HTTP response.
    ///
    /// Automatically determines the appropriate HTTP status code based on the
    /// error variant and returns a JSON response with the error message.
    ///
    /// # Arguments
    ///
    /// * `error` - The error to handle
    ///
    /// # Returns
    ///
    /// An HTTP response with the appropriate status code and error message
    ///
    /// # Examples
    ///
    /// ```rust
    /// use product_os_router::ProductOSRouterError;
    ///
    /// let error = ProductOSRouterError::Authentication("Login required".to_string());
    /// let response = ProductOSRouterError::handle_error(&error);
    /// ```
    pub fn handle_error(error: &ProductOSRouterError) -> Response<Body> {
        Self::build_error_response(error.message(), &error.status_code())
    }
}

/// Main router enum for Product OS Router
///
/// Dispatches to either Axum or Flow implementation based on feature flags.
/// Provides a high-level API for defining routes, handlers, and middleware.
/// Supports both stateful and stateless routing.
///
/// # Type Parameters
///
/// * `S` - The state type (defaults to `()` for stateless routing)
///
/// # Examples
///
/// ## Stateless Router
///
/// ```rust
/// use product_os_router::ProductOSRouter;
///
/// async fn handler() -> &'static str {
///     "Hello!"
/// }
///
/// let mut router = ProductOSRouter::new();
/// router.add_get_no_state("/", handler);
/// ```
///
/// ## Stateful Router
///
/// ```rust
/// use product_os_router::{ProductOSRouter, State};
///
/// #[derive(Clone)]
/// struct AppState {
///     message: String,
/// }
///
/// async fn handler(State(state): State<AppState>) -> String {
///     state.message
/// }
///
/// let state = AppState { message: "Hello!".to_string() };
/// let mut router = ProductOSRouter::new_with_state(state);
/// router.add_get("/", handler);
/// ```
#[cfg(any(feature = "framework_axum", feature = "framework_flow"))]
pub enum ProductOSRouter<S = ()> {
    #[cfg(feature = "framework_axum")]
    Axum(axum_impl::AxumRouter<S>),

    #[cfg(feature = "framework_flow")]
    Flow(flow_impl::Router<S>),
}

#[cfg(feature = "framework_axum")]
impl<S: Clone + Send + Sync + 'static> Clone for ProductOSRouter<S> {
    fn clone(&self) -> Self {
        match self {
            Self::Axum(r) => Self::Axum(r.clone()),
        }
    }
}

#[cfg(feature = "framework_flow")]
impl<S: Clone + Send + Sync + 'static> Clone for ProductOSRouter<S> {
    fn clone(&self) -> Self {
        match self {
            Self::Flow(r) => Self::Flow(r.clone()),
        }
    }
}

#[cfg(any(feature = "framework_axum", feature = "framework_flow"))]
impl<S: Clone + Send + Sync + 'static> ProductOSRouter<S> {
    /// Create a new router with the given shared state.
    pub fn new_with_state(state: S) -> Self {
        #[cfg(feature = "framework_axum")]
        return Self::Axum(axum_impl::AxumRouter::new_with_state(state));

        #[cfg(feature = "framework_flow")]
        return Self::Flow(flow_impl::Router::new_with_state(state));
    }
}

#[cfg(any(feature = "framework_axum", feature = "framework_flow"))]
impl ProductOSRouter<()> {
    /// Create a new stateless router
    ///
    /// Creates a router without any shared state. Use [`new_with_state`](Self::new_with_state)
    /// if you need to share state between handlers.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use product_os_router::ProductOSRouter;
    ///
    /// let router = ProductOSRouter::new();
    /// ```
    pub fn new() -> Self {
        ProductOSRouter::new_with_state(())
    }
}

#[cfg(any(feature = "framework_axum", feature = "framework_flow"))]
impl Default for ProductOSRouter<()> {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(any(feature = "framework_axum", feature = "framework_flow"))]
impl ProductOSRouter<()> {
    /// Convert path parameters from Express-style to Axum format
    pub fn param_to_field(path: &str) -> String {
        #[cfg(feature = "framework_axum")]
        {
            axum_impl::AxumRouter::<()>::param_to_field(path)
        }
        #[cfg(all(feature = "framework_flow", not(feature = "framework_axum")))]
        {
            flow_impl::Router::<()>::param_to_field(path)
        }
    }

    /// Returns a 501 Not Implemented JSON response
    pub fn not_implemented() -> Response<Body> {
        Response::builder()
            .status(StatusCode::NOT_IMPLEMENTED.as_u16())
            .header("content-type", "application/json")
            .body(Body::from("{}"))
            .unwrap()
    }

    /// Returns a 404 Not Found JSON response
    #[deprecated(
        since = "0.0.41",
        note = "Use not_implemented() which correctly returns 501, or build a 404 response explicitly"
    )]
    pub fn not_implemented_legacy() -> Response<Body> {
        Response::builder()
            .status(StatusCode::NOT_FOUND.as_u16())
            .header("content-type", "application/json")
            .body(Body::from("{}"))
            .unwrap()
    }

    /// Add a GET route handler without state
    pub fn add_get_no_state<H, T>(&mut self, path: &str, handler: H)
    where
        H: Handler<T, ()> + Send + Sync + 'static,
        T: 'static,
    {
        match self {
            #[cfg(feature = "framework_axum")]
            Self::Axum(r) => r.add_get_no_state(path, handler),
            #[cfg(feature = "framework_flow")]
            Self::Flow(r) => r.add_get_no_state(path, handler),
        }
    }

    /// Add a POST route handler without state
    pub fn add_post_no_state<H, T>(&mut self, path: &str, handler: H)
    where
        H: Handler<T, ()> + Send + Sync + 'static,
        T: 'static,
    {
        match self {
            #[cfg(feature = "framework_axum")]
            Self::Axum(r) => r.add_post_no_state(path, handler),
            #[cfg(feature = "framework_flow")]
            Self::Flow(r) => r.add_post_no_state(path, handler),
        }
    }

    /// Add a handler for the given HTTP method without state
    pub fn add_handler_no_state<H, T>(&mut self, path: &str, method: Method, handler: H)
    where
        H: Handler<T, ()> + Send + Sync + 'static,
        T: 'static,
    {
        match self {
            #[cfg(feature = "framework_axum")]
            Self::Axum(r) => r.add_handler_no_state(path, method, handler),
            #[cfg(feature = "framework_flow")]
            Self::Flow(r) => r.add_handler_no_state(path, method, handler),
        }
    }

    /// Add a default header without state
    pub fn add_default_header_no_state(&mut self, header_name: &str, header_value: &str) {
        match self {
            #[cfg(feature = "framework_axum")]
            Self::Axum(r) => r.add_default_header_no_state(header_name, header_value),
            #[cfg(feature = "framework_flow")]
            Self::Flow(r) => r.add_default_header_no_state(header_name, header_value),
        }
    }

    /// Set a fallback handler without state
    pub fn set_fallback_handler_no_state<H, T>(&mut self, handler: H)
    where
        H: Handler<T, ()> + Send + Sync + 'static,
        T: 'static,
    {
        match self {
            #[cfg(feature = "framework_axum")]
            Self::Axum(r) => r.set_fallback_handler_no_state(handler),
            #[cfg(feature = "framework_flow")]
            Self::Flow(r) => r.set_fallback_handler_no_state(handler),
        }
    }

    /// Add a CORS handler without state
    #[cfg(feature = "cors")]
    pub fn add_cors_handler_no_state<H, T>(&mut self, path: &str, method: Method, handler: H)
    where
        H: Handler<T, ()> + Send + Sync + 'static,
        T: 'static,
    {
        match self {
            #[cfg(feature = "framework_axum")]
            Self::Axum(r) => r.add_cors_handler_no_state(path, method, handler),
            #[cfg(feature = "framework_flow")]
            Self::Flow(r) => r.add_cors_handler(path, method, handler),
        }
    }

    /// Add an SSE handler without state
    #[cfg(feature = "sse")]
    pub fn add_sse_handler_no_state<H, T>(&mut self, path: &str, handler: H)
    where
        H: Handler<T, ()> + Send + Sync + 'static,
        T: 'static,
    {
        match self {
            #[cfg(feature = "framework_axum")]
            Self::Axum(r) => r.add_sse_handler_no_state(path, handler),
            #[cfg(feature = "framework_flow")]
            Self::Flow(r) => r.add_sse_handler(path, handler),
        }
    }

    /// Add a WebSocket handler without state (Axum / flow_axum_compat)
    #[cfg(all(
        feature = "ws",
        any(feature = "framework_axum", feature = "flow_axum_compat")
    ))]
    pub fn add_ws_handler_no_state<H, T>(&mut self, path: &str, handler: H)
    where
        H: Handler<T, ()> + Send + Sync + 'static,
        T: 'static,
    {
        match self {
            #[cfg(feature = "framework_axum")]
            Self::Axum(r) => r.add_ws_handler_no_state(path, handler),
            #[cfg(feature = "framework_flow")]
            Self::Flow(r) => r.add_ws_handler(path, handler),
        }
    }

    #[cfg(all(
        feature = "ws",
        feature = "framework_flow",
        not(feature = "framework_axum"),
        not(feature = "flow_axum_compat")
    ))]
    pub fn add_ws_handler_no_state<F, Fut>(&mut self, path: &str, handler: F)
    where
        F: Fn(flow_impl::ws::WebSocket) -> Fut + Send + Sync + Clone + 'static,
        Fut: std::future::Future<Output = ()> + Send + 'static,
    {
        match self {
            Self::Flow(r) => r.add_ws_handler(path, handler),
        }
    }

    /// Add multiple handlers without state
    pub fn add_handlers_no_state<H, T>(&mut self, path: &str, handlers: BTreeMap<Method, H>)
    where
        H: Handler<T, ()> + Send + Sync + 'static,
        T: 'static,
    {
        match self {
            #[cfg(feature = "framework_axum")]
            Self::Axum(r) => r.add_handlers_no_state(path, handlers),
            #[cfg(feature = "framework_flow")]
            Self::Flow(r) => r.add_handlers_no_state(path, handlers),
        }
    }

    /// Add multiple CORS-enabled handlers without state
    #[cfg(feature = "cors")]
    pub fn add_cors_handlers_no_state<H, T>(&mut self, path: &str, handlers: BTreeMap<Method, H>)
    where
        H: Handler<T, ()> + Send + Sync + 'static,
        T: 'static,
    {
        match self {
            #[cfg(feature = "framework_axum")]
            Self::Axum(r) => r.add_cors_handlers_no_state(path, handlers),
            #[cfg(feature = "framework_flow")]
            Self::Flow(r) => r.add_cors_handlers_no_state(path, handlers),
        }
    }

    /// Add a CORS middleware without state
    #[cfg(feature = "cors")]
    pub fn add_cors_middleware_no_state(&mut self) {
        match self {
            #[cfg(feature = "framework_axum")]
            Self::Axum(r) => r.add_cors_middleware_no_state(),
            #[cfg(feature = "framework_flow")]
            Self::Flow(r) => r.add_cors_middleware(),
        }
    }
}

#[cfg(any(feature = "framework_axum", feature = "framework_flow"))]
impl<S> ProductOSRouter<S>
where
    S: Clone + Send + Sync + 'static,
{
    /// Return a clone of the router.
    ///
    /// For both Axum and Flow, returns `ProductOSRouter<S>`.
    /// Use `into_axum_router()` when you need the raw `axum::Router`.
    pub fn get_router(&self) -> ProductOSRouter<S> {
        self.clone()
    }

    /// Extract the inner `axum::Router`, consuming the wrapper.
    ///
    /// This is needed when handing the router to axum-specific server functions
    /// that expect a raw `axum::Router`.
    #[cfg(feature = "framework_axum")]
    pub fn into_axum_router(self) -> Router {
        match self {
            Self::Axum(r) => r.get_router(),
        }
    }

    /// Returns a clone of the shared state stored in the router.
    #[cfg(feature = "framework_flow")]
    pub fn get_state(&self) -> S {
        match self {
            Self::Flow(r) => r.get_state(),
        }
    }

    /// Add a GET route handler
    pub fn add_get<H, T>(&mut self, path: &str, handler: H)
    where
        H: Handler<T, S> + Send + Sync + 'static,
        T: 'static,
    {
        match self {
            #[cfg(feature = "framework_axum")]
            Self::Axum(r) => r.add_get(path, handler),
            #[cfg(feature = "framework_flow")]
            Self::Flow(r) => r.add_get(path, handler),
        }
    }

    /// Add a POST route handler
    pub fn add_post<H, T>(&mut self, path: &str, handler: H)
    where
        H: Handler<T, S> + Send + Sync + 'static,
        T: 'static,
    {
        match self {
            #[cfg(feature = "framework_axum")]
            Self::Axum(r) => r.add_post(path, handler),
            #[cfg(feature = "framework_flow")]
            Self::Flow(r) => r.add_post(path, handler),
        }
    }

    /// Add a handler for the given HTTP method and path
    pub fn add_handler<H, T>(&mut self, path: &str, method: Method, handler: H)
    where
        H: Handler<T, S> + Send + Sync + 'static,
        T: 'static,
    {
        match self {
            #[cfg(feature = "framework_axum")]
            Self::Axum(r) => r.add_handler(path, method, handler),
            #[cfg(feature = "framework_flow")]
            Self::Flow(r) => r.add_handler(path, method, handler),
        }
    }

    /// Add a default header
    pub fn add_default_header(&mut self, header_name: &str, header_value: &str) {
        match self {
            #[cfg(feature = "framework_axum")]
            Self::Axum(r) => r.add_default_header(header_name, header_value),
            #[cfg(feature = "framework_flow")]
            Self::Flow(r) => r.add_default_header(header_name, header_value),
        }
    }

    /// Add CORS middleware
    #[cfg(feature = "cors")]
    pub fn add_cors_middleware(&mut self) {
        match self {
            #[cfg(feature = "framework_axum")]
            Self::Axum(r) => r.add_cors_middleware(),
            #[cfg(feature = "framework_flow")]
            Self::Flow(r) => r.add_cors_middleware(),
        }
    }

    /// Add multiple HTTP method handlers at a single path
    pub fn add_handlers<H, T>(&mut self, path: &str, handlers: BTreeMap<Method, H>)
    where
        H: Handler<T, S> + Send + Sync + 'static,
        T: 'static,
    {
        match self {
            #[cfg(feature = "framework_axum")]
            Self::Axum(r) => r.add_handlers(path, handlers),
            #[cfg(feature = "framework_flow")]
            Self::Flow(r) => r.add_handlers(path, handlers),
        }
    }

    /// Add CORS-enabled handlers for different methods on the same path
    #[cfg(feature = "cors")]
    pub fn add_cors_handlers<H, T>(&mut self, path: &str, handlers: BTreeMap<Method, H>)
    where
        H: Handler<T, S> + Send + Sync + 'static,
        T: 'static,
    {
        match self {
            #[cfg(feature = "framework_axum")]
            Self::Axum(r) => r.add_cors_handlers(path, handlers),
            #[cfg(feature = "framework_flow")]
            Self::Flow(r) => r.add_cors_handlers(path, handlers),
        }
    }

    /// Add a CORS handler for a specific HTTP method and path
    #[cfg(feature = "cors")]
    pub fn add_cors_handler<H, T>(&mut self, path: &str, method: Method, handler: H)
    where
        H: Handler<T, S> + Send + Sync + 'static,
        T: 'static,
    {
        match self {
            #[cfg(feature = "framework_axum")]
            Self::Axum(r) => r.add_cors_handler(path, method, handler),
            #[cfg(feature = "framework_flow")]
            Self::Flow(r) => r.add_cors_handler(path, method, handler),
        }
    }

    /// Add a WebSocket handler for the specified path (Axum / flow_axum_compat)
    #[cfg(all(
        feature = "ws",
        any(feature = "framework_axum", feature = "flow_axum_compat")
    ))]
    pub fn add_ws_handler<H, T>(&mut self, path: &str, handler: H)
    where
        H: Handler<T, S> + Send + Sync + 'static,
        T: 'static,
    {
        match self {
            #[cfg(feature = "framework_axum")]
            Self::Axum(r) => r.add_ws_handler(path, handler),
            #[cfg(feature = "framework_flow")]
            Self::Flow(r) => r.add_ws_handler(path, handler),
        }
    }

    /// Add a native Flow WebSocket handler for the specified path.
    ///
    /// The handler receives a `WebSocket` after the HTTP upgrade has been
    /// performed by the server layer.
    #[cfg(all(
        feature = "ws",
        feature = "framework_flow",
        not(feature = "framework_axum"),
        not(feature = "flow_axum_compat")
    ))]
    pub fn add_ws_handler<F, Fut>(&mut self, path: &str, handler: F)
    where
        F: Fn(flow_impl::ws::WebSocket) -> Fut + Send + Sync + Clone + 'static,
        Fut: std::future::Future<Output = ()> + Send + 'static,
    {
        match self {
            Self::Flow(r) => r.add_ws_handler(path, handler),
        }
    }

    /// Look up a native WebSocket handler for the given path.
    #[cfg(all(
        feature = "ws",
        feature = "framework_flow",
        not(feature = "framework_axum"),
        not(feature = "flow_axum_compat")
    ))]
    pub fn lookup_ws_handler(
        &mut self,
        path: &str,
    ) -> Option<flow_impl::router::WsHandlerCallback> {
        match self {
            Self::Flow(r) => r.lookup_ws_handler(path),
        }
    }

    /// Add an SSE handler for the specified path
    #[cfg(feature = "sse")]
    pub fn add_sse_handler<H, T>(&mut self, path: &str, handler: H)
    where
        H: Handler<T, S> + Send + Sync + 'static,
        T: 'static,
    {
        match self {
            #[cfg(feature = "framework_axum")]
            Self::Axum(r) => r.add_sse_handler(path, handler),
            #[cfg(feature = "framework_flow")]
            Self::Flow(r) => r.add_sse_handler(path, handler),
        }
    }

    /// Add a route with an Axum `MethodRouter`.
    #[cfg(feature = "framework_axum")]
    pub fn add_route(&mut self, path: &str, method_router: axum::routing::MethodRouter<S>)
    where
        S: Clone + Send + Sync + 'static,
    {
        match self {
            Self::Axum(r) => r.add_route(path, method_router),
        }
    }

    /// Set the fallback service for unmatched routes (Axum `MethodRouter`).
    #[cfg(feature = "framework_axum")]
    pub fn set_fallback(&mut self, service_handler: axum::routing::MethodRouter<S>)
    where
        S: Clone + Send + Sync + 'static,
    {
        match self {
            Self::Axum(r) => r.set_fallback(service_handler),
        }
    }

    /// Set the fallback handler for unmatched routes.
    pub fn set_fallback_handler<H, T>(&mut self, handler: H)
    where
        H: Handler<T, S>,
        T: 'static,
    {
        match self {
            #[cfg(feature = "framework_axum")]
            Self::Axum(r) => r.set_fallback_handler(handler),
            #[cfg(feature = "framework_flow")]
            Self::Flow(r) => r.set_fallback_handler(handler),
        }
    }

    /// Add an Axum-compatible Tower middleware layer.
    #[cfg(feature = "framework_axum")]
    pub fn add_middleware<L, ResBody>(&mut self, middleware: L)
    where
        L: tower::Layer<axum::routing::Route> + Clone + Send + Sync + 'static,
        L::Service: tower::Service<Request<Body>, Response = Response<ResBody>>
            + Clone
            + Send
            + Sync
            + 'static,
        <L::Service as tower::Service<Request<Body>>>::Future: Send,
        <L::Service as tower::Service<Request<Body>>>::Error: Into<axum::BoxError>,
        ResBody: axum::body::HttpBody<Data = product_os_http_body::Bytes> + Send + 'static,
        ResBody::Error: Into<axum::BoxError>,
    {
        match self {
            Self::Axum(r) => r.add_middleware(middleware),
        }
    }

    /// Add a Tower middleware layer for the Flow framework.
    ///
    /// Uses body-erasing to convert arbitrary `Response<ResBody>` into
    /// `Response<Body>`, enabling standard Tower layers (CSRF, compression, etc.)
    /// to be used with Flow's middleware pipeline.
    #[cfg(all(feature = "framework_flow", feature = "flow_axum_compat"))]
    pub fn add_middleware<L, ResBody>(&mut self, middleware: L)
    where
        L: tower::Layer<flow_impl::router::NextService<S>> + Clone + Send + Sync + 'static,
        L::Service: tower::Service<Request<Body>, Response = product_os_http::Response<ResBody>>
            + Clone
            + Send
            + Sync
            + 'static,
        <L::Service as tower::Service<Request<Body>>>::Future: Send + 'static,
        <L::Service as tower::Service<Request<Body>>>::Error: Into<axum::BoxError> + Send + 'static,
        ResBody: axum::body::HttpBody<Data = product_os_http_body::Bytes> + Send + 'static,
        ResBody::Error: Into<axum::BoxError>,
    {
        match self {
            Self::Flow(r) => r.add_middleware_erased(middleware),
        }
    }

    /// Add a Tower middleware layer for pure Flow (without `flow_axum_compat`).
    ///
    /// The layer wraps `NextService<S>` and must produce
    /// `Response<Body>` with `Error = Infallible`.
    #[cfg(all(
        feature = "framework_flow",
        not(feature = "flow_axum_compat"),
        not(feature = "framework_axum")
    ))]
    pub fn add_middleware<L>(&mut self, middleware: L)
    where
        L: tower::Layer<flow_impl::router::NextService<S>> + Clone + Send + Sync + 'static,
        L::Service: tower::Service<
                Request<Body>,
                Response = flow_impl::core::body::Response,
                Error = core::convert::Infallible,
            > + Clone
            + Send
            + Sync
            + 'static,
        <L::Service as tower::Service<Request<Body>>>::Future: Send + 'static,
    {
        match self {
            Self::Flow(r) => r.add_middleware(middleware),
        }
    }
}

#[cfg(any(feature = "framework_axum", feature = "framework_flow"))]
impl<S: Clone + Send + Sync + 'static> ProductOSRouter<S> {
    /// Inserts a value into request extensions for all subsequent handlers.
    pub fn add_extension<T: Clone + Send + Sync + 'static>(&mut self, value: T) {
        #[cfg(feature = "framework_axum")]
        {
            match self {
                Self::Axum(r) => {
                    let router = r.take_router().layer(axum::Extension(value));
                    *r.get_router_mut() = router;
                }
            }
        }
        #[cfg(all(feature = "framework_flow", not(feature = "framework_axum")))]
        {
            match self {
                Self::Flow(r) => {
                    r.add_extension(value);
                }
            }
        }
    }
}

// Delegate tower::Service to the inner router variant.
// This enables `ProductOSRouter` to be used directly with hyper.
#[cfg(feature = "framework_flow")]
impl<S> tower::Service<Request<Body>> for ProductOSRouter<S>
where
    S: Clone + Send + Sync + 'static,
{
    type Response = flow_impl::core::body::Response;
    type Error = core::convert::Infallible;
    type Future = core::pin::Pin<
        Box<dyn core::future::Future<Output = Result<Self::Response, Self::Error>> + Send>,
    >;

    fn poll_ready(
        &mut self,
        cx: &mut core::task::Context<'_>,
    ) -> core::task::Poll<Result<(), Self::Error>> {
        match self {
            Self::Flow(r) => tower::Service::poll_ready(r, cx),
        }
    }

    fn call(&mut self, req: Request<Body>) -> Self::Future {
        match self {
            Self::Flow(r) => tower::Service::call(r, req),
        }
    }
}

#[cfg(feature = "std")]
mod extension_impl {
    use crate::{Body, FromRequestParts, Request};
    use std::ops::{Deref, DerefMut};
    use std::prelude::v1::*;

    /// Extractor that gets a value from request extensions.
    ///
    /// This is commonly used with middleware that inserts data into extensions.
    /// In axum 0.7+, prefer using `State` for application-wide state.
    ///
    /// # Example
    ///
    /// ```ignore
    /// use product_os_router::{Extension, ServiceBuilder, MethodRouter};
    /// use std::sync::Arc;
    ///
    /// async fn handler(Extension(value): Extension<Arc<String>>) -> String {
    ///     (*value).clone()
    /// }
    ///
    /// let shared_data = Arc::new("Hello".to_string());
    /// let route = MethodRouter::new()
    ///     .get(handler)
    ///     .layer(ServiceBuilder::new().layer(tower::layer::util::AddExtensionLayer::new(shared_data)));
    /// ```
    #[derive(Debug, Clone, Copy, Default)]
    pub struct Extension<T>(pub T);

    impl<T> Deref for Extension<T> {
        type Target = T;

        fn deref(&self) -> &Self::Target {
            &self.0
        }
    }

    impl<T> DerefMut for Extension<T> {
        fn deref_mut(&mut self) -> &mut Self::Target {
            &mut self.0
        }
    }

    impl<T, S> FromRequestParts<S> for Extension<T>
    where
        T: Clone + Send + Sync + 'static,
        S: Send + Sync,
    {
        type Rejection = (product_os_http::StatusCode, &'static str);

        fn from_request_parts(
            parts: &mut product_os_http::request::Parts,
            _state: &S,
        ) -> impl core::future::Future<Output = Result<Self, Self::Rejection>> + Send {
            let value = parts.extensions.get::<T>().cloned().map(Extension).ok_or((
                product_os_http::StatusCode::INTERNAL_SERVER_ERROR,
                "Extension not found in request",
            ));

            async move { value }
        }
    }

    pub fn extension_layer<T>(
        value: T,
    ) -> tower::util::MapRequestLayer<impl Fn(Request<Body>) -> Request<Body> + Clone>
    where
        T: Clone + Send + Sync + 'static,
    {
        tower::util::MapRequestLayer::new(move |mut req: Request<Body>| {
            req.extensions_mut().insert(value.clone());
            req
        })
    }
}

#[cfg(feature = "std")]
pub use extension_impl::{extension_layer, Extension};