//! Object store (spec §13.1, §13.3, §13.4): loose objects, one file per

//! object under `objects/loose/<hh>/<rest>`, and packs, `pack-<id>.pack`

//! with a sorted `pack-<id>.idx`, holding zstd-compressed objects. Reads

//! check loose objects first, then every pack; every read re-hashes.

//!

//! Pack file: `OXYPACK1` then entries of

//! `type u8 | flags u8 (bit 0: zstd) | stored_len u32 | raw_len u32 | ref 32 | body`.

//! Index file: `OXYIDX01` then `count u32` then sorted entries of

//! `ref 32 | offset u64 | stored_len u32 | raw_len u32 | type u8 | flags u8`.

use crate::attest::Attest;

use crate::hash::{hash_object, header, ObjectRef, ObjectType};

use crate::identity::{Identity, Signature};

use crate::objects::{Change, ObjectError, StateObject};

use rand::RngCore;

use std::cell::RefCell;

use std::fs;

use std::io::{self, Read, Seek, SeekFrom, Write};

use std::path::{Path, PathBuf};

const PACK_MAGIC: &[u8; 8] = b"OXYPACK1";

const IDX_MAGIC: &[u8; 8] = b"OXYIDX01";

const IDX_ENTRY: usize = 32 + 8 + 4 + 4 + 1 + 1;

const FLAG_ZSTD: u8 = 1;

const ZSTD_LEVEL: i32 = 3;

/// Largest object a pack entry may claim; anything bigger is corruption,

/// and refusing it keeps a bad index from driving huge allocations.

const MAX_OBJECT_LEN: u32 = 1 << 30;

#[derive(Clone, Copy, Debug)]

struct PackEntry {

    offset: u64,

    stored_len: u32,

    raw_len: u32,

    ty: ObjectType,

    flags: u8,

}

/// One pack's index, fully in memory and sorted by ObjectRef.

#[derive(Debug)]

struct PackIndex {

    path: PathBuf,

    refs: Vec<ObjectRef>,

    entries: Vec<PackEntry>,

}

impl PackIndex {

    fn load(idx_path: &Path, pack_path: &Path) -> Result<PackIndex> {

        let data = fs::read(idx_path)?;

        let bad = |m: &str| StoreError::Corrupt(ObjectRef::ZERO, format!("{}: {m}", idx_path.display()));

        if data.len() < 12 || &data[..8] != IDX_MAGIC {

            return Err(bad("bad index header"));

        }

        let count = u32::from_le_bytes(data[8..12].try_into().unwrap()) as usize;

        if data.len() != 12 + count * IDX_ENTRY {

            return Err(bad("index length mismatch"));

        }

        let pack_len = fs::metadata(pack_path)?.len();

        let mut refs = Vec::with_capacity(count);

        let mut entries = Vec::with_capacity(count);

        for i in 0..count {

            let e = &data[12 + i * IDX_ENTRY..12 + (i + 1) * IDX_ENTRY];

            let r = ObjectRef::from_slice(&e[..32]).unwrap();

            if let Some(prev) = refs.last() {

                if r <= *prev {

                    return Err(bad("index not sorted"));

                }

            }

            refs.push(r);

            let entry = PackEntry {

                offset: u64::from_le_bytes(e[32..40].try_into().unwrap()),

                stored_len: u32::from_le_bytes(e[40..44].try_into().unwrap()),

                raw_len: u32::from_le_bytes(e[44..48].try_into().unwrap()),

                ty: ObjectType::ALL.get(e[48] as usize).copied().ok_or_else(|| bad("unknown object type"))?,

                flags: e[49],

            };

            if entry.stored_len > MAX_OBJECT_LEN || entry.raw_len > MAX_OBJECT_LEN || entry.offset.saturating_add(entry.stored_len as u64) > pack_len {

                return Err(bad("index entry points outside the pack"));

            }

            entries.push(entry);

        }

        Ok(PackIndex { path: pack_path.to_path_buf(), refs, entries })

    }

    fn find(&self, r: &ObjectRef) -> Option<PackEntry> {

        self.refs.binary_search(r).ok().map(|i| self.entries[i])

    }

