use fred::prelude::FredResult;
use http::{Response, StatusCode, header::RETRY_AFTER};
use pin_project::pin_project;
use std::{
    pin::Pin,
    task::{Context, Poll, ready},
};
#[cfg(feature = "tracing")]
use tracing::{Span, error, info};

type RateLimitStatus = (bool, u32); // allowed, retry_after
type CheckFuture = Pin<Box<dyn Future<Output = FredResult<RateLimitStatus>> + Send + 'static>>;

#[pin_project(project = ResponseFutureProj)]
pub enum ResponseFuture<F, ErrorBody> {
    Pending {
        #[pin]
        inner: F,
        #[pin]
        check_future: CheckFuture,
        check_result: Option<RateLimitStatus>,
        fail_open: bool,
        #[cfg(feature = "tracing")]
        span: Span,
    },
    Ready(Option<Response<ErrorBody>>),
    Passthrough(#[pin] F),
}

impl<F, ErrorBody> ResponseFuture<F, ErrorBody> {
    pub(crate) fn ready(inner: Response<ErrorBody>) -> Self {
        Self::Ready(Some(inner))
    }

    pub(crate) fn passthrough(inner: F) -> Self {
        Self::Passthrough(inner)
    }

    pub(crate) fn pending(
        inner: F,
        check_future: CheckFuture,
        fail_open: bool,
        #[cfg(feature = "tracing")] span: Span,
    ) -> Self {
        Self::Pending {
            inner,
            check_future,
            check_result: None,
            fail_open,
            #[cfg(feature = "tracing")]
            span,
        }
    }
}

impl<F, B, E> Future for ResponseFuture<F, B>
where
    F: Future<Output = Result<Response<B>, E>>,
    B: Default,
{
    type Output = Result<Response<B>, E>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.project();

        match this {
            ResponseFutureProj::Passthrough(inner) => inner.poll(cx),
            ResponseFutureProj::Ready(inner) => Poll::Ready(Ok(inner
                .take()
                .expect("Completed future should not be polled again"))),
            ResponseFutureProj::Pending {
                inner,
                check_future,
                check_result,
                fail_open,
                #[cfg(feature = "tracing")]
                span,
            } => {
                if check_result.is_none() {
                    match ready!(check_future.poll(cx)) {
                        Ok(res) => {
                            *check_result = Some(res);
                        }
                        Err(e) => {
                            #[cfg(not(feature = "tracing"))]
                            let _ = e;
                            #[cfg(feature = "tracing")]
                            span.in_scope(
                                || error!(error = %e, "failed to check rate limit status"),
                            );
                            if *fail_open {
                                #[cfg(feature = "tracing")]
                                span.in_scope(|| info!("failing open"));
                                *check_result = Some((true, 0))
                            } else {
                                return Poll::Ready(Ok(respond(StatusCode::INTERNAL_SERVER_ERROR)));
                            }
                        }
                    }
                }

                let (allowed, retry_after) =
                    check_result.expect("check_result is None after check_future completed");
                if allowed {
                    inner.poll(cx)
                } else {
                    #[cfg(feature = "tracing")]
                    span.in_scope(|| info!(retry_after, "rate limit exceeded"));
                    let mut resp = respond(StatusCode::TOO_MANY_REQUESTS);
                    resp.headers_mut().insert(RETRY_AFTER, retry_after.into());
                    Poll::Ready(Ok(resp))
                }
            }
        }
    }
}

fn respond<B>(status_code: StatusCode) -> Response<B>
where
    B: Default,
{
    let mut res = Response::new(B::default());
    *res.status_mut() = status_code;
    res
}
