//
// Syd: rock-solid unikernel
// src/hash.rs: Utilities for hashing
//
// Copyright (c) 2024 Ali Polatel <alip@chesswob.org>
//
// SPDX-License-Identifier: GPL-3.0-or-later

use std::{
    hash::{BuildHasher, Hasher},
    io::BufRead,
};

use nix::errno::Errno;

use crate::err::SydResult;

/// Defines hash functions supported by syd.
#[derive(Debug, Clone, Copy)]
pub enum HashAlgorithm {
    /// Md5
    Md5,
    /// SHA-1.
    Sha1,
    /// SHA3-256
    Sha256,
    /// SHA3-384
    Sha384,
    /// SHA3-512
    Sha512,
}

impl TryFrom<usize> for HashAlgorithm {
    type Error = Errno;

    fn try_from(len: usize) -> Result<Self, Self::Error> {
        match len {
            16 => Ok(HashAlgorithm::Md5),
            20 => Ok(HashAlgorithm::Sha1),
            32 => Ok(HashAlgorithm::Sha256),
            48 => Ok(HashAlgorithm::Sha384),
            64 => Ok(HashAlgorithm::Sha512),
            _ => Err(Errno::EINVAL),
        }
    }
}

/// Calculate sha{1,256,512} of the given buffered reader.
/// Returns a byte array.
pub fn hash<R: BufRead>(mut reader: R, func: HashAlgorithm) -> SydResult<Vec<u8>> {
    match func {
        HashAlgorithm::Md5 => {
            let mut context = md5::Context::new();
            loop {
                let consumed = {
                    let buf = reader.fill_buf()?;
                    if buf.is_empty() {
                        break;
                    }
                    context.consume(buf);
                    buf.len()
                };
                reader.consume(consumed);
            }
            let digest = context.compute();
            let dbytes: [u8; 16] = digest.into();
            Ok(dbytes.to_vec())
        }
        HashAlgorithm::Sha1 => {
            let mut hasher = rs_sha1::Sha1Hasher::default();
            loop {
                let consumed = {
                    let buf = reader.fill_buf()?;
                    if buf.is_empty() {
                        break;
                    }
                    hasher.write(buf);
                    buf.len()
                };
                reader.consume(consumed);
            }
            Ok(rs_sha1::HasherContext::finish(&mut hasher)
                .as_ref()
                .to_vec())
        }
        HashAlgorithm::Sha256 => {
            let mut hasher = rs_sha3_256::Sha3_256Hasher::default();
            loop {
                let consumed = {
                    let buf = reader.fill_buf()?;
                    if buf.is_empty() {
                        break;
                    }
                    hasher.write(buf);
                    buf.len()
                };
                reader.consume(consumed);
            }
            Ok(rs_sha3_256::HasherContext::finish(&mut hasher)
                .as_ref()
                .to_vec())
        }
        HashAlgorithm::Sha384 => {
            let mut hasher = rs_sha3_384::Sha3_384Hasher::default();
            loop {
                let consumed = {
                    let buf = reader.fill_buf()?;
                    if buf.is_empty() {
                        break;
                    }
                    hasher.write(buf);
                    buf.len()
                };
                reader.consume(consumed);
            }
            Ok(rs_sha3_384::HasherContext::finish(&mut hasher)
                .as_ref()
                .to_vec())
        }
        HashAlgorithm::Sha512 => {
            let mut hasher = rs_sha3_512::Sha3_512Hasher::default();
            loop {
                let consumed = {
                    let buf = reader.fill_buf()?;
                    if buf.is_empty() {
                        break;
                    }
                    hasher.write(buf);
                    buf.len()
                };
                reader.consume(consumed);
            }
            Ok(rs_sha3_512::HasherContext::finish(&mut hasher)
                .as_ref()
                .to_vec())
        }
    }
}

/// Avoid duplicate hashing while using `HashSet` with u64 keys.
pub(crate) struct NoHasher {
    value: u64,
}

impl Hasher for NoHasher {
    fn write(&mut self, _bytes: &[u8]) {
        unreachable!("NoHasher should only be used for u64 keys");
    }

    fn write_u64(&mut self, i: u64) {
        self.value = i;
    }

    fn finish(&self) -> u64 {
        self.value
    }
}

/// A builder for creating instances of `NoHasher`.
#[derive(Clone)]
pub(crate) struct NoHasherBuilder;

impl Default for NoHasherBuilder {
    fn default() -> Self {
        Self
    }
}

impl BuildHasher for NoHasherBuilder {
    type Hasher = NoHasher;

    fn build_hasher(&self) -> Self::Hasher {
        NoHasher { value: 0 }
    }
}

// A `HashSet` with no hashers.
// pub(crate) type NoHashSet = HashSet<u64, NoHasherBuilder>;

#[cfg(test)]
mod tests {
    use std::io::Cursor;

    use hex::DisplayHex;

    use super::*;

    struct TestCase(&'static [u8], &'static str, HashAlgorithm);

