use crate::{KeyResolver, config::Config, future::ResponseFuture, rate_limit};
use fred::prelude::LuaInterface;
use http::{Request, Response};
use std::task::{Context, Poll};
use tower::Service;
#[cfg(feature = "tracing")]
use tracing::{Instrument, field::Empty, info_span};

#[derive(Debug, Clone)]
pub struct RateLimit<S, L, K>
where
    L: LuaInterface,
    K: KeyResolver,
{
    inner: S,
    config: Config<L, K>,
}

impl<S, L, K> RateLimit<S, L, K>
where
    L: LuaInterface,
    K: KeyResolver,
{
    pub(crate) fn new(inner: S, config: Config<L, K>) -> Self {
        RateLimit { inner, config }
    }
}

impl<S, L, K, ReqBody, ResBody> Service<Request<ReqBody>> for RateLimit<S, L, K>
where
    S: Service<Request<ReqBody>, Response = Response<ResBody>>,
    L: LuaInterface + 'static,
    K: KeyResolver<ErrorBody = ResBody>,
    ResBody: Default,
{
    type Response = S::Response;
    type Error = S::Error;
    type Future = ResponseFuture<S::Future, K::ErrorBody>;

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

    fn call(&mut self, req: Request<ReqBody>) -> Self::Future {
        #[cfg(feature = "tracing")]
        let span = info_span!("rate limit check", key = Empty);

        let key = {
            #[cfg(feature = "tracing")]
            let _span_guard = span.enter();
            match self.config.key_resolver.key(&req) {
                Ok(Some(key)) => key,
                Ok(None) => return ResponseFuture::passthrough(self.inner.call(req)),
                Err(res) => return ResponseFuture::ready(res),
            }
        };

        #[cfg(feature = "tracing")]
        match key.as_str() {
            Some(s) => span.record("key", s),
            None => span.record("key", key.as_bytes()),
        };

        let check_future = {
            let conn = self.config.conn.clone();
            let key = key.clone();
            let capacity = self.config.capacity;
            let emission_interval = self.config.emission_interval;

            let fut = async move {
                rate_limit::run_script(&conn, key, capacity, emission_interval).await
            };

            #[cfg(not(feature = "tracing"))]
            {
                Box::pin(fut)
            }

            #[cfg(feature = "tracing")]
            Box::pin(fut.instrument(span.clone().or_current()))
        };

        #[cfg(not(feature = "tracing"))]
        return ResponseFuture::pending(self.inner.call(req), check_future, self.config.fail_open);

        #[cfg(feature = "tracing")]
        ResponseFuture::pending(
            self.inner.call(req),
            check_future,
            self.config.fail_open,
            span,
        )
    }
}
