use std::convert::Infallible;

use axum::{
    extract::FromRequestParts,
    http::request::Parts,
    response::{IntoResponseParts, ResponseParts},
};
use axum_extra::extract::cookie::{Cookie, CookieJar};

// The usage pattern here - receive the extractor as an argument, return it in
// the response - is heavily modelled after CookieJar's own intended usage.
pub struct IdentityToken {
    cookies: CookieJar,
}

impl IdentityToken {
    /// Get the identity token sent in the request, if any. If the identity was
    /// not sent, or if it has previously been [clear]ed, then this will return
    /// [None]. If the identity has previously been [set], then this will return
    /// that token.
    pub fn token(&self) -> Option<&str> {
        self.cookies.get(IDENTITY_COOKIE).map(Cookie::value)
    }

    /// Positively set the identity token, and ensure that it will be sent back
    /// to the client when this extractor is included in a response.
    pub fn set(self, token: &str) -> Self {
        let identity_cookie = Cookie::build((IDENTITY_COOKIE, String::from(token)))
            .http_only(true)
            .permanent()
            .build();

        IdentityToken {
            cookies: self.cookies.add(identity_cookie),
        }
    }

    /// Remove the identity token and ensure that it will be cleared when this
    /// extractor is included in a response.
    pub fn clear(self) -> Self {
        IdentityToken {
            cookies: self.cookies.remove(IDENTITY_COOKIE),
        }
    }
}

const IDENTITY_COOKIE: &str = "identity";

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

    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
        let cookies = CookieJar::from_request_parts(parts, state).await?;
        Ok(IdentityToken { cookies })
    }
}

impl IntoResponseParts for IdentityToken {
    type Error = Infallible;

    fn into_response_parts(self, res: ResponseParts) -> Result<ResponseParts, Self::Error> {
        let IdentityToken { cookies } = self;
        cookies.into_response_parts(res)
    }
}