use std::future::Future;
use std::pin::Pin;

use axum::body::BodyDataStream;
use axum::response::IntoResponse;
use tower::Layer;
use tower_service::Service;
use tracing::info;

#[derive(Default, Clone, Copy)]
pub struct LambdaLayer;

impl<S> Layer<S> for LambdaLayer {
    type Service = LambdaService<S>;

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

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

impl<S> Service<lambda_http::Request> for LambdaService<S>
where
    S: Service<axum::http::Request<axum::body::Body>>,
    S::Response: axum::response::IntoResponse + Send + 'static,
    S::Error: std::error::Error + Send + Sync + 'static,
    S::Future: Send + 'static,
{
    type Response = lambda_http::Response<BodyDataStream>;
    type Error = lambda_http::Error;
    type Future =
        Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send + 'static>>;

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

    fn call(&mut self, req: lambda_http::Request) -> Self::Future {
        let (parts, body) = req.into_parts();
        let body = match body {
            lambda_http::Body::Empty => axum::body::Body::default(),
            lambda_http::Body::Text(t) => t.into(),
            lambda_http::Body::Binary(v) => v.into(),
            _ => todo!(),
        };

        let request = axum::http::Request::from_parts(parts, body);
        info!("Received request: {request:?}");
        let fut = self.inner.call(request);

        let fut = async move {
            let resp = fut.await?;
            let (parts, body) = resp.into_response().into_parts();
            // We support streaming responses with this lambda
            let stream = body.into_data_stream();
            let resp: hyper::Response<BodyDataStream> = hyper::Response::from_parts(parts, stream);
            info!("Generated response: {resp:?}");
            Ok(resp)
        };
        Box::pin(fut)
    }
}