//! Images by hash (oxygenflow#39, docs/jetstream.md §8).
//!
//! A rootfs image is named by the BLAKE3 of its bytes, as pipelines pin it
//! (`cas://blob/HASH`). To move it, it is cut into chunks with Oxygen's
//! FastCDC (32 KiB to 256 KiB, `oxy_core::chunk`), each named by its own
//! BLAKE3, and described by a [`Manifest`]: the image's hash, size and
//! architecture, and its chunks in order. A sender sends only the chunks a
//! receiver lacks; a receiver writes the image back, leaving runs of zeros
//! as holes without fetching them, and checks the whole hash before using
//! it. Chunk boundaries depend only on content, so an image rebuilt with a
//! small change shares nearly all its chunks with the one before.

use oxy_core::chunk::{cut_point, MAX_SIZE};
use serde::{Deserialize, Serialize};
use std::collections::HashSet;
use std::io::{Read, Seek, SeekFrom, Write};
use std::path::Path;

/// The manifest format this crate writes and reads.
pub const VERSION: u32 = 1;
/// Largest image accepted: 64 GiB.
pub const MAX_IMAGE: u64 = 64 << 30;

#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("{0}")]
    Io(#[from] std::io::Error),
    #[error("{0}")]
    Bad(String),
}

pub type Result<T> = std::result::Result<T, Error>;

/// One chunk: its BLAKE3 (hex) and length.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct Chunk(pub String, pub u32);

impl Chunk {
    pub fn hash(&self) -> &str {
        &self.0
    }

    pub fn len(&self) -> usize {
        self.1 as usize
    }

    pub fn is_empty(&self) -> bool {
        self.1 == 0
    }

    /// Whether this is a chunk of zeros, which a receiver writes as a hole
    /// without fetching it.
    pub fn is_zeros(&self) -> bool {
        self.0 == zeros_hash(self.len())
    }
}

/// An image: its hash, size, architecture and chunks.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct Manifest {
    pub version: u32,
    /// BLAKE3 of the whole image, hex: its name.
    pub hash: String,
    pub size: u64,
    /// `x86_64`, `riscv64`, …, as its import recorded (OxygenFlow spec §11.2).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub arch: Option<String>,
    pub chunks: Vec<Chunk>,
}

impl Manifest {
    /// The distinct chunks a receiver must fetch: all but zeros.
    pub fn needed(&self) -> Vec<&str> {
        let mut seen = HashSet::new();
        self.chunks.iter().filter(|c| !c.is_zeros() && seen.insert(c.hash())).map(|c| c.hash()).collect()
    }

    /// Check its shape: a known version, hex names, chunks in the
    /// chunker's bounds that add up to the size.
    pub fn check(&self) -> Result<()> {
        if self.version != VERSION {
            return Err(Error::Bad(format!("image manifest version {} is not known; this reads {VERSION}", self.version)));
        }
        if !is_hash(&self.hash) {
            return Err(Error::Bad(format!("{:?} is not an image hash", self.hash)));
        }
        if self.size > MAX_IMAGE {
            return Err(Error::Bad(format!("the image is {} GiB; at most {} GiB", self.size >> 30, MAX_IMAGE >> 30)));
        }
        let mut total = 0u64;
        for c in &self.chunks {
            if !is_hash(c.hash()) || c.is_empty() || c.len() > MAX_SIZE {
                return Err(Error::Bad(format!("chunk {:?} of {} bytes is not a chunk", c.hash(), c.len())));
            }
            total += c.1 as u64;
        }
        if total != self.size {
            return Err(Error::Bad(format!("the chunks hold {total} bytes; the image is {}", self.size)));
        }
        Ok(())
    }
}

/// Whether `s` is 64 lower-case hex digits.
pub fn is_hash(s: &str) -> bool {
    s.len() == 64 && s.bytes().all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b))
}

/// The BLAKE3 of `len` zero bytes, hex; remembered per length.
pub fn zeros_hash(len: usize) -> String {
    use std::sync::Mutex;
    static MEMO: Mutex<Vec<(usize, String)>> = Mutex::new(Vec::new());
    let mut memo = MEMO.lock().unwrap_or_else(|p| p.into_inner());
    if let Some((_, h)) = memo.iter().find(|(l, _)| *l == len) {
        return h.clone();
    }
    let h = blake3::hash(&vec![0u8; len]).to_hex().to_string();
    if memo.len() < 64 {
        memo.push((len, h.clone()));
    }
    h
}

/// Cut `r` into chunks as Oxygen's chunker would cut it whole, calling
/// `each` with every chunk's bytes, and answer the manifest. Reads in
/// bounded memory: the chunker never looks past [`MAX_SIZE`].
pub fn scan<R: Read>(mut r: R, arch: Option<String>, mut each: impl FnMut(&Chunk, &[u8]) -> Result<()>) -> Result<Manifest> {
    const FILL: usize = 8 << 20;
    let mut buf: Vec<u8> = Vec::with_capacity(FILL + MAX_SIZE);
    let (mut pos, mut eof, mut size) = (0usize, false, 0u64);
    let mut whole = blake3::Hasher::new();
    let mut chunks = Vec::new();
    loop {
        if !eof && buf.len() - pos < MAX_SIZE {
            buf.drain(..pos);
            pos = 0;
            let start = buf.len();
            buf.resize(start + FILL, 0);
            let mut got = 0;
            while got < FILL {
                match r.read(&mut buf[start + got..])? {
                    0 => {
                        eof = true;
                        break;
                    }
                    n => got += n,
                }
            }
            buf.truncate(start + got);
        }
        if pos == buf.len() {
            break;
        }
        let n = cut_point(&buf[pos..]);
        let bytes = &buf[pos..pos + n];
        whole.update(bytes);
        let c = Chunk(blake3::hash(bytes).to_hex().to_string(), n as u32);
        each(&c, bytes)?;
        chunks.push(c);
        size += n as u64;
        pos += n;
    }
    Ok(Manifest { version: VERSION, hash: whole.finalize().to_hex().to_string(), size, arch, chunks })
}