    fn read(&self, r: &ObjectRef, e: PackEntry) -> Result<(ObjectType, Vec<u8>)> {

        let mut f = fs::File::open(&self.path)?;

        f.seek(SeekFrom::Start(e.offset))?;

        let mut stored = vec![0u8; e.stored_len as usize];

        f.read_exact(&mut stored)?;

        let body = if e.flags & FLAG_ZSTD != 0 {

            zstd::bulk::decompress(&stored, e.raw_len as usize).map_err(|err| StoreError::Corrupt(*r, format!("zstd: {err}")))?

        } else {

            stored

        };

        if body.len() != e.raw_len as usize || hash_object(e.ty, &body) != *r {

            return Err(StoreError::Corrupt(*r, "packed object does not match its reference".into()));

        }

        Ok((e.ty, body))

    }

}

fn type_index(ty: ObjectType) -> u8 {

    ObjectType::ALL.iter().position(|t| *t == ty).unwrap() as u8

}

/// What `pack_loose` did.

#[derive(Debug, Clone, Default, PartialEq, Eq)]

pub struct PackReport {

    pub objects: usize,

    pub raw_bytes: u64,

    pub packed_bytes: u64,

}

/// Packs tolerated before `merge_packs` folds them into one (spec §13.4).

pub const MAX_PACKS: usize = 16;

/// What `prune` did.

#[derive(Debug, Clone, Default, PartialEq, Eq)]

pub struct PruneReport {

    pub objects: usize,

    pub bytes: u64,

    /// Unreachable objects left alone because they are younger than the grace period.

    pub deferred: usize,

}

#[derive(Debug, thiserror::Error)]

pub enum StoreError {

    #[error("io: {0}")]

    Io(#[from] io::Error),

    #[error("corrupt object {0}: {1}")]

    Corrupt(ObjectRef, String),

    #[error("object {0}")]

    Object(#[from] ObjectError),

    #[error("object {0} is a {1}, expected {2}")]

    WrongType(ObjectRef, &'static str, &'static str),

    #[error("object {0} not found")]

    NotFound(ObjectRef),

}

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

#[derive(Debug)]

pub struct Store {

    dir: PathBuf,

    packs: RefCell<Vec<PackIndex>>,

    /// An open write batch (spec §13.4): objects go here instead of to loose

    /// files, and are written out with one `fsync` when the batch finishes.

    pending: RefCell<Option<Pending>>,

    batch_depth: std::cell::Cell<u32>,

}

/// Objects written since `begin_batch`. Small batches stay in memory and

/// become loose objects; once `SPILL_OBJECTS` or `SPILL_BYTES` is exceeded

/// the batch streams into a pack file instead.

#[derive(Debug)]

struct Pending {

    buffered: Vec<(ObjectType, Vec<u8>)>,

    buffered_bytes: usize,

    lookup: std::collections::HashMap<ObjectRef, PendingEntry>,

    pack: Option<PendingPack>,

    report: PackReport,

}

#[derive(Debug, Clone, Copy)]

enum PendingEntry {

    Buffered(usize),

    Packed(PackEntry),

}

#[derive(Debug)]

struct PendingPack {

    id: String,

    tmp_path: PathBuf,

    writer: io::BufWriter<fs::File>,

    offset: u64,

    index: Vec<(ObjectRef, PackEntry)>,

}

const SPILL_OBJECTS: usize = 256;

const SPILL_BYTES: usize = 8 << 20;

/// Finishes the enclosing batch when dropped; call `finish` to see errors.

pub struct Batch<'a> {

    store: &'a Store,

    done: bool,

}

impl Batch<'_> {

    pub fn finish(mut self) -> Result<()> {

        self.done = true;

        self.store.finish_batch()

    }

}

impl Drop for Batch<'_> {

    fn drop(&mut self) {

        if !self.done {

            let _ = self.store.finish_batch();

        }

    }

}

impl Store {

    pub fn open(dir: impl Into<PathBuf>) -> io::Result<Store> {

        let dir = dir.into();

        fs::create_dir_all(dir.join("loose"))?;

        let store = Store { dir, packs: RefCell::new(Vec::new()), pending: RefCell::new(None), batch_depth: std::cell::Cell::new(0) };

        store.reload_packs().map_err(|e| io::Error::other(e.to_string()))?;

        Ok(store)

    }

    fn reload_packs(&self) -> Result<()> {

        let mut packs = Vec::new();

        let mut names: Vec<PathBuf> = fs::read_dir(&self.dir)?

            .filter_map(|e| e.ok().map(|e| e.path()))

            .filter(|p| p.extension().is_some_and(|x| x == "idx") && p.file_name().is_some_and(|n| n.to_string_lossy().starts_with("pack-")))

            .collect();

        names.sort();

        for idx in names {

            let pack = idx.with_extension("pack");

            if pack.is_file() {

                packs.push(PackIndex::load(&idx, &pack)?);

            }

        }

        *self.packs.borrow_mut() = packs;

        Ok(())

    }

