pandora_lib_promethion/src/uniqueness/cache.rs

use std::fs::{self, File};
use std::io::{self, BufReader, BufWriter, Read, Write};
use std::path::Path;
use std::time::SystemTime;

use memmap2::Mmap;

use super::UniqueLen;

/// Magic + version guard
const MAGIC: &[u8; 8] = b"UNIQIDX1";
const VERSION: u32 = 2;
const HEADER_LEN: usize = 20;

/// Identifies a cache entry unambiguously.
/// Derive from FASTA path + contig name + genome length + content hash.
pub struct CacheKey {
    pub contig: String,
    pub genome_len: usize,
    /// Blake3 hex hash of the raw genome bytes (first 16 hex chars sufficient)
    pub hash: String,
}

impl CacheKey {
    /// Compute from raw genome bytes + contig name.
    pub fn from_genome(contig: &str, genome: &[u8]) -> Self {
        // Blake3 is fast (~3 GB/s), ideal for large genomes.
        // Use `blake3` crate or fall back to a simple FNV for no-dep builds.
        let hash = blake3_hex_short(genome);
        Self {
            contig: contig.to_owned(),
            genome_len: genome.len(),
            hash,
        }
    }

    /// Compute a genome-wide cache key without reading the FASTA sequence.
    ///
    /// `metadata` should include all selected contig names/lengths and any
    /// indexing parameters that affect the cached vectors.
    pub fn from_genome_metadata(
        contig: &str,
        genome_len: usize,
        fasta_path: &Path,
        metadata: &[u8],
    ) -> io::Result<Self> {
        let meta = std::fs::metadata(fasta_path)?;
        let mtime = meta
            .modified()?
            .duration_since(SystemTime::UNIX_EPOCH)
            .map(|d| d.as_secs())
            .unwrap_or(0);
        let fsize = meta.len();

        let mut hash_input = Vec::with_capacity(metadata.len() + 32);
        hash_input.extend_from_slice(&mtime.to_le_bytes());
        hash_input.extend_from_slice(&fsize.to_le_bytes());
        hash_input.extend_from_slice(metadata);

        Ok(Self {
            contig: contig.to_owned(),
            genome_len,
            hash: blake3_hex_short(&hash_input),
        })
    }

    /// Construit la clé depuis le .fai (noodles) — ne lit pas le FASTA.
    /// `fai_entry` contient déjà la longueur du contig.
    pub fn from_fai(contig: &str, fai_len: usize, fasta_path: &Path) -> io::Result<Self> {
        let meta = std::fs::metadata(fasta_path)?;
        let mtime = meta
            .modified()?
            .duration_since(SystemTime::UNIX_EPOCH)
            .map(|d| d.as_secs())
            .unwrap_or(0);
        let fsize = meta.len();

        // Hash = mtime + fsize — suffisant pour invalider sur modification
        let hash = format!("{mtime:016x}{fsize:016x}");

        Ok(Self {
            contig: contig.to_owned(),
            genome_len: fai_len,
            hash,
        })
    }

    pub fn file_name(&self) -> String {
        format!(
            "{}__{}__{}.uniqidx",
            self.contig, self.genome_len, self.hash
        )
    }
}

/// Binary format (little-endian):
///
/// [0..8]   magic   = b"UNIQIDX1"
/// [8..12]  version = 2u32
/// [12..20] n       = genome_len as u64
/// [20..]   min_unique_len: n × u32  (`NOT_UNIQUE` means no unique interval)
pub fn save_to_cache(path: &Path, min_unique_len: &[UniqueLen]) -> io::Result<()> {
    let tmp = path.with_extension("uniqidx.tmp");
    {
        let f = File::create(&tmp)?;
        let mut w = BufWriter::with_capacity(8 * 1024 * 1024, f);

        w.write_all(MAGIC)?;
        w.write_all(&VERSION.to_le_bytes())?;
        w.write_all(&(min_unique_len.len() as u64).to_le_bytes())?;

        // Bulk-write as u32 array — avoids per-element call overhead.
        let buf: Vec<u8> = min_unique_len
            .iter()
            .flat_map(|&v| v.to_le_bytes())
            .collect();
        w.write_all(&buf)?;
        w.flush()?;
    }
    // Atomic rename: avoids corrupt cache on crash mid-write
    fs::rename(&tmp, path)?;
    Ok(())
}

