//! Family presence-mask annex: a compact, read-only-after-build, per-slot
//! derived value used by the central-position SNP distance estimator (see
//! `docmd/theory/evolutionary_distances.md`, "Step 2b" and "Definitions:
//! family, and the canonical form of a family").
//!
//! One byte is stored per MPHF slot of a partition/layer, its low 4 bits
//! encoding a **presence mask** for the slot's k-mer's "family" (the up to 4
//! k-mers sharing the same flanks, differing only at the central base):
//! bit `b` (`b` = 0..3, in the fixed A/C/G/T = 0/1/2/3 encoding already used
//! for a single nucleotide) is set iff the family member whose *own* central
//! base — in its own canonical orientation — is `b`, is observed anywhere in
//! the current multi-genome index. This is a property of the whole index,
//! not of any one genome.
//!
//! Both facts the earlier (superseded) 3-bit design stored explicitly are
//! derived from the mask instead, not stored:
//! - sibling count = `popcount(mask) - 1`;
//! - minorant = regenerate the family's 4 canonical forms from the slot's
//!   own k-mer (`CanonicalKmerOf::central_canonical_neighbors`, cheap, no
//!   lookup), compare the raw encodings of whichever are set in the mask,
//!   take the smallest — see `obikindex::siblings`.
//!
//! Mask value 0 is logically unreachable as a real result (a slot's own base
//! is always present in its own family) and is reused as the "not yet
//! computed" sentinel: annex files are pre-initialised to all-zero, and a
//! real value is only ever written once, by the computation pass.
//!
//! Deliberately simpler than a true 4-bit pack (1 byte/slot instead of 4
//! bits/slot): correctness and simplicity first, for a first implementation.
//! Packing to 4 bits/slot is a pure storage-density follow-up, not a
//! behavioural change, left for later.

use std::fs::{File, OpenOptions};
use std::io;
use std::path::{Path, PathBuf};

use memmap2::{Mmap, MmapMut};

const MAGIC: [u8; 4] = *b"PSIB";

// Header: magic(4) + _pad(4) + n(8) = 16 bytes. Data (1 byte/slot) follows.
const HEADER_SIZE: usize = 16;

/// A family presence mask: bit `b` set iff the member whose own canonical
/// central base is `b` (0=A, 1=C, 2=G, 3=T) is observed in the index.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct FamilyMask(u8);

impl FamilyMask {
    /// The empty mask — never a valid *computed* result (a slot's own base
    /// is always present in its own family) — used only to build up a mask
    /// via repeated [`with`](Self::with) calls before storing it.
    pub const EMPTY: FamilyMask = FamilyMask(0);

    /// Set bit `base` (0=A, 1=C, 2=G, 3=T).
    #[inline]
    pub fn with(self, base: u8) -> Self {
        debug_assert!(base < 4, "base out of range: {base}");
        FamilyMask(self.0 | (1 << base))
    }

    /// Is the member with central base `base` (0..3) present?
    #[inline]
    pub fn has(self, base: u8) -> bool {
        debug_assert!(base < 4, "base out of range: {base}");
        self.0 & (1 << base) != 0
    }

    /// Number of family members observed anywhere in the index (1..=4).
    #[inline]
    pub fn family_size(self) -> u32 {
        self.0.count_ones()
    }

    /// Number of *other* members observed (0..=3) — `family_size() - 1`.
    #[inline]
    pub fn siblings(self) -> u32 {
        self.family_size() - 1
    }

    /// Raw bitmask (bit `b` = base `b` present) — for callers that build up
    /// a mask via their own bit operations (e.g. concurrently, via an
    /// `AtomicU8`) and only need the `FamilyMask` wrapper at the end.
    #[inline]
    pub fn bits(self) -> u8 {
        self.0
    }

    /// Construct from a raw bitmask (only the low 4 bits are kept).
    #[inline]
    pub fn from_bits(bits: u8) -> Self {
        FamilyMask(bits & 0b1111)
    }

    #[inline]
    fn encode(self) -> u8 {
        self.0
    }

    #[inline]
    fn decode(byte: u8) -> Option<Self> {
        if byte == 0 {
            // Unreachable for a real result — reserved as the "not yet
            // computed" sentinel.
            return None;
        }
        Some(FamilyMask(byte & 0b1111))
    }
}

// ── SiblingAnnex (reader) ───────────────────────────────────────────────────

pub struct SiblingAnnex {
    mmap: Mmap,
    n: usize,
    path: PathBuf,
}

