use axum::{
    extract::{Extension, FromRequestParts, Request},
    http::{request::Parts, StatusCode},
    middleware::Next,
    response::Response,
};
use chrono::{DateTime, 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<Utc>);

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

    pub fn timestamp(&self) -> DateTime<Utc> {
        self.0
    }
}

#[async_trait::async_trait]
impl<S> FromRequestParts<S> for RequestedAt
where
    S: Send + Sync,
{
    type Rejection = <Extension<RequestedAt> 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::<RequestedAt>::from_request_parts(parts, state).await?;

        Ok(requested_at)
    }
}

/// 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)
}