h3x 0.2.0

use std::{
    collections::HashMap,
    sync::{Arc, RwLock},
    task::{Context, Poll},
};

use futures::future::BoxFuture;
use http::{Method, StatusCode};

use crate::server::{
    BoxService, BoxServiceFuture, IntoBoxService, MessageStreamError, Request, Response, Service,
    UnresolvedRequest, box_service,
};

#[tracing::instrument(skip_all)]
pub async fn default_fallback(_request: &mut Request, response: &mut Response) {
    tracing::debug!("call default fallback service (404 Not Found)");
    _ = response.set_status(StatusCode::NOT_FOUND)
}

#[derive(Debug, Clone)]
struct Fallback(Arc<RwLock<BoxService>>);

impl Fallback {
    pub fn new(service: BoxService) -> Self {
        Self(Arc::new(RwLock::new(service)))
    }

    pub fn set(&mut self, service: BoxService) {
        *self.0.write().expect("lock is not poisoned") = service;
    }
}

impl Service for Fallback {
    type Future<'s> = BoxServiceFuture<'s>;

    fn serve<'s>(&self, request: &'s mut Request, response: &'s mut Response) -> Self::Future<'s> {
        tracing::debug!("call fallback service");
        self.0
            .read()
            .expect("lock is not poisoned")
            .serve(request, response)
    }
}

#[derive(Debug, Clone)]
struct RouterInner {
    router: matchit::Router<BoxService>,
    fallback: Fallback,
}

impl Default for RouterInner {
    fn default() -> Self {
        Self {
            router: Default::default(),
            fallback: Fallback::new(box_service(default_fallback)),
        }
    }
}

impl RouterInner {
    fn route(&mut self, path: &str, service: impl IntoBoxService) {
        self.router
            .insert(path, service.into_box_service())
            .expect("failed to register route");
    }

    pub fn on(&mut self, method: Method, path: &str, service: impl IntoBoxService) {
        match self.router.at_mut(path) {
            Ok(exist_service) => {
                if let Some(router) = exist_service
                    .value
                    .downcast_mut::<MethodRouter<BoxService>>()
                {
                    router.set(method, service.into_box_service());
                } else {
                    let fallback = exist_service.value.clone();
                    let mut router = MethodRouter::new(fallback);
                    router.set(method, service.into_box_service());
                    *exist_service.value = router.into_box_service();
                }
            }
            Err(..) => {
                let mut router = MethodRouter::new(self.fallback.clone().into_box_service());
                router.set(method, service.into_box_service());
                self.route(path, router)
            }
        }
    }
}

impl Service for RouterInner {
    type Future<'s> = BoxServiceFuture<'s>;

    fn serve<'s>(
        &self,
        request: &'s mut Request,
        response: &'s mut Response,
    ) -> BoxServiceFuture<'s> {
        let Some(path_and_query) = request.path() else {
            tracing::debug!("missing path in request URI, call fallback service");
            return self.fallback.serve(request, response);
        };
        let path = path_and_query.path();
        let Ok(endpoint) = self.router.at(path) else {
            tracing::debug!(path, "path route: not found, call fallback service");
            return self.fallback.serve(request, response);
        };

        tracing::debug!(path, "path route found, call matched service");
        endpoint.value.serve(request, response)
    }
}

#[derive(Debug, Default, Clone)]
pub struct Router {
    inner: Arc<RouterInner>,
}

impl Router {
    pub fn new() -> Self {
        Self::default()
    }

    fn inner_ref(&self) -> &RouterInner {
        &self.inner
    }

    fn inner_mut(&mut self) -> &mut RouterInner {
        Arc::make_mut(&mut self.inner)
    }

    pub fn route(mut self, path: &str, service: impl IntoBoxService) -> Self {
        self.inner_mut().route(path, service.into_box_service());
        self
    }

    pub fn on(mut self, method: Method, path: &str, service: impl IntoBoxService) -> Self {
        self.inner_mut()
            .on(method, path, service.into_box_service());
        self
    }

    pub fn fallback(mut self, service: impl IntoBoxService) -> Self {
        self.inner_mut().fallback.set(service.into_box_service());
        self
    }

    pub fn options(self, path: &str, service: impl IntoBoxService) -> Self {
        self.on(Method::OPTIONS, path, service)
    }
    pub fn get(self, path: &str, service: impl IntoBoxService) -> Self {
        self.on(Method::GET, path, service)
    }
    pub fn post(self, path: &str, service: impl IntoBoxService) -> Self {
        self.on(Method::POST, path, service)
    }
    pub fn put(self, path: &str, service: impl IntoBoxService) -> Self {
        self.on(Method::PUT, path, service)
    }
    pub fn delete(self, path: &str, service: impl IntoBoxService) -> Self {
        self.on(Method::DELETE, path, service)
    }
    pub fn head(self, path: &str, service: impl IntoBoxService) -> Self {
        self.on(Method::HEAD, path, service)
    }
    pub fn trace(self, path: &str, service: impl IntoBoxService) -> Self {
        self.on(Method::TRACE, path, service)
    }
    pub fn connect(self, path: &str, service: impl IntoBoxService) -> Self {
        self.on(Method::CONNECT, path, service)
    }
    pub fn patch(self, path: &str, service: impl IntoBoxService) -> Self {
        self.on(Method::PATCH, path, service)
    }