    // Source: https://www.di-mgt.com.au/sha_testvectors.html
    // MD5 test vectors were calculated with python-3.11.8's hashlib.md5
    const TEST_CASES: &[TestCase] = &[
    TestCase(
        b"",
        "D41D8CD98F00B204E9800998ECF8427E",
        HashAlgorithm::Md5,
    ),
    TestCase(
        b"",
        "DA39A3EE5E6B4B0D3255BFEF95601890AFD80709",
        HashAlgorithm::Sha1,
    ),
    TestCase(
        b"",
        "A7FFC6F8BF1ED76651C14756A061D662F580FF4DE43B49FA82D80A4B80F8434A",
        HashAlgorithm::Sha256,
    ),
    TestCase(
        b"",
        "0C63A75B845E4F7D01107D852E4C2485C51A50AAAA94FC61995E71BBEE983A2AC3713831264ADB47FB6BD1E058D5F004",
        HashAlgorithm::Sha384,
    ),
    TestCase(
        b"",
        "A69F73CCA23A9AC5C8B567DC185A756E97C982164FE25859E0D1DCC1475C80A615B2123AF1F5F94C11E3E9402C3AC558F500199D95B6D3E301758586281DCD26",
        HashAlgorithm::Sha512,
    ),
    TestCase(
        b"abc",
        "900150983CD24FB0D6963F7D28E17F72",
        HashAlgorithm::Md5,
    ),
    TestCase(
        b"abc",
        "A9993E364706816ABA3E25717850C26C9CD0D89D",
        HashAlgorithm::Sha1,
    ),
    TestCase(
        b"abc",
        "3A985DA74FE225B2045C172D6BD390BD855F086E3E9D525B46BFE24511431532",
        HashAlgorithm::Sha256,
    ),
    TestCase(
        b"abc",
        "EC01498288516FC926459F58E2C6AD8DF9B473CB0FC08C2596DA7CF0E49BE4B298D88CEA927AC7F539F1EDF228376D25",
        HashAlgorithm::Sha384,
    ),
    TestCase(
        b"abc",
        "B751850B1A57168A5693CD924B6B096E08F621827444F70D884F5D0240D2712E10E116E9192AF3C91A7EC57647E3934057340B4CF408D5A56592F8274EEC53F0",
        HashAlgorithm::Sha512
    ),
    TestCase(
        b"abcdbcdecdefdefgefghfghighijhijkijkljklmklmnlmnomnopnopq",
        "8215EF0796A20BCAAAE116D3876C664A",
        HashAlgorithm::Md5,
    ),
    TestCase(
        b"abcdbcdecdefdefgefghfghighijhijkijkljklmklmnlmnomnopnopq",
        "84983E441C3BD26EBAAE4AA1F95129E5E54670F1",
        HashAlgorithm::Sha1,
    ),
    TestCase(
        b"abcdbcdecdefdefgefghfghighijhijkijkljklmklmnlmnomnopnopq",
        "41C0DBA2A9D6240849100376A8235E2C82E1B9998A999E21DB32DD97496D3376",
        HashAlgorithm::Sha256,
    ),
    TestCase(
        b"abcdbcdecdefdefgefghfghighijhijkijkljklmklmnlmnomnopnopq",
        "991C665755EB3A4B6BBDFB75C78A492E8C56A22C5C4D7E429BFDBC32B9D4AD5AA04A1F076E62FEA19EEF51ACD0657C22",
        HashAlgorithm::Sha384,
    ),
    TestCase(
        b"abcdbcdecdefdefgefghfghighijhijkijkljklmklmnlmnomnopnopq",
        "04A371E84ECFB5B8B77CB48610FCA8182DD457CE6F326A0FD3D7EC2F1E91636DEE691FBE0C985302BA1B0D8DC78C086346B533B49C030D99A27DAF1139D6E75E",
        HashAlgorithm::Sha512,
    ),
    TestCase(
        b"abcdefghbcdefghicdefghijdefghijkefghijklfghijklmghijklmnhijklmnoijklmnopjklmnopqklmnopqrlmnopqrsmnopqrstnopqrstu",
        "03DD8807A93175FB062DFB55DC7D359C",
        HashAlgorithm::Md5,
    ),
    TestCase(
        b"abcdefghbcdefghicdefghijdefghijkefghijklfghijklmghijklmnhijklmnoijklmnopjklmnopqklmnopqrlmnopqrsmnopqrstnopqrstu",
        "A49B2446A02C645BF419F995B67091253A04A259",
        HashAlgorithm::Sha1,
    ),
    TestCase(
        b"abcdefghbcdefghicdefghijdefghijkefghijklfghijklmghijklmnhijklmnoijklmnopjklmnopqklmnopqrlmnopqrsmnopqrstnopqrstu",
        "916F6061FE879741CA6469B43971DFDB28B1A32DC36CB3254E812BE27AAD1D18",
        HashAlgorithm::Sha256,
    ),
    TestCase(
        b"abcdefghbcdefghicdefghijdefghijkefghijklfghijklmghijklmnhijklmnoijklmnopjklmnopqklmnopqrlmnopqrsmnopqrstnopqrstu",
        "79407D3B5916B59C3E30B09822974791C313FB9ECC849E406F23592D04F625DC8C709B98B43B3852B337216179AA7FC7",
        HashAlgorithm::Sha384,
    ),
    TestCase(
        b"abcdefghbcdefghicdefghijdefghijkefghijklfghijklmghijklmnhijklmnoijklmnopjklmnopqklmnopqrlmnopqrsmnopqrstnopqrstu",
        "AFEBB2EF542E6579C50CAD06D2E578F9F8DD6881D7DC824D26360FEEBF18A4FA73E3261122948EFCFD492E74E82E2189ED0FB440D187F382270CB455F21DD185",
        HashAlgorithm::Sha512,
    ),
    ];