/// [`scan`] of a file.
pub fn scan_file(path: &Path, arch: Option<String>, each: impl FnMut(&Chunk, &[u8]) -> Result<()>) -> Result<Manifest> {
    scan(std::io::BufReader::with_capacity(1 << 20, std::fs::File::open(path)?), arch, each)
}

/// Write image `m` to `out`, getting each chunk's bytes from `fetch`
/// (zeros are holes, never fetched; a chunk seen before is fetched once
/// per run of repeats), checking every chunk and then the whole hash. The
/// image is written beside `out` and renamed into place only when whole.
pub fn assemble(m: &Manifest, out: &Path, mut fetch: impl FnMut(&Chunk) -> Result<Vec<u8>>) -> Result<()> {
    m.check()?;
    let tmp = out.with_extension("part");
    let mut f = std::fs::File::create(&tmp)?;
    let mut whole = blake3::Hasher::new();
    let mut last: Option<(String, Vec<u8>)> = None;
    for c in &m.chunks {
        if c.is_zeros() {
            whole.update(&vec![0u8; c.len()]);
            f.seek(SeekFrom::Current(c.len() as i64))?;
            continue;
        }
        let bytes = match &last {
            Some((h, b)) if h == c.hash() => b.clone(),
            _ => {
                let b = fetch(c)?;
                if b.len() != c.len() || blake3::hash(&b).to_hex().as_str() != c.hash() {
                    let _ = std::fs::remove_file(&tmp);
                    return Err(Error::Bad(format!("chunk {} came back with other bytes", &c.hash()[..12])));
                }
                last = Some((c.0.clone(), b.clone()));
                b
            }
        };
        whole.update(&bytes);
        f.write_all(&bytes)?;
    }
    f.set_len(m.size)?;
    f.sync_all()?;
    drop(f);
    let got = whole.finalize().to_hex().to_string();
    if got != m.hash {
        let _ = std::fs::remove_file(&tmp);
        return Err(Error::Bad(format!("the image came out as {}, not {}", &got[..12], &m.hash[..12])));
    }
    std::fs::rename(&tmp, out)?;
    Ok(())
}

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

    fn image(seed: u8) -> Vec<u8> {
        // Some text-like bytes, a long run of zeros (as an ext4 image has),
        // then more bytes.
        let mut v: Vec<u8> = (0..300_000u32).map(|i| ((i * 7 + seed as u32) % 251) as u8).collect();
        v.extend(std::iter::repeat_n(0u8, 3 << 20));
        v.extend((0..200_000u32).map(|i| ((i * 13 + seed as u32) % 241) as u8));
        v
    }

    #[test]
    fn a_scanned_image_comes_back_whole_and_zeros_are_never_fetched() {
        let data = image(1);
        let mut stored = std::collections::HashMap::new();
        let m = scan(&data[..], Some("x86_64".into()), |c, b| {
            stored.insert(c.hash().to_string(), b.to_vec());
            Ok(())
        })
        .unwrap();
        m.check().unwrap();
        assert_eq!(m.hash, blake3::hash(&data).to_hex().to_string(), "named as pipelines pin it");
        assert_eq!(m.size, data.len() as u64);
        assert_eq!(
            m.chunks.iter().map(|c| c.len()).collect::<Vec<_>>(),
            oxy_core::chunk::chunks(&data).iter().map(|c| c.len()).collect::<Vec<_>>(),
            "the same cuts as Oxygen's chunker over the whole"
        );
        assert!(m.chunks.iter().any(|c| c.is_zeros()));
        let dir = tempfile::tempdir().unwrap();
        let out = dir.path().join("img.ext4");
        let mut fetched = 0;
        assemble(&m, &out, |c| {
            fetched += 1;
            assert!(!c.is_zeros());
            Ok(stored[c.hash()].clone())
        })
        .unwrap();
        assert_eq!(std::fs::read(&out).unwrap(), data);
        assert!(fetched <= m.needed().len() + 1);

        // A rebuild with a small change shares nearly every chunk.
        let mut changed = data.clone();
        changed[100] ^= 0xff;
        let m2 = scan(&changed[..], None, |_, _| Ok(())).unwrap();
        let have: HashSet<&str> = m.chunks.iter().map(|c| c.hash()).collect();
        let new = m2.needed().into_iter().filter(|h| !have.contains(h)).count();
        assert!(new <= 2, "{new} new chunks for a one-byte change");
    }

    #[test]
    fn bad_bytes_and_bad_manifests_are_refused() {
        let data = image(2);
        let m = scan(&data[..], None, |_, _| Ok(())).unwrap();
        let dir = tempfile::tempdir().unwrap();
        let out = dir.path().join("img.ext4");
        let e = assemble(&m, &out, |c| Ok(vec![1u8; c.len()])).unwrap_err();
        assert!(e.to_string().contains("other bytes"), "{e}");
        assert!(!out.exists() && !out.with_extension("part").exists());
        let mut short = m.clone();
        short.size += 1;
        assert!(short.check().unwrap_err().to_string().contains("chunks hold"));
        let mut odd = m.clone();
        odd.hash = "nope".into();
        assert!(odd.check().is_err());
    }
}