#![deny(clippy::pedantic)]
#![deny(clippy::restriction)]
#![allow(clippy::blanket_clippy_restriction_lints)]
#![allow(clippy::implicit_return)]
#![allow(clippy::missing_docs_in_private_items)]
#![allow(clippy::missing_inline_in_public_items)]
#![allow(clippy::shadow_reuse)]

use std::fmt::Debug;

use hmac::{Hmac, Mac};
use password_hash::{Output, Salt};
use rand::RngCore;
use scrypt::scrypt;
use serde::{Deserialize, Serialize};
use sha2::Sha256;

use crate::{
    serde::as_hex,
    types::{AccountResetID, HawkKey, KeyFetchID, PasswordChangeID, SecretKey, SessionID},
};

const NAMESPACE: &[u8] = b"identity.mozilla.com/picl/v1/";

pub fn random_bytes<const N: usize>() -> [u8; N] {
    let mut result = [0; N];
    rand::rngs::OsRng.fill_bytes(&mut result);
    result
}

fn xor<const N: usize>(l: &[u8; N], r: &[u8; N]) -> [u8; N] {
    let mut result = *l;
    for (a, b) in result.iter_mut().zip(r.iter()) {
        *a ^= b;
    }
    result
}

mod from_hkdf {
    use hkdf::Hkdf;
    use sha2::Sha256;

    pub trait FromHkdf {
        const SIZE: usize;
        fn from_hkdf(bytes: &[u8]) -> Self;
    }

    impl<const N: usize> FromHkdf for [u8; N] {
        const SIZE: usize = N;
        fn from_hkdf(bytes: &[u8]) -> Self {
            #[allow(clippy::expect_used)]
            bytes.try_into().expect("hkdf failed")
        }
    }

    impl<L: FromHkdf, R: FromHkdf> FromHkdf for (L, R) {
        const SIZE: usize = L::SIZE + R::SIZE;
        #[allow(clippy::indexing_slicing)]
        fn from_hkdf(bytes: &[u8]) -> Self {
            (L::from_hkdf(&bytes[0..L::SIZE]), R::from_hkdf(&bytes[L::SIZE..]))
        }
    }

    pub fn from_hkdf<O: FromHkdf>(key: &[u8], info: &[&[u8]]) -> O {
        let hk = Hkdf::<Sha256>::new(None, key);
        let mut buf = vec![0; O::SIZE];
        #[allow(clippy::expect_used)]
        // worth keeping an eye out for very large results (>255*32 bytes)
        hk.expand_multi_info(info, buf.as_mut_slice()).expect("hkdf failed");
        O::from_hkdf(&buf)
    }
}

use from_hkdf::from_hkdf;

#[derive(Debug, Deserialize, Serialize)]
pub(crate) struct AuthPW(#[serde(with = "as_hex")] [u8; 32]);

pub(crate) struct StretchedPW {
    pw: Output,
}

impl AuthPW {
    pub fn stretch(&self, salt: Salt) -> anyhow::Result<StretchedPW> {
        let mut result = [0; 32];
        let mut buf = [0; Salt::MAX_LENGTH];
        let params = scrypt::Params::new(16, 8, 1)?;
        let salt = salt.b64_decode(&mut buf)?;
        scrypt(&self.0, salt, &params, &mut result)?;
        Ok(StretchedPW { pw: Output::new(&result)? })
    }
}

impl StretchedPW {
    pub fn verify_hash(&self) -> Output {
        let raw: [u8; 32] = from_hkdf(self.pw.as_bytes(), &[NAMESPACE, b"verifyHash"]);
        #[allow(clippy::unwrap_used)]
        Output::new(&raw).unwrap()
    }

    fn wrap_wrap_key(&self) -> [u8; 32] {
        from_hkdf(self.pw.as_bytes(), &[NAMESPACE, b"wrapwrapKey"])
    }

    pub fn decrypt_wwkb(&self, wwkb: &SecretKey) -> SecretKey {
        SecretKey(xor(&wwkb.0, &self.wrap_wrap_key()))
    }

    pub fn rewrap_wkb(&self, wkb: &SecretKey) -> SecretKey {
        SecretKey(xor(&wkb.0, &self.wrap_wrap_key()))
    }
}

#[derive(Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub(crate) struct SessionToken(#[serde(with = "as_hex")] [u8; 32]);

impl Debug for SessionToken {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_tuple("SessionToken").field(&hex::encode(self.0)).finish()
    }
}

