use hmac::{Hmac, Mac};
use sha2::Sha512;
use subtle::ConstantTimeEq;

use crate::error::{OpaqueError, Result};

pub const MAC_LEN: usize = 64;

type HmacSha512 = Hmac<Sha512>;

pub fn compute(key: &[u8], data: &[u8]) -> [u8; MAC_LEN] {
    let mut mac = HmacSha512::new_from_slice(key).expect("HMAC accepts any key length");
    mac.update(data);
    mac.finalize().into_bytes().into()
}

pub fn verify(key: &[u8], data: &[u8], expected: &[u8]) -> Result<()> {
    if expected.len() != MAC_LEN {
        return Err(OpaqueError::MacVerificationFailed);
    }

    let computed = compute(key, data);
    if computed.ct_eq(expected).into() {
        Ok(())
    } else {
        Err(OpaqueError::MacVerificationFailed)
    }
}

pub fn compute_multi(key: &[u8], parts: &[&[u8]]) -> [u8; MAC_LEN] {
    let mut mac = HmacSha512::new_from_slice(key).expect("HMAC accepts any key length");
    for part in parts {
        mac.update(part);
    }
    mac.finalize().into_bytes().into()
}

pub struct HmacContext {
    mac: HmacSha512,
}

impl HmacContext {
    #[must_use]
    pub fn new(key: &[u8]) -> Self {
        let mac = HmacSha512::new_from_slice(key).expect("HMAC accepts any key length");
        Self { mac }
    }

    pub fn update(&mut self, data: &[u8]) {
        self.mac.update(data);
    }

    #[must_use]
    pub fn finalize(self) -> [u8; MAC_LEN] {
        self.mac.finalize().into_bytes().into()
    }

    pub fn verify(self, expected: &[u8]) -> Result<()> {
        if expected.len() != MAC_LEN {
            return Err(OpaqueError::MacVerificationFailed);
        }

        let computed = self.finalize();
        if computed.ct_eq(expected).into() {
            Ok(())
        } else {
            Err(OpaqueError::MacVerificationFailed)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_compute_and_verify() {
        let key = b"secret key";
        let data = b"message to authenticate";

        let tag = compute(key, data);
        assert_eq!(tag.len(), MAC_LEN);

        assert!(verify(key, data, &tag).is_ok());
    }

    #[test]
    fn test_verify_wrong_tag() {
        let key = b"secret key";
        let data = b"message";
        let wrong_tag = [0u8; MAC_LEN];

        assert!(verify(key, data, &wrong_tag).is_err());
    }

    #[test]
    fn test_verify_wrong_length() {
        let key = b"secret key";
        let data = b"message";
        let short_tag = [0u8; 32];

        assert!(verify(key, data, &short_tag).is_err());
    }

    #[test]
    fn test_compute_multi() {
        let key = b"secret key";
        let part1 = b"hello ";
        let part2 = b"world";

        let tag_multi = compute_multi(key, &[part1, part2]);
        let tag_single = compute(key, b"hello world");

        assert_eq!(tag_multi, tag_single);
    }

    #[test]
    fn test_hmac_context() {
        let key = b"secret key";
        let data = b"message to authenticate";

        let mut ctx = HmacContext::new(key);
        ctx.update(data);
        let tag1 = ctx.finalize();

        let tag2 = compute(key, data);

        assert_eq!(tag1, tag2);
    }

    #[test]
    fn test_hmac_context_verify() {
        let key = b"secret key";
        let data = b"message";

        let tag = compute(key, data);

        let mut ctx = HmacContext::new(key);
        ctx.update(data);
        assert!(ctx.verify(&tag).is_ok());
    }

    #[test]
    fn test_different_keys_different_tags() {
        let data = b"message";
        let tag1 = compute(b"key1", data);
        let tag2 = compute(b"key2", data);

        assert_ne!(tag1, tag2);
    }
}