impl SiblingAnnex {
    pub fn open(path: &Path) -> io::Result<Self> {
        let mmap = unsafe { Mmap::map(&File::open(path)?)? };
        if mmap.len() < HEADER_SIZE {
            return Err(io::Error::new(io::ErrorKind::InvalidData, "PSIB file too short"));
        }
        if mmap[0..4] != MAGIC {
            return Err(io::Error::new(io::ErrorKind::InvalidData, "bad PSIB magic"));
        }
        let n = u64::from_le_bytes(mmap[8..16].try_into().unwrap()) as usize;
        if mmap.len() < HEADER_SIZE + n {
            return Err(io::Error::new(io::ErrorKind::InvalidData, "PSIB file truncated"));
        }
        Ok(Self { mmap, n, path: path.to_path_buf() })
    }

    pub fn path(&self) -> &Path { &self.path }
    pub fn len(&self) -> usize { self.n }
    pub fn is_empty(&self) -> bool { self.n == 0 }

    /// `None` means the slot has not (yet) been computed — see module docs.
    pub fn get(&self, slot: usize) -> Option<FamilyMask> {
        FamilyMask::decode(self.mmap[HEADER_SIZE + slot])
    }
}

// ── SiblingAnnexBuilder (writer) ────────────────────────────────────────────

pub struct SiblingAnnexBuilder {
    mmap: MmapMut,
    n: usize,
    path: PathBuf,
}

impl SiblingAnnexBuilder {
    /// Create a new annex of `n` slots at `path`, pre-initialised to the
    /// "not yet computed" sentinel (all-zero).
    pub fn new(n: usize, path: &Path) -> io::Result<Self> {
        let file_size = HEADER_SIZE + n;
        let file = OpenOptions::new()
            .read(true).write(true).create(true).truncate(true)
            .open(path)?;
        file.set_len(file_size as u64)?;
        let mut mmap = unsafe { MmapMut::map_mut(&file)? };
        mmap[0..4].copy_from_slice(&MAGIC);
        mmap[4..8].copy_from_slice(&[0u8; 4]);
        mmap[8..16].copy_from_slice(&(n as u64).to_le_bytes());
        // Data region left at 0 by `set_len`/mmap — the sentinel value.
        Ok(Self { mmap, n, path: path.to_path_buf() })
    }

    pub fn len(&self) -> usize { self.n }
    pub fn is_empty(&self) -> bool { self.n == 0 }

    pub fn get(&self, slot: usize) -> Option<FamilyMask> {
        FamilyMask::decode(self.mmap[HEADER_SIZE + slot])
    }

    pub fn set(&mut self, slot: usize, mask: FamilyMask) {
        // Redundant concurrent writes from independent recomputation paths
        // converge to the same encoded byte for a given slot, so a plain
        // store here is safe even without external synchronisation, as long
        // as the byte write itself is atomic (true for a single aligned
        // byte on every platform this project targets).
        self.mmap[HEADER_SIZE + slot] = mask.encode();
    }

    pub fn close(self) -> io::Result<()> { self.mmap.flush() }

    pub fn finish(self) -> io::Result<SiblingAnnex> {
        let path = self.path.clone();
        self.close()?;
        SiblingAnnex::open(&path)
    }
}

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

    #[test]
    fn sentinel_is_zero_and_unset_slots_read_as_uncomputed() {
        let dir = tempdir().unwrap();
        let path = dir.path().join("test.psib");
        let builder = SiblingAnnexBuilder::new(4, &path).unwrap();
        for slot in 0..4 {
            assert_eq!(builder.get(slot), None);
        }
        builder.close().unwrap();
    }

    #[test]
    fn roundtrip_all_valid_masks() {
        let dir = tempdir().unwrap();
        let path = dir.path().join("test.psib");
        let mut builder = SiblingAnnexBuilder::new(4, &path).unwrap();

        let masks = [
            FamilyMask::EMPTY.with(0),                   // just A: family size 1
            FamilyMask::EMPTY.with(0).with(3),            // A + T: size 2
            FamilyMask::EMPTY.with(1).with(2).with(3),    // C+G+T: size 3
            FamilyMask::EMPTY.with(0).with(1).with(2).with(3), // all 4
        ];
        for (slot, mask) in masks.iter().enumerate() {
            builder.set(slot, *mask);
        }
        let annex = builder.finish().unwrap();
        for (slot, mask) in masks.iter().enumerate() {
            assert_eq!(annex.get(slot), Some(*mask));
        }
        assert_eq!(annex.get(0).unwrap().siblings(), 0);
        assert_eq!(annex.get(1).unwrap().siblings(), 1);
        assert_eq!(annex.get(2).unwrap().siblings(), 2);
        assert_eq!(annex.get(3).unwrap().siblings(), 3);
        assert_eq!(annex.get(3).unwrap().family_size(), 4);
    }

    #[test]
    fn has_reflects_individual_bits() {
        let mask = FamilyMask::EMPTY.with(0).with(2);
        assert!(mask.has(0));
        assert!(!mask.has(1));
        assert!(mask.has(2));
        assert!(!mask.has(3));
    }
}
