use aws_lc_rs::digest::{Context, SHA256, SHA384};

use crate::message::HashAlgorithm;

/// A hash context that supports SHA256 and SHA384 algorithms
pub struct Hash {
    context: Context,
}

impl Hash {
    /// Create a new hash context with the specified algorithm
    pub fn new(algorithm: HashAlgorithm) -> Self {
        let context = match algorithm {
            HashAlgorithm::SHA256 => Context::new(&SHA256),
            HashAlgorithm::SHA384 => Context::new(&SHA384),
            _ => panic!("Unsupported hash algorithm for handshake: {:?}", algorithm),
        };
        Hash { context }
    }

    /// Update the hash with new data
    pub fn update(&mut self, data: &[u8]) {
        self.context.update(data);
    }

    /// Finalize the hash and return the result. This clones the state, so
    /// it is possible to continue the hashing.
    pub fn clone_and_finalize(&self) -> Vec<u8> {
        self.context.clone().finish().as_ref().to_vec()
    }
}

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

    #[test]
    fn test_sha256() {
        let mut hash = Hash::new(HashAlgorithm::SHA256);
        hash.update(b"hello");
        hash.update(b" ");
        hash.update(b"world");
        let result = hash.clone_and_finalize();

        // Expected SHA256 hash of "hello world"
        let expected = [
            0xb9, 0x4d, 0x27, 0xb9, 0x93, 0x4d, 0x3e, 0x08, 0xa5, 0x2e, 0x52, 0xd7, 0xda, 0x7d,
            0xab, 0xfa, 0xc4, 0x84, 0xef, 0xe3, 0x7a, 0x53, 0x80, 0xee, 0x90, 0x88, 0xf7, 0xac,
            0xe2, 0xef, 0xcd, 0xe9,
        ];

        assert_eq!(result, expected);
    }

    #[test]
    fn test_sha384() {
        let mut hash = Hash::new(HashAlgorithm::SHA384);
        hash.update(b"hello");
        hash.update(b" ");
        hash.update(b"world");
        let result = hash.clone_and_finalize();

        // Expected SHA384 hash of "hello world"
        let expected = [
            0xfd, 0xbd, 0x8e, 0x75, 0xa6, 0x7f, 0x29, 0xf7, 0x01, 0xa4, 0xe0, 0x40, 0x38, 0x5e,
            0x2e, 0x23, 0x98, 0x63, 0x03, 0xea, 0x10, 0x23, 0x92, 0x11, 0xaf, 0x90, 0x7f, 0xcb,
            0xb8, 0x35, 0x78, 0xb3, 0xe4, 0x17, 0xcb, 0x71, 0xce, 0x64, 0x6e, 0xfd, 0x08, 0x19,
            0xdd, 0x8c, 0x08, 0x8d, 0xe1, 0xbd,
        ];

        assert_eq!(result, expected);
    }

    #[test]
    fn test_clone_and_finalize_continues_hashing() {
        let mut hash = Hash::new(HashAlgorithm::SHA256);

        // First part of the message
        hash.update(b"hello");
        hash.update(b" ");

        // Get intermediate hash
        let intermediate = hash.clone_and_finalize();

        // Continue hashing
        hash.update(b"world");
        let final_result = hash.clone_and_finalize();

        // Create a new hash for comparison
        let mut expected_hash = Hash::new(HashAlgorithm::SHA256);
        expected_hash.update(b"hello");
        expected_hash.update(b" ");
        expected_hash.update(b"world");
        let expected = expected_hash.clone_and_finalize();

        // Verify the final result matches what we expect
        assert_eq!(final_result, expected);

        // Verify the intermediate hash is different from the final hash
        assert_ne!(intermediate, final_result);
    }
}