    #[test]
    fn test_hash_simple() {
        let mut errors = Vec::new();

        for case in TEST_CASES {
            let input_cursor = Cursor::new(case.0);
            let result = match hash(input_cursor, case.2) {
                Ok(hash) => hash.to_upper_hex_string(),
                Err(e) => {
                    errors.push(format!(
                        "Hashing failed for {:?} with error: {:?}",
                        case.2, e
                    ));
                    continue;
                }
            };

            if result != case.1 {
                errors.push(format!(
                    "Mismatch for {:?}: expected {}, got {}",
                    case.2, case.1, result
                ));
            }
        }

        assert!(errors.is_empty(), "Errors encountered: {:?}", errors);
    }

    #[test]
    fn test_hash_long() {
        let mut errors = Vec::new();

        let input = b"a".repeat(1_000_000);
        let cases = &[
            (HashAlgorithm::Md5, "7707D6AE4E027C70EEA2A935C2296F21"),
            (HashAlgorithm::Sha1, "34AA973CD4C4DAA4F61EEB2BDBAD27316534016F"),
            (HashAlgorithm::Sha256, "5C8875AE474A3634BA4FD55EC85BFFD661F32ACA75C6D699D0CDCB6C115891C1"),
            (HashAlgorithm::Sha384, "EEE9E24D78C1855337983451DF97C8AD9EEDF256C6334F8E948D252D5E0E76847AA0774DDB90A842190D2C558B4B8340"),
            (HashAlgorithm::Sha512, "3C3A876DA14034AB60627C077BB98F7E120A2A5370212DFFB3385A18D4F38859ED311D0A9D5141CE9CC5C66EE689B266A8AA18ACE8282A0E0DB596C90B0A7B87"),
        ];

        for case in cases {
            let input_cursor = Cursor::new(input.clone());
            let result = match hash(input_cursor, case.0) {
                Ok(hash) => hash.to_upper_hex_string(),
                Err(e) => {
                    errors.push(format!(
                        "Hashing failed for {:?} with error: {:?}",
                        case.0, e
                    ));
                    continue;
                }
            };

            if result != case.1 {
                errors.push(format!(
                    "Mismatch for {:?}: expected {}, got {}",
                    case.0, case.1, result
                ));
            }
        }

        assert!(errors.is_empty(), "Errors encountered: {:?}", errors);
    }

    #[test]
    #[ignore] // it is too expensive.
    fn test_hash_extremely_long() {
        let mut errors = Vec::new();

        let input =
            b"abcdefghbcdefghicdefghijdefghijkefghijklfghijklmghijklmnhijklmno".repeat(16_777_216);
        let cases = &[
            (HashAlgorithm::Md5, "D338139169D50F55526194C790EC0448"),
            (HashAlgorithm::Sha1, "7789F0C9EF7BFC40D93311143DFBE69E2017F592"),
            (HashAlgorithm::Sha256, "ECBBC42CBF296603ACB2C6BC0410EF4378BAFB24B710357F12DF607758B33E2B"),
            (HashAlgorithm::Sha384, "A04296F4FCAAE14871BB5AD33E28DCF69238B04204D9941B8782E816D014BCB7540E4AF54F30D578F1A1CA2930847A12"),
            (HashAlgorithm::Sha512, "235FFD53504EF836A1342B488F483B396EABBFE642CF78EE0D31FEEC788B23D0D18D5C339550DD5958A500D4B95363DA1B5FA18AFFC1BAB2292DC63B7D85097C"),
        ];

        for case in cases {
            let input_cursor = Cursor::new(input.clone());
            let result = match hash(input_cursor, case.0) {
                Ok(hash) => hash.to_upper_hex_string(),
                Err(e) => {
                    errors.push(format!(
                        "Hashing failed for {:?} with error: {:?}",
                        case.0, e
                    ));
                    continue;
                }
            };

            if result != case.1 {
                errors.push(format!(
                    "Mismatch for {:?}: expected {}, got {}",
                    case.0, case.1, result
                ));
            }
        }

        assert!(errors.is_empty(), "Errors encountered: {:?}", errors);
    }
}
