use sha2::{Digest, Sha256};

/// Hex-encoded independent SHA-256 digest of transferred bytes.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct IntegrityDigest(String);

impl IntegrityDigest {
    #[must_use]
    pub fn as_str(&self) -> &str {
        &self.0
    }
}

impl std::fmt::Display for IntegrityDigest {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.0.fmt(f)
    }
}

/// Running integrity state for one transfer member.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct IntegrityState {
    bytes_hashed: u64,
    digest: Option<IntegrityDigest>,
}

impl IntegrityState {
    #[must_use]
    pub fn bytes_hashed(&self) -> u64 {
        self.bytes_hashed
    }

    #[must_use]
    pub fn digest(&self) -> Option<&IntegrityDigest> {
        self.digest.as_ref()
    }
}

/// Independent SHA-256 accountant updated as bytes leave the copy loop.
#[derive(Clone)]
pub struct Sha256Accountant {
    hasher: Sha256,
    bytes: u64,
}

impl Default for Sha256Accountant {
    fn default() -> Self {
        Self::new()
    }
}

impl Sha256Accountant {
    #[must_use]
    pub fn new() -> Self {
        Self {
            hasher: Sha256::new(),
            bytes: 0,
        }
    }

    #[must_use]
    pub fn bytes_hashed(&self) -> u64 {
        self.bytes
    }

    pub fn update(&mut self, data: &[u8]) {
        if data.is_empty() {
            return;
        }
        self.hasher.update(data);
        self.bytes = self.bytes.saturating_add(data.len() as u64);
    }

    #[must_use]
    pub fn snapshot(&self) -> IntegrityState {
        let digest = {
            let hasher = self.hasher.clone();
            let finalized = hasher.finalize();
            IntegrityDigest(format_sha256(finalized.as_ref()))
        };
        IntegrityState {
            bytes_hashed: self.bytes,
            digest: Some(digest),
        }
    }

    #[must_use]
    pub fn finalize(self) -> IntegrityState {
        let digest = IntegrityDigest(format_sha256(self.hasher.finalize().as_ref()));
        IntegrityState {
            bytes_hashed: self.bytes,
            digest: Some(digest),
        }
    }
}

fn format_sha256(bytes: &[u8]) -> String {
    let mut out = String::with_capacity(7 + bytes.len() * 2);
    out.push_str("sha256:");
    for byte in bytes {
        use std::fmt::Write as _;
        let _ = write!(out, "{byte:02x}");
    }
    out
}

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

    #[test]
    fn empty_digest_matches_sha256_of_empty() {
        let state = Sha256Accountant::new().finalize();
        assert_eq!(state.bytes_hashed(), 0);
        assert_eq!(
            state.digest().map(IntegrityDigest::as_str),
            Some("sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855")
        );
    }

    #[test]
    fn known_vector_abc() {
        let mut accountant = Sha256Accountant::new();
        accountant.update(b"abc");
        let state = accountant.finalize();
        assert_eq!(state.bytes_hashed(), 3);
        assert_eq!(
            state.digest().map(IntegrityDigest::as_str),
            Some("sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad")
        );
    }
}