    pub fn pack_count(&self) -> usize {

        self.packs.borrow().len()

    }

    fn find_packed(&self, r: &ObjectRef) -> Option<(usize, PackEntry)> {

        self.packs.borrow().iter().enumerate().find_map(|(i, p)| p.find(r).map(|e| (i, e)))

    }

    /// Loose objects on disk, with their paths.

    fn loose_objects(&self) -> io::Result<Vec<(ObjectRef, PathBuf)>> {

        let mut out = Vec::new();

        let loose = self.dir.join("loose");

        for d in fs::read_dir(&loose)? {

            let d = d?;

            if !d.file_type()?.is_dir() {

                continue;

            }

            let prefix = d.file_name().to_string_lossy().to_string();

            for f in fs::read_dir(d.path())? {

                let f = f?;

                let name = f.file_name().to_string_lossy().to_string();

                if let Some(r) = ObjectRef::from_hex(&format!("{prefix}{name}")) {

                    out.push((r, f.path()));

                }

            }

        }

        out.sort();

        Ok(out)

    }

    pub fn loose_count(&self) -> io::Result<usize> {

        Ok(self.loose_objects()?.len())

    }

    /// Move every loose object into a new pack (spec §13.4). Loose files are

    /// deleted only after the pack and its index are durable.

    pub fn pack_loose(&self) -> Result<PackReport> {

        let loose = self.loose_objects()?;

        self.write_pack(&loose.iter().map(|(r, _)| *r).collect::<Vec<_>>(), &loose.iter().map(|(_, p)| p.clone()).collect::<Vec<_>>(), &[])

    }

    /// Fold every pack into one when there are more than `MAX_PACKS`

    /// (spec §13.4). Returns how many packs were merged.

    pub fn merge_packs(&self) -> Result<usize> {

        let (refs, paths): (Vec<ObjectRef>, Vec<PathBuf>) = {

            let packs = self.packs.borrow();

            if packs.len() <= MAX_PACKS {

                return Ok(0);

            }

            (packs.iter().flat_map(|p| p.refs.iter().copied()).collect(), packs.iter().map(|p| p.path.clone()).collect())

        };

        let n = paths.len();

        self.write_pack(&refs, &[], &paths)?;

        Ok(n)

    }

    /// Delete every object not in `keep` (spec §21). Loose objects and packs

    /// younger than `grace` are left alone, since a concurrent writer may not

    /// have recorded its references yet. Packs holding unreachable objects

    /// are rewritten without them. Signature sidecars and seal-index entries

    /// of deleted objects go too.

    pub fn prune(&self, keep: &std::collections::HashSet<ObjectRef>, grace: std::time::Duration) -> Result<PruneReport> {

        let now = std::time::SystemTime::now();

        let old_enough = |m: &fs::Metadata| m.modified().map(|t| now.duration_since(t).unwrap_or_default() >= grace).unwrap_or(true);

        let mut report = PruneReport::default();

        // Loose objects.

        for (r, path) in self.loose_objects()? {

            if keep.contains(&r) {

                continue;

            }

            let meta = fs::metadata(&path)?;

            if !old_enough(&meta) {

                report.deferred += 1;

                continue;

            }

            report.bytes += meta.len();

            report.objects += 1;

            fs::remove_file(&path)?;

            let _ = fs::remove_dir(path.parent().unwrap());

            self.remove_sidecars(&r)?;

        }

        // Packs: rewrite those that hold unreachable objects.

        let mut keep_refs: Vec<ObjectRef> = Vec::new();

        let mut old_paths: Vec<PathBuf> = Vec::new();

        let mut rewrite = false;

        for p in self.packs.borrow().iter() {

            let meta = fs::metadata(&p.path)?;

            let unreachable = p.refs.iter().filter(|r| !keep.contains(r)).count();

            if unreachable == 0 {

                continue;

            }

            if !old_enough(&meta) {

                report.deferred += unreachable;

                continue;

            }

            rewrite = true;

            old_paths.push(p.path.clone());

            for (r, e) in p.refs.iter().zip(p.entries.iter()) {

                if keep.contains(r) {

                    keep_refs.push(*r);

                } else {

                    report.objects += 1;

                    report.bytes += e.stored_len as u64;

                }

            }

        }

        if rewrite {

            let dropped: Vec<ObjectRef> = {

                let packs = self.packs.borrow();

                packs.iter().filter(|p| old_paths.contains(&p.path)).flat_map(|p| p.refs.iter().copied()).filter(|r| !keep.contains(r)).collect()

            };

            self.write_pack(&keep_refs, &[], &old_paths)?;

            for r in dropped {

                self.remove_sidecars(&r)?;

            }

        }

        Ok(report)

    }