    pub fn serve<'s>(
        &self,
        request: &'s mut Request,
        response: &'s mut Response,
    ) -> BoxServiceFuture<'s> {
        self.inner_ref().serve(request, response)
    }

    #[tracing::instrument(skip(self, req), fields(method = tracing::field::Empty, uri = tracing::field::Empty))]
    pub async fn handle(&self, req: UnresolvedRequest) -> Result<(), MessageStreamError> {
        let (mut request, mut response) = req.resolve().await?;

        tracing::Span::current()
            .record("method", request.method().as_str())
            .record("uri", request.uri().to_string());

        self.serve(&mut request, &mut response).await;

        // Drop response in place to avoid spawning another tokio task
        // FIXME: remove this when async drop is stabilized (https://github.com/rust-lang/rust/issues/126482)
        if let Some(drop_future) = response.drop() {
            drop_future.await;
        }

        Ok(())
    }
}

impl Service for Router {
    type Future<'s> = BoxServiceFuture<'s>;

    fn serve<'s>(&self, request: &'s mut Request, response: &'s mut Response) -> Self::Future<'s> {
        Router::serve(self, request, response)
    }
}

impl tower_service::Service<UnresolvedRequest> for Router {
    type Response = ();

    type Error = MessageStreamError;

    type Future = BoxFuture<'static, Result<(), MessageStreamError>>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        _ = cx;
        Poll::Ready(Ok(()))
    }

    fn call(&mut self, req: UnresolvedRequest) -> Self::Future {
        let router = self.clone();
        Box::pin(async move { router.handle(req).await })
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MethodRouter<S> {
    // most used methods are stored separately for faster access
    options: Option<S>,
    get: Option<S>,
    post: Option<S>,
    put: Option<S>,
    delete: Option<S>,
    head: Option<S>,
    trace: Option<S>,
    connect: Option<S>,
    patch: Option<S>,
    // other
    extensions: HashMap<Method, S>,
    // fallback service when no method match
    fallback: S,
}

impl<S> MethodRouter<S> {
    pub fn new(fallback: S) -> Self {
        Self {
            options: None,
            get: None,
            post: None,
            put: None,
            delete: None,
            head: None,
            trace: None,
            connect: None,
            patch: None,
            extensions: HashMap::new(),
            fallback,
        }
    }

    pub fn service(&self, method: Method) -> Option<&S> {
        match method {
            Method::OPTIONS => self.options.as_ref(),
            Method::GET => self.get.as_ref(),
            Method::POST => self.post.as_ref(),
            Method::PUT => self.put.as_ref(),
            Method::DELETE => self.delete.as_ref(),
            Method::HEAD => self.head.as_ref(),
            Method::TRACE => self.trace.as_ref(),
            Method::CONNECT => self.connect.as_ref(),
            Method::PATCH => self.patch.as_ref(),
            _ => self.extensions.get(&method),
        }
    }

    pub fn service_mut(&mut self, method: Method) -> Option<&mut S> {
        match method {
            Method::OPTIONS => self.options.as_mut(),
            Method::GET => self.get.as_mut(),
            Method::POST => self.post.as_mut(),
            Method::PUT => self.put.as_mut(),
            Method::DELETE => self.delete.as_mut(),
            Method::HEAD => self.head.as_mut(),
            Method::TRACE => self.trace.as_mut(),
            Method::CONNECT => self.connect.as_mut(),
            Method::PATCH => self.patch.as_mut(),
            _ => self.extensions.get_mut(&method),
        }
    }

    pub fn set(&mut self, method: Method, service: S) {
        match method {
            Method::OPTIONS => self.options = Some(service),
            Method::GET => self.get = Some(service),
            Method::POST => self.post = Some(service),
            Method::PUT => self.put = Some(service),
            Method::DELETE => self.delete = Some(service),
            Method::HEAD => self.head = Some(service),
            Method::TRACE => self.trace = Some(service),
            Method::CONNECT => self.connect = Some(service),
            Method::PATCH => self.patch = Some(service),
            _ => _ = self.extensions.insert(method, service),
        }
    }

    pub fn set_fallback(&mut self, service: S) {
        self.fallback = service;
    }
}

impl<S> Service for MethodRouter<S>
where
    S: Clone + for<'s> Service<Future<'s>: Send> + Send + 'static,
{
    type Future<'s> = BoxServiceFuture<'s>;

    fn serve<'s>(
        &self,
        request: &'s mut super::Request,
        response: &'s mut super::Response,
    ) -> Self::Future<'s> {
        let method = request.method();
        let service = match method {
            Method::OPTIONS => self.options.as_ref().unwrap_or(&self.fallback),
            Method::GET => self.get.as_ref().unwrap_or(&self.fallback),
            Method::POST => self.post.as_ref().unwrap_or(&self.fallback),
            Method::PUT => self.put.as_ref().unwrap_or(&self.fallback),
            Method::DELETE => self.delete.as_ref().unwrap_or(&self.fallback),
            Method::HEAD => self.head.as_ref().unwrap_or(&self.fallback),
            Method::TRACE => self.trace.as_ref().unwrap_or(&self.fallback),
            Method::CONNECT => self.connect.as_ref().unwrap_or(&self.fallback),
            Method::PATCH => self.patch.as_ref().unwrap_or(&self.fallback),
            _ => self.extensions.get(&method).unwrap_or(&self.fallback),
        }
        .clone();
        Box::pin(async move { service.serve(request, response).await })
    }
}