impl SessionToken {
    pub fn generate() -> Self {
        Self(random_bytes())
    }
}

pub(crate) struct SessionCredentials {
    pub token_id: SessionID,
    pub req_hmac_key: HawkKey,
}

impl SessionCredentials {
    pub fn derive_from(seed: &SessionToken) -> Self {
        let (token_id, req_hmac_key) = from_hkdf(&seed.0, &[NAMESPACE, b"sessionToken"]);
        Self { token_id: SessionID(token_id), req_hmac_key: HawkKey(req_hmac_key) }
    }
}

#[derive(Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub(crate) struct KeyFetchToken(#[serde(with = "as_hex")] [u8; 32]);

impl Debug for KeyFetchToken {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_tuple("KeyFetchToken").field(&hex::encode(self.0)).finish()
    }
}

impl KeyFetchToken {
    pub fn generate() -> Self {
        Self(random_bytes())
    }
}

pub(crate) struct KeyFetchReq {
    pub token_id: KeyFetchID,
    pub req_hmac_key: HawkKey,
    key_request_key: [u8; 32],
}

impl KeyFetchReq {
    pub fn derive_from(key_fetch_token: &KeyFetchToken) -> Self {
        let (token_id, (req_hmac_key, key_request_key)) =
            from_hkdf(&key_fetch_token.0, &[NAMESPACE, b"keyFetchToken"]);
        Self {
            token_id: KeyFetchID(token_id),
            req_hmac_key: HawkKey(req_hmac_key),
            key_request_key,
        }
    }

    pub fn derive_resp(&self) -> KeyFetchResp {
        let (resp_hmac_key, resp_xor_key) =
            from_hkdf(&self.key_request_key, &[NAMESPACE, b"account/keys"]);
        KeyFetchResp { resp_hmac_key, resp_xor_key }
    }
}

pub(crate) struct KeyFetchResp {
    resp_hmac_key: [u8; 32],
    resp_xor_key: [u8; 64],
}

impl KeyFetchResp {
    pub fn wrap_keys(&self, keys: &KeyBundle) -> WrappedKeyBundle {
        let ciphertext = xor(&self.resp_xor_key, &keys.to_bytes());
        #[allow(clippy::unwrap_used)]
        let mut hmac = Hmac::<Sha256>::new_from_slice(&self.resp_hmac_key).unwrap();
        hmac.update(&ciphertext);
        let hmac = *hmac.finalize().into_bytes().as_ref();
        WrappedKeyBundle { ciphertext, hmac }
    }
}

pub(crate) struct KeyBundle {
    pub ka: SecretKey,
    pub wrap_kb: SecretKey,
}

impl KeyBundle {
    pub fn to_bytes(&self) -> [u8; 64] {
        let mut result = [0; 64];
        result[0..32].copy_from_slice(&self.ka.0);
        result[32..].copy_from_slice(&self.wrap_kb.0);
        result
    }
}

#[derive(Debug)]
pub struct WrappedKeyBundle {
    ciphertext: [u8; 64],
    hmac: [u8; 32],
}

impl WrappedKeyBundle {
    pub fn to_bytes(&self) -> [u8; 96] {
        let mut result = [0; 96];
        result[0..64].copy_from_slice(&self.ciphertext);
        result[64..].copy_from_slice(&self.hmac);
        result
    }
}

#[derive(Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub(crate) struct PasswordChangeToken(#[serde(with = "as_hex")] [u8; 32]);

impl Debug for PasswordChangeToken {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_tuple("PasswordChangeToken").field(&hex::encode(self.0)).finish()
    }
}

impl PasswordChangeToken {
    pub fn generate() -> Self {
        Self(random_bytes())
    }
}

pub(crate) struct PasswordChangeReq {
    pub token_id: PasswordChangeID,
    pub req_hmac_key: HawkKey,
}

impl PasswordChangeReq {
    pub fn derive_from_change_token(token: &PasswordChangeToken) -> Self {
        let (token_id, req_hmac_key) = from_hkdf(&token.0, &[NAMESPACE, b"passwordChangeToken"]);
        Self { token_id: PasswordChangeID(token_id), req_hmac_key: HawkKey(req_hmac_key) }
    }

