use async_session::{
    base64,
    hmac::{Hmac, Mac, NewMac},
    sha2::Sha256,
};
pub use async_session::{CookieStore, MemoryStore, Session, SessionStore};
use cookie::{Cookie, Key, SameSite};
use hypers_core::{
    async_trait,
    prelude::{Next, Request, Response},
    Error, Hook,
};
use std::time::Duration;

pub trait SessionDepotExt {
    /// Sets session
    fn set_session(&mut self, session: Session) -> &mut Self;
    /// Take session
    fn take_session(&mut self) -> Option<Session>;
    /// Get session reference
    fn session(&self) -> Option<&Session>;
    /// Get session mutable reference
    fn session_mut(&mut self) -> Option<&mut Session>;
}
impl SessionDepotExt for Request {
    #[inline]
    fn set_session(&mut self, session: Session) -> &mut Self {
        self.extensions_mut().insert(session);
        self
    }
    #[inline]
    fn take_session(&mut self) -> Option<Session> {
        self.extensions_mut().remove()
    }
    #[inline]
    fn session(&self) -> Option<&Session> {
        self.extensions().get()
    }
    #[inline]
    fn session_mut(&mut self) -> Option<&mut Session> {
        self.extensions_mut().get_mut()
    }
}
impl SessionDepotExt for Response {
    #[inline]
    fn set_session(&mut self, session: Session) -> &mut Self {
        self.extensions_mut().insert(session);
        self
    }
    #[inline]
    fn take_session(&mut self) -> Option<Session> {
        self.extensions_mut().remove()
    }
    #[inline]
    fn session(&self) -> Option<&Session> {
        self.extensions().get()
    }
    #[inline]
    fn session_mut(&mut self) -> Option<&mut Session> {
        self.extensions_mut().get_mut()
    }
}