#[cfg(test)]
mod tests {
    use http::Method;

    use super::MethodRouter;

    fn make_router() -> MethodRouter<&'static str> {
        let mut router = MethodRouter::new("fallback");
        router.set(Method::GET, "get_handler");
        router.set(Method::POST, "post_handler");
        router.set(Method::PUT, "put_handler");
        router.set(Method::DELETE, "delete_handler");
        router
    }

    #[test]
    fn method_router_service_lookup() {
        let router = make_router();
        assert_eq!(router.service(Method::GET), Some(&"get_handler"));
        assert_eq!(router.service(Method::POST), Some(&"post_handler"));
        assert_eq!(router.service(Method::PUT), Some(&"put_handler"));
        assert_eq!(router.service(Method::DELETE), Some(&"delete_handler"));
    }

    #[test]
    fn method_router_unset_returns_none() {
        let router = make_router();
        assert_eq!(router.service(Method::PATCH), None);
        assert_eq!(router.service(Method::HEAD), None);
        assert_eq!(router.service(Method::OPTIONS), None);
        assert_eq!(router.service(Method::TRACE), None);
        assert_eq!(router.service(Method::CONNECT), None);
    }

    #[test]
    fn method_router_fallback() {
        let router = make_router();
        assert_eq!(router.fallback, "fallback");
    }

    #[test]
    fn method_router_set_all_standard_methods() {
        let mut router = MethodRouter::new("fb");
        router.set(Method::OPTIONS, "opt");
        router.set(Method::GET, "get");
        router.set(Method::POST, "post");
        router.set(Method::PUT, "put");
        router.set(Method::DELETE, "del");
        router.set(Method::HEAD, "head");
        router.set(Method::TRACE, "trace");
        router.set(Method::CONNECT, "connect");
        router.set(Method::PATCH, "patch");

        assert_eq!(router.service(Method::OPTIONS), Some(&"opt"));
        assert_eq!(router.service(Method::GET), Some(&"get"));
        assert_eq!(router.service(Method::POST), Some(&"post"));
        assert_eq!(router.service(Method::PUT), Some(&"put"));
        assert_eq!(router.service(Method::DELETE), Some(&"del"));
        assert_eq!(router.service(Method::HEAD), Some(&"head"));
        assert_eq!(router.service(Method::TRACE), Some(&"trace"));
        assert_eq!(router.service(Method::CONNECT), Some(&"connect"));
        assert_eq!(router.service(Method::PATCH), Some(&"patch"));
    }

    #[test]
    fn method_router_service_mut() {
        let mut router = make_router();
        if let Some(handler) = router.service_mut(Method::GET) {
            *handler = "updated_get";
        }
        assert_eq!(router.service(Method::GET), Some(&"updated_get"));
    }

    #[test]
    fn method_router_set_fallback() {
        let mut router = make_router();
        router.set_fallback("new_fallback");
        assert_eq!(router.fallback, "new_fallback");
    }

    #[test]
    fn method_router_overwrite() {
        let mut router = make_router();
        router.set(Method::GET, "overwritten");
        assert_eq!(router.service(Method::GET), Some(&"overwritten"));
    }

    #[test]
    fn router_builder_chain() {
        use super::Router;

        async fn dummy(_req: &mut super::Request, _resp: &mut super::Response) {}

        // Just test that the builder pattern compiles and doesn't panic
        let _router = Router::new()
            .route("/exact", dummy)
            .get("/api/users", dummy)
            .post("/api/users", dummy)
            .fallback(dummy);
    }

    #[test]
    fn matchit_path_matching() {
        // Directly test the underlying matchit router to verify path matching logic
        let mut router = matchit::Router::new();
        router.insert("/", "root").unwrap();
        router.insert("/users", "users").unwrap();
        router.insert("/users/{id}", "user_by_id").unwrap();
        router.insert("/files/{*path}", "files_catch_all").unwrap();

        // Exact matches
        assert_eq!(*router.at("/").unwrap().value, "root");
        assert_eq!(*router.at("/users").unwrap().value, "users");

        // Parameterized match
        let m = router.at("/users/42").unwrap();
        assert_eq!(*m.value, "user_by_id");
        assert_eq!(m.params.get("id"), Some("42"));

        // Catch-all match
        let m = router.at("/files/docs/readme.md").unwrap();
        assert_eq!(*m.value, "files_catch_all");
        assert_eq!(m.params.get("path"), Some("docs/readme.md"));

        // No match
        assert!(router.at("/nonexistent").is_err());
        assert!(router.at("").is_err());
    }
}