pub fn load_from_cache(path: &Path, expected_len: usize) -> io::Result<Vec<UniqueLen>> {
    let f = File::open(path)?;
    let mut r = BufReader::with_capacity(8 * 1024 * 1024, f);

    // Magic
    let mut magic = [0u8; 8];
    r.read_exact(&mut magic)?;
    if &magic != MAGIC {
        return Err(io::Error::new(io::ErrorKind::InvalidData, "bad magic"));
    }

    // Version
    let mut vbuf = [0u8; 4];
    r.read_exact(&mut vbuf)?;
    if u32::from_le_bytes(vbuf) != VERSION {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            "version mismatch",
        ));
    }

    // Length
    let mut nbuf = [0u8; 8];
    r.read_exact(&mut nbuf)?;
    let n = u64::from_le_bytes(nbuf) as usize;
    if n != expected_len {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            format!("length mismatch: cache has {n}, genome is {expected_len}"),
        ));
    }

    // Payload: read all at once then cast
    let byte_len = n * 4;
    let mut buf = vec![0u8; byte_len];
    r.read_exact(&mut buf)?;

    let result: Vec<UniqueLen> = buf
        .chunks_exact(4)
        .map(|c| u32::from_le_bytes(c.try_into().unwrap()))
        .collect();

    Ok(result)
}

pub(crate) struct MappedUniqueLenCache {
    mmap: Mmap,
    len: usize,
}

impl MappedUniqueLenCache {
    pub(crate) fn get(&self, index: usize) -> UniqueLen {
        assert!(index < self.len);
        let offset = HEADER_LEN + index * 4;
        UniqueLen::from_le_bytes(
            self.mmap[offset..offset + 4]
                .try_into()
                .expect("cache payload offset should be in bounds"),
        )
    }

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

pub(crate) fn mmap_from_cache(
    path: &Path,
    expected_len: usize,
) -> io::Result<MappedUniqueLenCache> {
    let f = File::open(path)?;
    let file_len = f.metadata()?.len() as usize;
    let expected_file_len =
        HEADER_LEN
            .checked_add(expected_len.checked_mul(4).ok_or_else(|| {
                io::Error::new(io::ErrorKind::InvalidInput, "cache length overflow")
            })?)
            .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "cache length overflow"))?;
    if file_len != expected_file_len {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            format!("file length mismatch: cache has {file_len}, expected {expected_file_len}"),
        ));
    }

    // SAFETY: The mapping is read-only and the file is never mutated through
    // this handle. Cache files are written atomically before being mapped.
    let mmap = unsafe { Mmap::map(&f)? };
    validate_header(&mmap[..HEADER_LEN], expected_len)?;
    Ok(MappedUniqueLenCache {
        mmap,
        len: expected_len,
    })
}

fn validate_header(header: &[u8], expected_len: usize) -> io::Result<()> {
    if &header[0..8] != MAGIC {
        return Err(io::Error::new(io::ErrorKind::InvalidData, "bad magic"));
    }

    let version = u32::from_le_bytes(header[8..12].try_into().unwrap());
    if version != VERSION {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            "version mismatch",
        ));
    }

    let n = u64::from_le_bytes(header[12..20].try_into().unwrap()) as usize;
    if n != expected_len {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            format!("length mismatch: cache has {n}, genome is {expected_len}"),
        ));
    }

    Ok(())
}

// ── Minimal Blake3-like hash (FNV-1a 64bit, no deps) ─────────────────────────
// Replace with `blake3` crate for collision resistance on production.
fn blake3_hex_short(data: &[u8]) -> String {
    const FNV_OFFSET: u64 = 14695981039346656037;
    const FNV_PRIME: u64 = 1099511628211;
    let mut h = FNV_OFFSET;
    for &b in data {
        h ^= b as u64;
        h = h.wrapping_mul(FNV_PRIME);
    }
    format!("{h:016x}")
}