    pub fn derive_from_forgot_token(token: &PasswordChangeToken) -> Self {
        let (token_id, req_hmac_key) = from_hkdf(&token.0, &[NAMESPACE, b"passwordForgotToken"]);
        Self { token_id: PasswordChangeID(token_id), req_hmac_key: HawkKey(req_hmac_key) }
    }
}

#[derive(Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub(crate) struct AccountResetToken(#[serde(with = "as_hex")] [u8; 32]);

impl Debug for AccountResetToken {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_tuple("AccountResetToken").field(&hex::encode(self.0)).finish()
    }
}

impl AccountResetToken {
    pub fn generate() -> Self {
        Self(random_bytes())
    }
}

pub(crate) struct AccountResetReq {
    pub token_id: AccountResetID,
    pub req_hmac_key: HawkKey,
}

impl AccountResetReq {
    pub fn derive_from(token: &AccountResetToken) -> Self {
        let (token_id, req_hmac_key) = from_hkdf(&token.0, &[NAMESPACE, b"accountResetToken"]);
        Self { token_id: AccountResetID(token_id), req_hmac_key: HawkKey(req_hmac_key) }
    }
}

#[cfg(test)]
mod test {
    use hex_literal::hex;
    use password_hash::{Output, SaltString};

    use crate::{
        crypto::{
            AccountResetReq, AccountResetToken, KeyBundle, KeyFetchReq, KeyFetchToken,
            PasswordChangeReq, PasswordChangeToken, SessionCredentials, SessionToken,
        },
        types::SecretKey,
    };

    use super::AuthPW;

    #[test]
    fn test_derive_session() {
        let creds = SessionCredentials::derive_from(&SessionToken(hex!(
            "a0a1a2a3a4a5a6a7 a8a9aaabacadaeaf b0b1b2b3b4b5b6b7 b8b9babbbcbdbebf"
        )));
        assert_eq!(
            creds.token_id.0,
            hex!("c0a29dcf46174973da1378696e4c82ae10f723cf4f4d9f75e39f4ae3851595ab")
        );
        assert_eq!(
            creds.req_hmac_key.0,
            hex!("9d8f22998ee7f579 8b887042466b72d5 3e56ab0c094388bf 65831f702d2febc0")
        );
    }

    #[test]
    fn test_key_fetch() {
        let key_fetch = KeyFetchReq::derive_from(&KeyFetchToken(hex!(
            "8081828384858687 88898a8b8c8d8e8f 9091929394959697 98999a9b9c9d9e9f"
        )));
        assert_eq!(
            key_fetch.token_id.0,
            hex!("3d0a7c02a15a62a2882f76e39b6494b500c022a8816e048625a495718998ba60")
        );
        assert_eq!(
            key_fetch.req_hmac_key.0,
            hex!("87b8937f61d38d0e 29cd2d5600b3f4da 0aa48ac41de36a0e fe84bb4a9872ceb7")
        );
        assert_eq!(
            key_fetch.key_request_key,
            hex!("14f338a9e8c6324d 9e102d4e6ee83b20 9796d5c74bb734a4 10e729e014a4a546")
        );

        let resp = key_fetch.derive_resp();
        assert_eq!(
            resp.resp_hmac_key,
            hex!("f824d2953aab9faf 51a1cb65ba9e7f9e 5bf91c8d8fd1ac1c 8c2d31853a8a1210")
        );
        assert_eq!(
            resp.resp_xor_key,
            hex!(
                "ce7d7aa77859b235 9932970bbe2101f2 e80d01faf9191bd5 ee52181d2f0b7809
                 8281ba8cff392543 3a89f7c3095e0c89 900a469d60790c83 3281c4df1a11c763"
            )
        );

        let bundle = KeyBundle {
            ka: SecretKey(hex!(
                "2021222324252627 28292a2b2c2d2e2f 3031323334353637 38393a3b3c3d3e3f"
            )),
            wrap_kb: SecretKey(hex!(
                "7effe354abecbcb2 34a8dfc2d7644b4a d339b525589738f2 d27341bb8622ecd8"
            )),
        };
        assert_eq!(
            bundle.to_bytes(),
            hex!(
                "2021222324252627 28292a2b2c2d2e2f 3031323334353637 38393a3b3c3d3e3f
                 7effe354abecbcb2 34a8dfc2d7644b4a d339b525589738f2 d27341bb8622ecd8"
            )
        );

        let wrapped = resp.wrap_keys(&bundle);
        assert_eq!(
            wrapped.ciphertext,
            hex!(
                "ee5c58845c7c9412 b11bbd20920c2fdd d83c33c9cd2c2de2 d66b222613364636
                 fc7e59d854d599f1 0e212801de3a47c3 4333f3b838ee3471 e0f285649c332bbb"
            )
        );
        assert_eq!(
            wrapped.hmac,
            hex!("4c17f42a0b319bbb a327d2b326ad23e9 37219b4de32e3ec7 b3e3f740522ad6ef")
        );
        assert_eq!(
            wrapped.to_bytes(),
            hex!(
                "ee5c58845c7c9412 b11bbd20920c2fdd d83c33c9cd2c2de2 d66b222613364636
                 fc7e59d854d599f1 0e212801de3a47c3 4333f3b838ee3471 e0f285649c332bbb
                 4c17f42a0b319bbb a327d2b326ad23e9 37219b4de32e3ec7 b3e3f740522ad6ef"
            )
        );
    }

