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, Serialize};
use time::OffsetDateTime;
use uuid::Uuid;

use super::{Error, JWT};

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

impl<const LIFETIME: i64> Claims<LIFETIME> {
    pub fn new(uuid: Uuid) -> Self {
        let now = OffsetDateTime::now_utc().unix_timestamp();
        Self {
            sub: uuid,
            iat: now,
            exp: now + LIFETIME,
            jti: uuid::Uuid::new_v4(),
        }
    }
}

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

pub type AccessClaims = Claims<ACCESS>;

impl From<RefreshClaims> for AccessClaims {
    fn from(value: RefreshClaims) -> Self {
        Claims::new(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(OffsetDateTime::from_unix_timestamp(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)?;

        Ok(JWT.decode(bearer.token())?.claims)
    }
}