/// `SessionHook` is a middleware for session.
#[derive(Debug)]
pub struct SessionHook<S> {
    store: S,
    cookie_path: String,
    cookie_name: String,
    cookie_value: Option<String>,
    cookie_domain: Option<String>,
    session_ttl: Option<Duration>,
    same_site: SameSite,
    hmac: Hmac<Sha256>,
    fallback_hmacs: Vec<Hmac<Sha256>>,
}
impl<S> SessionHook<S>
where
    S: SessionStore,
{
    /// Create new `SessionHook`
    #[inline]
    pub fn new(store: S, secret: &[u8]) -> Result<Self, Error> {
        let hmac = Hmac::<Sha256>::new_from_slice(Key::from(secret).signing())
            .map_err(|_| "SessionHook error: invalid key length")?;
        Ok(Self {
            store,
            cookie_path: "/".into(),
            cookie_name: "hypers.session_id".into(),
            cookie_value: None,
            cookie_domain: None,
            same_site: SameSite::Lax,
            session_ttl: Some(Duration::from_secs(24 * 60 * 60)),
            hmac,
            fallback_hmacs: Vec::new(),
        })
    }
    /// Create new `SessionHook` whith `Session`
    #[inline]
    pub async fn session(store: S, secret: &[u8], session: Session) -> Result<Self, Error> {
        let mut hook = Self::new(store, secret)?;
        hook.cookie_value = hook.store.store_session(session).await?;
        Ok(hook)
    }
    /// Sets a cookie path for this session middleware.The default for this value is "/".
    #[inline]
    pub fn cookie_path(mut self, cookie_path: impl Into<String>) -> Self {
        self.cookie_path = cookie_path.into();
        self
    }
    /// Sets a session ttl. This will be used both for the cookie expiry and also for the session-internal expiry.
    /// The default for this value is one day. Set this to None to not set a cookie or session expiry. This is not recommended.
    #[inline]
    pub fn session_ttl(mut self, session_ttl: Option<Duration>) -> Self {
        self.session_ttl = session_ttl;
        self
    }
    /// Sets the name of the cookie that the session is stored with or in.
    /// If you are running multiple tide applications on the same domain, you will need different values for each application.
    /// The default value is "hypers.session_id".
    #[inline]
    pub fn cookie_name(mut self, cookie_name: impl Into<String>) -> Self {
        self.cookie_name = cookie_name.into();
        self
    }
    /// Sets the same site policy for the session cookie. Defaults to SameSite::Lax.
    /// See [incrementally better cookies](https://tools.ietf.org/html/draft-west-cookie-incrementalism-01) for more information about this setting.
    #[inline]
    pub fn same_site_policy(mut self, policy: SameSite) -> Self {
        self.same_site = policy;
        self
    }
    /// Sets the domain of the cookie.
    #[inline]
    pub fn cookie_domain(mut self, cookie_domain: impl AsRef<str>) -> Self {
        self.cookie_domain = Some(cookie_domain.as_ref().to_owned());
        self
    }
    /// Add fallback secret.
    #[inline]
    pub fn fallback_key(mut self, key: impl Into<Key>) -> Result<Self, Error> {
        self.fallback_hmacs.push(
            Hmac::<Sha256>::new_from_slice(key.into().signing())
                .map_err(|_| "SessionHook error: invalid key length")?,
        );
        Ok(self)
    }
    #[inline]
    async fn load_session(
        &self,
        mut req: Request,
        next: Next<'_>,
        cookie_value: Option<String>,
    ) -> Response {
        match cookie_value {
            Some(cookie_value) => match self.store.load_session(cookie_value).await {
                Ok(session) => match session {
                    Some(mut session) => {
                        if let Some(ttl) = self.session_ttl {
                            session.expire_in(ttl);
                        };
                        req.set_session(session);
                        next.next(req).await
                    }
                    None => next.next(req).await,
                },
                Err(e) => next.next(req).await.render(500).render(e.to_string()),
            },
            None => {
                let secure_cookie =
                    req.uri().scheme() == Some(&hypers_core::hyper::http::uri::Scheme::HTTPS);
                let mut res = next.next(req).await;
                match res.take_session() {
                    Some(session) => match self.store.store_session(session).await {
                        Ok(cookie_value) => match cookie_value {
                            Some(cookie_value) => {
                                let cookie = self.build_cookie(secure_cookie, cookie_value);
                                res.cookie(cookie);
                                res
                            }
                            None => res,
                        },
                        Err(e) => res.render(500).render(e.to_string()),
                    },
                    None => res,
                }
            }
        }
    }
    #[inline]
    fn verify_signature(&self, cookie_value: &str) -> Result<String, Error> {
        if cookie_value.len() < 44 {
            return Err("length of value is <= 44".into());
        }
        // Split [MAC | original-value] into its two parts.
        let (digest_str, value) = cookie_value.split_at(44);
        let digest = base64::decode(digest_str)?;
        // Perform the verification.
        let mut hmac = self.hmac.clone();
        hmac.update(value.as_bytes());
        if hmac.verify(&digest).is_ok() {
            return Ok(value.to_string());
        }
        for hmac in &self.fallback_hmacs {
            let mut hmac = hmac.clone();
            hmac.update(value.as_bytes());
            if hmac.verify(&digest).is_ok() {
                return Ok(value.to_string());
            }
        }
        Err("value did not verify".into())
    }
    // signs the cookie's value providing integrity and authenticity.
    // The following is reused verbatim from https://github.com/SergioBenitez/cookie-rs/blob/master/src/secure/signed.rs#L37-46
    #[inline]
    fn sign_cookie(&self, cookie: &mut Cookie<'_>) {
        // Compute HMAC-SHA256 of the cookie's value.
        let mut mac = self.hmac.clone();
        mac.update(cookie.value().as_bytes());
        // Cookie's new value is [MAC | original-value].
        let mut new_value = base64::encode(mac.finalize().into_bytes());
        new_value.push_str(cookie.value());
        cookie.set_value(new_value);
    }
    #[inline]
    fn build_cookie(&self, secure: bool, cookie_value: String) -> Cookie<'static> {
        let mut cookie = Cookie::build((self.cookie_name.clone(), cookie_value))
            .http_only(true)
            .same_site(self.same_site)
            .secure(secure)
            .path(self.cookie_path.clone())
            .build();
        if let Some(ttl) = self.session_ttl {
            cookie.set_expires(Some((std::time::SystemTime::now() + ttl).into()));
        }
        if let Some(cookie_domain) = self.cookie_domain.clone() {
            cookie.set_domain(cookie_domain)
        }
        self.sign_cookie(&mut cookie);
        cookie
    }
}
#[async_trait]
impl<S> Hook for SessionHook<S>
where
    S: SessionStore,
{
    #[inline]
    async fn handle<'a>(&'a self, req: Request, next: Next<'a>) -> Response {
        match req.cookies.get(&self.cookie_name) {
            Some(cookie) => match self.verify_signature(cookie.value()) {
                Ok(cookie_value) => self.load_session(req, next, Some(cookie_value)).await,
                Err(e) => next.next(req).await.render(500).render(e.to_string()),
            },
            None => {
                self.load_session(req, next, self.cookie_value.clone())
                    .await
            }
        }
    }
}
#[test]
fn test_session_data() -> Result<(), Error> {
    let hook = SessionHook::new(
        async_session::CookieStore,
        b"secretabsecretabsecretabsecretabsecretabsecretabsecretabsecretab",
    )?
    .cookie_domain("test.domain")
    .cookie_name("test_cookie")
    .cookie_path("/abc")
    .same_site_policy(SameSite::Strict)
    .session_ttl(Some(Duration::from_secs(30)));
    assert!(format!("{:?}", hook).contains("test_cookie"));
    assert_eq!(hook.cookie_domain, Some("test.domain".into()));
    assert_eq!(hook.cookie_name, "test_cookie");
    assert_eq!(hook.cookie_path, "/abc");
    assert_eq!(hook.same_site, SameSite::Strict);
    assert_eq!(hook.session_ttl, Some(Duration::from_secs(30)));
    Ok(())
}