use axum::{
    extract::{Extension, FromRequestParts, Request},
    http::{request::Parts, StatusCode},
    middleware::Next,
    response::Response,
};
use chrono::Utc;

pub type DateTime = chrono::DateTime<chrono::Utc>;

// Extractor that provides the "current time" for a request. This time is
// calculated once per request, even if the extractor is used in multiple
// places. This requires the [middleware] function to be installed with
// [axum::middleware::from_fn] around the current route.
#[derive(Clone)]
pub struct RequestedAt(pub DateTime);

impl RequestedAt {
    fn now() -> Self {
        Self(Utc::now())
    }
}

#[async_trait::async_trait]
impl<S> FromRequestParts<S> for RequestedAt
where
    S: Send + Sync,
{
    type Rejection = <Extension<Self> as FromRequestParts<S>>::Rejection;

    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
        // This is purely for ergonomics: it allows `RequestedAt` to be extracted
        // without having to wrap it in `Extension<>`. Callers _can_ still do that,
        // but they aren't forced to.
        let Extension(requested_at) = Extension::<Self>::from_request_parts(parts, state).await?;

        Ok(requested_at)
    }
}

impl From<DateTime> for RequestedAt {
    fn from(timestamp: DateTime) -> Self {
        Self(timestamp)
    }
}

impl std::ops::Deref for RequestedAt {
    type Target = DateTime;

    fn deref(&self) -> &Self::Target {
        let Self(timestamp) = self;
        timestamp
    }
}

// Computes a canonical "requested at" time for each request it wraps. This
// time can be recovered using the [RequestedAt] extractor.
pub async fn middleware(mut req: Request, next: Next) -> Result<Response, StatusCode> {
    let now = RequestedAt::now();
    req.extensions_mut().insert(now);

    Ok(next.run(req).await)
}