use std::task::Context;

use http::{Request, StatusCode};
use http_body::Body;
use http_body_util::{BodyExt as _, LengthLimitError, Limited};
use tower_layer::Layer;
use tower_service::Service;

use crate::response::{
    Response as Res,
    json::{CreateJsonResponse, JsonResponseError, ResponseError},
};

/// Default maximum body size in bytes.
pub const REQUEST_BODY_LIMIT_DEFAULT: usize = 10 * 1024 * 1024;

#[derive(Debug, Clone, Copy)]
pub struct RequestBodyLimitService<S> {
    inner: S,
    limit: usize,
}

impl<B, S> Service<Request<B>> for RequestBodyLimitService<S>
where
    S: Service<Request<B>, Response = Res> + Clone + Send + 'static,
    S::Future: Send + 'static,
    B: Body + Send + From<bytes::Bytes> + 'static,
    B::Data: Send + 'static,
    B::Error: std::error::Error + Send + Sync + 'static,
{
    type Response = S::Response;
    type Error = S::Error;
    type Future = std::pin::Pin<
        Box<
            dyn std::future::Future<
                    Output = Result<Self::Response, Self::Error>,
                > + Send,
        >,
    >;

    fn poll_ready(
        &mut self,
        cx: &mut Context<'_>,
    ) -> std::task::Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(
        &mut self,
        req: Request<B>,
    ) -> Self::Future {
        let (parts, body) = req.into_parts();

        let limited_body: Limited<B> = Limited::new(body, self.limit);

        let mut inner: S = self.inner.clone();

        Box::pin(async move {
            match limited_body.collect().await {
                | Ok(collected) => {
                    let new_body: B = collected.to_bytes().into();

                    let req: Request<B> = Request::from_parts(parts, new_body);

                    inner.call(req).await
                },
                | Err(err)
                    if err.downcast_ref::<LengthLimitError>().is_some() =>
                {
                    let rer: ResponseError = ResponseError::TooLarge;

                    let res: Res = CreateJsonResponse::failure()
                        .status(StatusCode::PAYLOAD_TOO_LARGE)
                        .add_error(
                            JsonResponseError::new()
                                .code(rer.as_code())
                                .path(["request", "body"])
                                .message(rer.as_message()),
                        )
                        .create();

                    Ok(res)
                },
                | Err(_) => {
                    let rer: ResponseError = ResponseError::Parse;

                    let res: Res = CreateJsonResponse::failure()
                        .status(StatusCode::BAD_REQUEST)
                        .add_error(
                            JsonResponseError::new()
                                .code(rer.as_code())
                                .path(["request", "body"])
                                .message(rer.as_message()),
                        )
                        .create();

                    Ok(res)
                },
            }
        })
    }
}

/// Layer for configuring the request body limit.
///
/// Following error will be returned if the request body exceeds the limit:
///
/// ```jsonc
/// // Status: 413
/// {
///     "success": false,
///     "data": null,
///     "errors": [
///         {
///             "code": "too_large",
///             "path": [
///                 "request",
///                 "body"
///             ],
///             "message": "Request body is too large"
///         }
///     ]
/// }
/// ```
///
/// ## Examples
///
/// Create a a layer with default limit:
///
/// ```no_run
/// use axum::{
///     Router,
///     extract::DefaultBodyLimit,
/// };
/// use jder_axum::layers::RequestBodyLimit;
///
/// let app: Router = Router::new()
///     .layer(DefaultBodyLimit::disable())
///     .layer(RequestBodyLimit::new());
/// ```
///
/// Create a a layer with custom limit:
///
/// ```no_run
/// use axum::{
///     Router,
///     extract::DefaultBodyLimit,
/// };
/// use jder_axum::layers::RequestBodyLimit;
///
/// let app: Router = Router::new()
///     .layer(DefaultBodyLimit::disable())
///     .layer(RequestBodyLimit::max(10 * 1024 * 1024));
/// ```
#[derive(Debug, Clone, Copy)]
pub struct RequestBodyLimit {
    limit: usize,
}

impl RequestBodyLimit {
    /// Create a new `RequestBodyLimit` layer
    /// with [REQUEST_BODY_LIMIT_DEFAULT].
    ///
    /// ## Example
    ///
    /// ```no_run
    /// use jder_axum::layers::RequestBodyLimit;
    ///
    /// RequestBodyLimit::new();
    /// ```
    pub fn new() -> Self {
        Self { limit: REQUEST_BODY_LIMIT_DEFAULT }
    }

    /// Set the request body limit in bytes.
    ///
    /// ## Example
    ///
    /// ```no_run
    /// use jder_axum::layers::RequestBodyLimit;
    ///
    /// RequestBodyLimit::max(10 * 1024 * 1024); // 10MiB
    /// ```
    pub fn max(limit: usize) -> Self {
        Self { limit }
    }
}

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

impl<S> Layer<S> for RequestBodyLimit {
    type Service = RequestBodyLimitService<S>;

    fn layer(
        &self,
        inner: S,
    ) -> Self::Service {
        RequestBodyLimitService { inner, limit: self.limit }
    }
}
