use std::borrow::Cow;

use hyper::{Request, Response, Uri, body::Incoming, service::Service};
use nailbody::Body;

pub fn trim_trailing_slash(uri: &mut Uri) {
    if !uri.path().ends_with('/') && !uri.path().starts_with("//") {
        return;
    }

    let new_path = format!("/{}", uri.path().trim_matches('/'));

    let mut parts = uri.clone().into_parts();

    let new_path_and_query = if let Some(path_and_query) = &parts.path_and_query {
        let new_path_and_query = if let Some(query) = path_and_query.query() {
            Cow::Owned(format!("{}?{}", new_path, query))
        } else {
            new_path.into()
        }
        .parse()
        .unwrap();

        Some(new_path_and_query)
    } else {
        None
    };

    parts.path_and_query = new_path_and_query;
    if let Ok(new_uri) = Uri::from_parts(parts) {
        *uri = new_uri;
    }
}

pub struct NormalizeUrls<S> {
    inner: S,
}

impl<S> NormalizeUrls<S> {
    pub fn new(inner: S) -> Self {
        Self { inner }
    }
}

impl<S> Service<Request<Incoming>> for NormalizeUrls<S>
where
    S: Service<Request<Incoming>, Response = Response<Body>>,
{
    type Response = S::Response;

    type Error = S::Error;

    type Future = S::Future;

    fn call(&self, mut req: Request<Incoming>) -> Self::Future {
        trim_trailing_slash(req.uri_mut());

        self.inner.call(req)
    }
}