use axum::{
    async_trait,
    extract::FromRequestParts,
    http::{
        header::{AUTHORIZATION, SET_COOKIE},
        request::Parts,
        HeaderValue,
    },
    response::{IntoResponse, IntoResponseParts},
    RequestPartsExt,
};
use axum_extra::{
    extract::{cookie::Cookie, CookieJar},
    headers::{authorization::Bearer, Authorization},
    TypedHeader,
};
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use time::{Duration, OffsetDateTime};
use uuid::Uuid;

use super::{Error, JWT};

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(remote = "Self")]
pub struct Claims<const LIFETIME: i64 = ACCESS> {
    pub sub: Uuid,
    #[serde(with = "numeric_date")]
    iat: OffsetDateTime,
    #[serde(with = "numeric_date")]
    exp: OffsetDateTime,
    jti: Uuid,
}

impl<const LIFETIME: i64> Claims<LIFETIME> {
    pub fn new(sub: Uuid, iat: OffsetDateTime) -> Self {
        let iat = iat
            .replace_millisecond(0)
            .expect("Failed to remove millisecond from datetime. This should not have happened.");

        let exp = iat + Duration::new(LIFETIME, 0);
        let jti = uuid::Uuid::new_v4();

        Self { sub, iat, exp, jti }
    }

    pub fn issue(uuid: Uuid) -> Self {
        Self::new(uuid, OffsetDateTime::now_utc())
    }

    pub fn expired(&self) -> bool {
        self.exp > OffsetDateTime::now_utc()
    }
}

impl<const LIFETIME: i64> Serialize for Claims<LIFETIME> {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        Self::serialize(self, serializer)
    }
}

impl<'de, const LIFETIME: i64> Deserialize<'de> for Claims<LIFETIME> {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let claims = Self::deserialize(deserializer)?;

        if claims.exp - claims.iat != Duration::new(LIFETIME, 0) {
            return Err(serde::de::Error::custom(
                "Lifetime is invalid for Claim type",
            ));
        }

        Ok(claims)
    }
}

mod numeric_date {
    //! Custom serialization of OffsetDateTime to conform with the JWT spec (RFC 7519 section 2, "Numeric Date")
    use serde::{self, Deserialize, Deserializer, Serializer};
    use time::OffsetDateTime;

    /// Serializes an OffsetDateTime to a Unix timestamp (milliseconds since 1970/1/1T00:00:00T)
    pub fn serialize<S>(date: &OffsetDateTime, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        serializer.serialize_i64(date.unix_timestamp())
    }

    /// Attempts to deserialize an i64 and use as a Unix timestamp
    pub fn deserialize<'de, D>(deserializer: D) -> Result<OffsetDateTime, D::Error>
    where
        D: Deserializer<'de>,
    {
        OffsetDateTime::from_unix_timestamp(i64::deserialize(deserializer)?)
            .map_err(|_| serde::de::Error::custom("invalid Unix timestamp value"))
    }
}

// 1 day in seconds
const ACCESS: i64 = 86400;

pub type AccessClaims = Claims<ACCESS>;

impl From<RefreshClaims> for AccessClaims {
    fn from(value: RefreshClaims) -> Self {
        Claims::issue(value.sub)
    }
}

impl TryFrom<AccessClaims> for Cookie<'_> {
    type Error = Error;

    fn try_from(value: AccessClaims) -> Result<Self, Self::Error> {
        Ok(Cookie::build(("token", JWT.encode(&value)?))
            .expires(value.exp)
            .secure(true)
            .http_only(true)
            .build())
    }
}

impl TryFrom<AccessClaims> for HeaderValue {
    type Error = Error;

    fn try_from(value: AccessClaims) -> Result<Self, Self::Error> {
        Cookie::try_from(value)?
            .to_string()
            .parse()
            .map_err(Into::into)
    }
}

impl IntoResponse for AccessClaims {
    fn into_response(self) -> axum::response::Response {
        (self, ()).into_response()
    }
}

impl IntoResponseParts for AccessClaims {
    type Error = Error;

    fn into_response_parts(
        self,
        mut res: axum::response::ResponseParts,
    ) -> Result<axum::response::ResponseParts, Self::Error> {
        res.headers_mut()
            .append(SET_COOKIE, HeaderValue::try_from(self)?);

        Ok(res)
    }
}

#[async_trait]
impl<S> FromRequestParts<S> for AccessClaims
where
    S: Send + Sync,
{
    type Rejection = Error;

    async fn from_request_parts(parts: &mut Parts, _: &S) -> Result<Self, Self::Rejection> {
        let jar = parts
            .extract::<CookieJar>()
            .await
            .expect("Infallable result was in fact, fallable");

        JWT.decode(jar.get("token").ok_or(Error::JwtNotFound)?.value())
            .map(|t| t.claims)
    }
}

// 30 days in seconds
const REFRESH: i64 = 2_592_000;

pub type RefreshClaims = Claims<REFRESH>;

impl RefreshClaims {
    pub fn refresh(self) -> AccessClaims {
        self.into()
    }
}

impl IntoResponseParts for RefreshClaims {
    type Error = Error;

    fn into_response_parts(
        self,
        mut res: axum::response::ResponseParts,
    ) -> Result<axum::response::ResponseParts, Self::Error> {
        res.headers_mut().append(
            AUTHORIZATION,
            HeaderValue::try_from(format!("Bearer {}", JWT.encode(&self)?))?,
        );

        Ok(res)
    }
}

impl IntoResponse for RefreshClaims {
    fn into_response(self) -> axum::response::Response {
        JWT.encode(&self).into_response()
    }
}

#[async_trait]
impl<S> FromRequestParts<S> for RefreshClaims
where
    S: Send + Sync,
{
    type Rejection = Error;

    async fn from_request_parts(parts: &mut Parts, _: &S) -> Result<Self, Self::Rejection> {
        let TypedHeader(Authorization(bearer)) = parts
            .extract::<TypedHeader<Authorization<Bearer>>>()
            .await
            .map_err(|_| Error::JwtNotFound)?;

        JWT.decode(bearer.token()).map(|jwt| jwt.claims)
    }
}