//! 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());

    }

}