#![deny(unused, missing_docs)]
//! This is an actix-web middleware to associate every request with a unique ID. This ID can be
//! used to track errors in an application.

use std::{
    fmt,
    pin::Pin,
    task::{Context, Poll},
};

use actix_web::{
    dev::{Payload, Service, ServiceRequest, ServiceResponse, Transform},
    error::ResponseError,
    http::header::{HeaderName, HeaderValue},
    Error as ActixError, FromRequest, HttpMessage, HttpRequest,
};
use futures::{
    future::{ok, ready, Ready},
    Future,
};
use uuid::Uuid;

/// The default header used for the request ID.
pub const DEFAULT_HEADER: &str = "x-request-id";

/// Possible error types for the middleware.
#[derive(Debug, Clone)]
pub enum Error {
    /// There is no ID associated with this request.
    NoAssociatedId,
}

/// ID wrapper for requests.
pub struct RequestIdMiddleware<S> {
    service: S,
    header_name: HeaderName,
    id: RequestId,
}

type Generator = fn() -> HeaderValue;

/// A middleware for generating per-request unique IDs
pub struct RequestIdentifier {
    header_name: &'static str,
    id_generator: Generator,
}

/// Request ID that can be extracted in handlers.
#[derive(Clone)]
pub struct RequestId(HeaderValue);

impl ResponseError for Error {}

impl fmt::Display for Error {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        use Error::*;
        match self {
            NoAssociatedId => write!(fmt, "NoAssociatedId"),
        }
    }
}

impl RequestId {
    fn header_value(&self) -> &HeaderValue {
        &self.0
    }

    /// Get a string representation of this ID
    pub fn as_str(&self) -> &str {
        self.0.to_str().expect("Non-ASCII IDs are not supported")
    }
}

impl RequestIdentifier {
    /// Create a default middleware using [`DEFAULT_HEADER`](./constant.DEFAULT_HEADER.html) as the header name and
    /// UUID v4 for ID generation.
    pub fn with_uuid() -> Self {
        Default::default()
    }

    /// Create a middleware using a custom header name and UUID v4 for ID generation.
    pub fn with_header(header_name: &'static str) -> Self {
        Self {
            header_name,
            id_generator: default_generator,
        }
    }

    /// Change the header name for this middleware.
    pub fn header(self, header_name: &'static str) -> Self {
        Self {
            header_name,
            ..self
        }
    }

    /// Create a middleware using [`DEFAULT_HEADER`](./constant.DEFAULT_HEADER.html) as the header
    /// name and IDs as generated by `id_generator`. `id_generator` should return a unique ID
    /// (ASCII only), each time it is invoked.
    pub fn with_generator(id_generator: Generator) -> Self {
        Self {
            header_name: DEFAULT_HEADER,
            id_generator,
        }
    }

    /// Change the ID generator for this middleware.
    pub fn generator(self, id_generator: Generator) -> Self {
        Self {
            id_generator,
            ..self
        }
    }
}

impl Default for RequestIdentifier {
    fn default() -> Self {
        Self {
            header_name: DEFAULT_HEADER,
            id_generator: default_generator,
        }
    }
}

/// Default UUID v4 based ID generator.
fn default_generator() -> HeaderValue {
    let uuid = Uuid::new_v4();
    HeaderValue::from_str(&uuid.to_hyphenated().to_string())
        // This unwrap can never fail since UUID v4 generated IDs are ASCII-only
        .unwrap()
}

impl<S, B> Transform<S> for RequestIdentifier
where
    S: Service<Request = ServiceRequest, Response = ServiceResponse<B>, Error = ActixError>,
    S::Future: 'static,
    B: 'static,
{
    type Request = S::Request;
    type Response = S::Response;
    type Error = S::Error;
    type InitError = ();
    type Transform = RequestIdMiddleware<S>;
    type Future = Ready<Result<Self::Transform, Self::InitError>>;

    fn new_transform(&self, service: S) -> Self::Future {
        // generate a new ID
        let value = (self.id_generator)();

        // associate with the request
        ok(RequestIdMiddleware {
            service,
            header_name: HeaderName::from_static(self.header_name),
            id: RequestId(value),
        })
    }
}

#[allow(clippy::type_complexity)]
impl<S, B> Service for RequestIdMiddleware<S>
where
    S: Service<Request = ServiceRequest, Response = ServiceResponse<B>, Error = ActixError>,
    S::Future: 'static,
    B: 'static,
{
    type Request = S::Request;
    type Response = S::Response;
    type Error = S::Error;
    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>>>>;

    fn poll_ready(&mut self, ctx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.service.poll_ready(ctx)
    }

    fn call(&mut self, req: ServiceRequest) -> Self::Future {
        req.extensions_mut().insert(self.id.clone());
        let fut = self.service.call(req);
        // TODO: clone needed?
        let header_name = self.header_name.clone();
        let header_value = self.id.header_value().clone();

        Box::pin(async move {
            let mut res = fut.await?;

            res.headers_mut().insert(header_name, header_value);

            Ok(res)
        })
    }
}

impl FromRequest for RequestId {
    type Error = Error;
    type Future = Ready<Result<Self, Self::Error>>;
    type Config = ();

    fn from_request(req: &HttpRequest, _: &mut Payload) -> Self::Future {
        ready(
            req.extensions()
                .get::<RequestId>()
                .map(RequestId::clone)
                .ok_or(Error::NoAssociatedId),
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use actix_web::{test, web, App};
    use bytes::{Buf, Bytes};

    async fn handler(id: RequestId) -> String {
        id.as_str().to_string()
    }

    async fn test_get(middleware: RequestIdentifier) -> ServiceResponse {
        let mut service = test::init_service(
            App::new()
                .wrap(middleware)
                .route("/", web::get().to(handler)),
        )
        .await;
        test::call_service(&mut service, test::TestRequest::get().uri("/").to_request()).await
    }

    #[actix_rt::test]
    async fn default_identifier() {
        let resp = test_get(RequestIdentifier::with_uuid()).await;
        let uid = resp
            .headers()
            .get(HeaderName::from_static(DEFAULT_HEADER))
            .map(|v| v.to_str().unwrap().to_string())
            .unwrap();
        let body: Bytes = test::read_body(resp).await;
        let body = String::from_utf8_lossy(body.bytes());
        assert_eq!(uid, body);
    }

    #[actix_rt::test]
    async fn deterministic_identifier() {
        let resp = test_get(RequestIdentifier::with_generator(|| {
            HeaderValue::from_static("look ma, i'm an id")
        }))
        .await;
        let uid = resp
            .headers()
            .get(HeaderName::from_static(DEFAULT_HEADER))
            .map(|v| v.to_str().unwrap().to_string())
            .unwrap();
        let body: Bytes = test::read_body(resp).await;
        let body = String::from_utf8_lossy(body.bytes());
        assert_eq!(uid, body);
    }

    #[actix_rt::test]
    async fn custom_header() {
        let resp = test_get(RequestIdentifier::with_header("custom-header")).await;
        assert!(resp
            .headers()
            .get(HeaderName::from_static(DEFAULT_HEADER))
            .is_none());
        let uid = resp
            .headers()
            .get(HeaderName::from_static("custom-header"))
            .map(|v| v.to_str().unwrap().to_string())
            .unwrap();
        let body: Bytes = test::read_body(resp).await;
        let body = String::from_utf8_lossy(body.bytes());
        assert_eq!(uid, body);
    }
}