    fn remove_sidecars(&self, r: &ObjectRef) -> Result<()> {

        for path in [self.sig_path(r), self.seal_index_path(r)] {

            match fs::remove_file(&path) {

                Ok(()) => {

                    let _ = fs::remove_dir(path.parent().unwrap());

                }

                Err(e) if e.kind() == io::ErrorKind::NotFound => {}

                Err(e) => return Err(e.into()),

            }

        }

        Ok(())

    }

    /// Write `refs` into a new pack, then delete `loose_paths` and `old_packs`.

    /// Write the index for a complete, synced temporary pack file and move

    /// both into place; `index` must be sorted by ObjectRef.

    fn finish_pack_files(&self, id: &str, tmp_pack: &Path, index: &[(ObjectRef, PackEntry)]) -> Result<()> {

        let pack_path = self.dir.join(format!("pack-{id}.pack"));

        let idx_path = self.dir.join(format!("pack-{id}.idx"));

        let tmp_idx = self.dir.join(format!("tmp-{id}.idx"));

        {

            let mut f = io::BufWriter::new(fs::File::create(&tmp_idx)?);

            f.write_all(IDX_MAGIC)?;

            f.write_all(&(index.len() as u32).to_le_bytes())?;

            for (r, e) in index {

                f.write_all(&r.0)?;

                f.write_all(&e.offset.to_le_bytes())?;

                f.write_all(&e.stored_len.to_le_bytes())?;

                f.write_all(&e.raw_len.to_le_bytes())?;

                f.write_all(&[type_index(e.ty), e.flags])?;

            }

            f.flush()?;

            f.get_ref().sync_all()?;

        }

        fs::rename(tmp_pack, &pack_path)?;

        fs::rename(&tmp_idx, &idx_path)?;

        self.reload_packs()

    }

    fn write_pack(&self, refs: &[ObjectRef], loose_paths: &[PathBuf], old_packs: &[PathBuf]) -> Result<PackReport> {

        if refs.is_empty() {

            for p in old_packs {

                fs::remove_file(p)?;

                fs::remove_file(p.with_extension("idx"))?;

            }

            self.reload_packs()?;

            return Ok(PackReport::default());

        }

        let mut refs: Vec<ObjectRef> = refs.to_vec();

        refs.sort();

        refs.dedup();

        let mut nonce = [0u8; 8];

        rand::rngs::OsRng.fill_bytes(&mut nonce);

        let id = hex::encode(nonce);

        let tmp_pack = self.dir.join(format!("tmp-{id}.pack"));

        let mut report = PackReport::default();

        let mut index: Vec<(ObjectRef, PackEntry)> = Vec::with_capacity(refs.len());

        {

            let mut f = io::BufWriter::new(fs::File::create(&tmp_pack)?);

            f.write_all(PACK_MAGIC)?;

            let mut offset = PACK_MAGIC.len() as u64;

            for r in &refs {

                let (ty, body) = self.get(r)?;

                let (flags, stored) = if self.is_nocompress(r) {

                    (0, body.clone())

                } else {

                    let compressed = zstd::bulk::compress(&body, ZSTD_LEVEL).map_err(|e| StoreError::Corrupt(*r, format!("zstd: {e}")))?;

                    if compressed.len() < body.len() { (FLAG_ZSTD, compressed) } else { (0, body.clone()) }

                };

                let header_len = 1 + 1 + 4 + 4 + 32;

                f.write_all(&[type_index(ty), flags])?;

                f.write_all(&(stored.len() as u32).to_le_bytes())?;

                f.write_all(&(body.len() as u32).to_le_bytes())?;

                f.write_all(&r.0)?;

                f.write_all(&stored)?;

                index.push((*r, PackEntry { offset: offset + header_len, stored_len: stored.len() as u32, raw_len: body.len() as u32, ty, flags }));

                offset += header_len + stored.len() as u64;

                report.objects += 1;

                report.raw_bytes += body.len() as u64;

                report.packed_bytes += stored.len() as u64;

            }

            f.flush()?;

            f.get_ref().sync_all()?;

        }

        self.finish_pack_files(&id, &tmp_pack, &index)?;