webrune 0.1.2

use std::{collections::HashMap, sync::Arc};

use http::{Method, StatusCode};
use http_body_util::Full;
use hyper::body::Bytes;
use matchit::Match;

use crate::{
    Flow::{self, Continue, Exit},
    FromRequest, Handler, IntoResponse, Request, Response,
    body::Empty,
    router::endpoint::{Endpoint, EndpointHandler, MethodHandler},
};

mod endpoint;

/// Internal identifier for a registered route.
///
/// Routes are deduplicated by path; multiple HTTP methods
/// can be associated with the same `RouteId`.
#[derive(Clone, Eq, PartialEq, Hash, Default)]
struct RouteId(u32);

/// An HTTP request router.
///
/// `Router` matches incoming requests by:
/// 1. request path (using [`matchit`]), then
/// 2. HTTP method,
/// and dispatches the request to a corresponding endpoint handler.
///
/// Routes are registered using a builder-style API:
///
/// ```ignore
/// let router = Router::new()
///     .get("/users", list_users)
///     .post("/users", create_user);
/// ```
///
/// The router itself implements [`Handler`], allowing it to be
/// used directly as a server handler or composed with other handlers.
#[derive(Default)]
pub struct Router<S> {
    /// Path matcher used to resolve routes.
    inner: matchit::Router<RouteId>,

    /// Mapping from route identifiers to endpoint definitions.
    routes: HashMap<RouteId, Endpoint<S>>,

    /// Next available route identifier.
    next_id: RouteId,

    /// Mapping from path strings to route identifiers.
    ///
    /// This allows multiple methods to be registered for the same path.
    path_to_id: HashMap<String, RouteId>,
}