    #[test]
    fn test_stretch() -> anyhow::Result<()> {
        let auth_pw =
            AuthPW(hex!("247b675ffb4c4631 0bc87e26d712153a be5e1c90ef00a478 4594f97ef54f2375"));

        let stretched = auth_pw.stretch(
            SaltString::b64_encode(&hex!(
                "00f0000000000000 0000000000000000 0000000000000000 0000000000000000"
            ))?
            .as_salt(),
        )?;
        assert_eq!(
            stretched.pw,
            Output::new(&hex!(
                "441509e25c92ee10 3d5a1a874e6f155d f25a44d06e61c894 616c9e85181dba97"
            ))?
        );

        assert_eq!(
            stretched.verify_hash().as_bytes(),
            hex!("a4765bf103dc057f 4cf4bc2c131ddb67 16e8a4333cc55e1d 3c449f31f0eec4f1")
        );

        assert_eq!(
            stretched.wrap_wrap_key(),
            hex!("3ebea117efa9faf5 7ce195899b290505 8368e7760cc26ea5 8a2a1be0da7fb287")
        );
        Ok(())
    }

    #[test]
    fn test_password_change() {
        let req = PasswordChangeReq::derive_from_change_token(&PasswordChangeToken(hex!(
            "0000000000000000 0000000000000000 0000000000000000 0000000000000000"
        )));
        assert_eq!(
            req.token_id.0,
            hex!("5a9f93f66c26fd1c 1ea9826fafc422e9 4b9c9f833cd2bfa5 da18c8d3317224aa")
        );
        assert_eq!(
            req.req_hmac_key.0,
            hex!("31940008f942939a 22d7cf8ad38dc6ac 346a8148439c0d98 d4a6ae8352da4536")
        );
    }

    #[test]
    fn test_password_forgot() {
        let req = PasswordChangeReq::derive_from_forgot_token(&PasswordChangeToken(hex!(
            "0000000000000000 0000000000000000 0000000000000000 0000000000000000"
        )));
        assert_eq!(
            req.token_id.0,
            hex!("570e79050fd157a9 b8e7d7d6f88a3f67 e36207c5dfabe7d8 a80994502a624e07")
        );
        assert_eq!(
            req.req_hmac_key.0,
            hex!("fde451da23c03eec fe94e401eeef8bff 5c51742839b8058e c216214f78742d9a")
        );
    }

    #[test]
    fn test_account_reset() {
        let req = AccountResetReq::derive_from(&AccountResetToken(hex!(
            "0000000000000000 0000000000000000 0000000000000000 0000000000000000"
        )));
        assert_eq!(
            req.token_id.0,
            hex!("8ade842449ab0285 e7b22de9d428cd5b 3c38ea0aa78e2956 a6a69ec66818d864")
        );
        assert_eq!(
            req.req_hmac_key.0,
            hex!("d17d0a55c5a0451c f21efaf39f3611bc 54ebc530ceb1fe8c 9be330c1b68f989f")
        );
    }
}