jelq

use std::{collections::HashMap, fmt::Write as _, time::Instant};

use axum::{

    body::Bytes,

    extract::{FromRequest, FromRequestParts, Query, RawPathParams, Request},

    http::{HeaderMap, HeaderName, Method, StatusCode, Uri, Version, header},

    response::{IntoResponse, Response},

};

use hyper::upgrade::OnUpgrade;

use crate::problem::Problem;

pub const BODY_LIMIT: usize = 64 * 1024;

/// Everything a handler may look at, gathered once per request.

pub struct Ctx {

    method: Method,

    uri: Uri,

    version: Version,

    headers: HeaderMap,

    query: HashMap<String, String>,

    params: HashMap<String, String>,

    body: Bytes,

    now: Instant,

    upgrade: Option<OnUpgrade>,

}

impl Ctx {

    pub fn header(&self, name: HeaderName) -> Option<&str> {

        self.headers.get(name).and_then(|v| v.to_str().ok())

    }

    pub fn cookie(&self, name: &str) -> Option<&str> {

        self.headers

            .get_all(header::COOKIE)

            .iter()

            .filter_map(|v| v.to_str().ok())

            .flat_map(|v| v.split(';'))

            .filter_map(|pair| pair.trim().split_once('='))

            .find(|(key, _)| *key == name)

            .map(|(_, value)| value)

    }

    pub fn param(&self, name: &str) -> &str {

        self.params.get(name).map_or("", String::as_str)

    }

    pub fn query(&self, name: &str) -> Option<&str> {

        self.query.get(name).map(String::as_str)

    }

    pub fn method(&self) -> &Method {

        &self.method

    }

    pub fn body(&self) -> &[u8] {

        &self.body

    }

    pub fn now(&self) -> Instant {

        self.now

    }

    pub fn take_upgrade(&mut self) -> Option<OnUpgrade> {

        self.upgrade.take()

    }

    /// The request as it arrived, which is what TRACE sends back.

    pub fn echo(&self) -> String {

        let mut message = format!("{} {} {:?}\r\n", self.method, self.uri, self.version);

        for (name, value) in &self.headers {

            let _ = write!(

                message,

                "{name}: {}\r\n",

                value.to_str().unwrap_or("<binary>")

            );

        }

        message.push_str("\r\n");

        message

    }

}

impl<S: Send + Sync> FromRequest<S> for Ctx {

    type Rejection = Response;

    async fn from_request(mut request: Request, state: &S) -> Result<Self, Self::Rejection> {

        let upgrade = request.extensions_mut().remove::<OnUpgrade>();

        let (mut parts, body) = request.into_parts();

        let params = RawPathParams::from_request_parts(&mut parts, state)

            .await

            .map(|raw| {

                raw.iter()

                    .map(|(k, v)| (k.to_owned(), v.to_owned()))

                    .collect()

            })

            .unwrap_or_default();

        let query = Query::<HashMap<String, String>>::try_from_uri(&parts.uri)

            .map(|q| q.0)

            .unwrap_or_default();

        let body = axum::body::to_bytes(body, BODY_LIMIT).await.map_err(|_| {

            Problem::new(

                StatusCode::PAYLOAD_TOO_LARGE,

                "body-too-large",

                "The request body is larger than anything the server is prepared to receive.",

            )

            .into_response()

        })?;

        Ok(Self {

            method: parts.method,

            uri: parts.uri,

            version: parts.version,

            headers: parts.headers,

            query,

            params,

            body,

            now: Instant::now(),

            upgrade,

        })

    }

}