impl<S> Router<S>
where
    S: Clone + Send + Sync + 'static,
{
    /// Creates an empty router.
    pub fn new() -> Self {
        Self {
            inner: matchit::Router::new(),
            routes: HashMap::new(),
            next_id: RouteId(0),
            path_to_id: HashMap::new(),
        }
    }

    /// Generates a new unique route identifier.
    fn get_next_id(&mut self) -> RouteId {
        let id = self.next_id.clone();
        self.next_id = RouteId(id.0 + 1);
        id
    }

    /// Registers a `POST` handler for the given path.
    pub fn post<I, O>(
        self,
        path: &str,
        handler: impl EndpointHandler<I, S, Output = O, Future: Send> + Clone + Send + Sync + 'static,
    ) -> Self
    where
        I: FromRequest<Output = I> + Send + 'static,
        <I as FromRequest>::Error: Send,
        <I as FromRequest>::from_request(..): Send,
        O: IntoResponse + 'static,
    {
        self.insert(path, Method::POST, handler)
    }

    /// Registers a `PUT` handler for the given path.
    pub fn put<I, O>(
        self,
        path: &str,
        handler: impl EndpointHandler<I, S, Output = O, Future: Send> + Clone + Send + Sync + 'static,
    ) -> Self
    where
        I: FromRequest<Output = I> + Send + 'static,
        <I as FromRequest>::Error: Send,
        <I as FromRequest>::from_request(..): Send,
        O: IntoResponse + 'static,
    {
        self.insert(path, Method::PUT, handler)
    }

    /// Registers a `GET` handler for the given path.
    pub fn get<I, O>(
        self,
        path: &str,
        handler: impl EndpointHandler<I, S, Output = O, Future: Send> + Clone + Send + Sync + 'static,
    ) -> Self
    where
        I: FromRequest<Output = I> + Send + 'static,
        <I as FromRequest>::Error: Send,
        <I as FromRequest>::from_request(..): Send,
        O: IntoResponse + 'static,
    {
        self.insert(path, Method::GET, handler)
    }

    /// Registers a `PATCH` handler for the given path.
    pub fn patch<I, O>(
        self,
        path: &str,
        handler: impl EndpointHandler<I, S, Output = O, Future: Send> + Clone + Send + Sync + 'static,
    ) -> Self
    where
        I: FromRequest<Output = I> + Send + 'static,
        <I as FromRequest>::Error: Send,
        <I as FromRequest>::from_request(..): Send,
        O: IntoResponse + 'static,
    {
        self.insert(path, Method::PATCH, handler)
    }

    /// Registers a `DELETE` handler for the given path.
    pub fn delete<I, O>(
        self,
        path: &str,
        handler: impl EndpointHandler<I, S, Output = O, Future: Send> + Clone + Send + Sync + 'static,
    ) -> Self
    where
        I: FromRequest<Output = I> + Send + 'static,
        <I as FromRequest>::Error: Send,
        <I as FromRequest>::from_request(..): Send,
        O: IntoResponse + 'static,
    {
        self.insert(path, Method::DELETE, handler)
    }

    /// Registers a handler for a specific HTTP method and path.
    ///
    /// If the path has not been seen before, it is inserted into the
    /// path matcher and assigned a new route identifier.
    ///
    /// # Panics
    ///
    /// Panics if a handler is already registered for the given
    /// path and HTTP method.
    pub fn insert<I, O, H>(mut self, path: &str, method: Method, handler: H) -> Self
    where
        H: EndpointHandler<I, S, Output = O> + Clone + Send + Sync + 'static,
        H::Future: Send,
        I: FromRequest<Output = I> + Send + 'static,
        <I as FromRequest>::Error: Send,
        <I as FromRequest>::from_request(..): Send,
        O: IntoResponse + 'static,
    {
        let id = match self.path_to_id.get(path) {
            Some(existing_id) => existing_id.clone(),
            None => {
                let new_id = self.get_next_id();
                self.inner.insert(path, new_id.clone()).unwrap();
                self.path_to_id.insert(path.to_string(), new_id.clone());
                new_id
            }
        };

        // Wrap the endpoint handler in a method-specific adapter that:
        // - extracts request data,
        // - executes the handler,
        // - and normalizes output into a `Response`.
        let method_handler: Arc<MethodHandler<S>> = Arc::new(move |request, state| {
            let handler = Arc::new(handler.clone());
            let future = I::from_request(request);
            Box::pin(async move {
                match future.await {
                    Ok(input) => match handler.handle(input, state).await {
                        Continue(output) => Continue(output.into_response()),
                        Exit(exception) => Exit(exception.into_response()),
                    },
                    Err(_) => {
                        let mut response = Empty.into_response();
                        *response.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
                        Exit(response)
                    }
                }
            })
        });

        let endpoint = self.routes.entry(id).or_insert_with(|| Endpoint {
            methods: HashMap::new(),
        });

        if endpoint
            .methods
            .insert(method.clone(), method_handler)
            .is_some()
        {
            panic!("Route `{path}` already has handler for method `{method}`");
        }

        self
    }
}

/// [`Handler`] implementation for [`Router`].
///
/// This allows a router to be used anywhere a handler is expected,
/// including directly as the server handler.
impl<S> Handler<Request, S> for Router<S>
where
    S: Clone + Send + Sync + 'static,
{
    type Output = Response;
    type Exception = Response;
    type Future = impl Future<Output = Flow<Self::Output, Self::Exception>> + Send;

    // Attempt to match the request path and method.
    fn handle(&self, mut req: Request, state: S) -> Self::Future {
        let (method, path) = (req.method(), req.uri().path().to_string());

        let result = {
            let mut response = Response::new(Full::new(Bytes::new()));
            *response.status_mut() = StatusCode::NOT_FOUND;

            if let Ok(Match {
                value: route_id,
                params,
            }) = self.inner.at(&path)
            {
                if let Some(endpoint) = self.routes.get(route_id) {
                    if let Some(handler) = endpoint.methods.get(method) {
                        let params = params
                            .iter()
                            .map(|(k, v)| (k.to_string(), v.to_string()))
                            .collect::<HashMap<String, String>>();

                        Ok((handler.clone(), params))
                    } else {
                        *response.status_mut() = StatusCode::METHOD_NOT_ALLOWED;
                        Err(response)
                    }
                } else {
                    Err(response)
                }
            } else {
                Err(response)
            }
        };

        async move {
            match result {
                Ok((h, params)) => {
                    // Store route parameters in request extensions.
                    req.extensions_mut().insert(params);
                    h(req, state).await
                }
                Err(res) => Exit(res),
            }
        }
    }
}