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

#[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 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);
            entries.push(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],
            });
        }
        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,
}

#[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>>,
}

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()) };
        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()?;
        if loose.is_empty() {
            return Ok(PackReport::default());
        }
        let mut nonce = [0u8; 8];
        rand::rngs::OsRng.fill_bytes(&mut nonce);
        let id = hex::encode(nonce);
        let pack_path = self.dir.join(format!("pack-{id}.pack"));
        let idx_path = self.dir.join(format!("pack-{id}.idx"));
        let tmp_pack = self.dir.join(format!("tmp-{id}.pack"));
        let tmp_idx = self.dir.join(format!("tmp-{id}.idx"));
        let mut report = PackReport::default();
        let mut index: Vec<(ObjectRef, PackEntry)> = Vec::with_capacity(loose.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 &loose {
                let (ty, body) = self.get(r)?;
                let compressed = zstd::bulk::compress(&body, ZSTD_LEVEL).map_err(|e| StoreError::Corrupt(*r, format!("zstd: {e}")))?;
                let (flags, stored) = 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()?;
        }
        {
            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()?;
        for (_, path) in &loose {
            fs::remove_file(path)?;
            if let Some(dir) = path.parent() {
                let _ = fs::remove_dir(dir);
            }
        }
        Ok(report)
    }

    pub fn dir(&self) -> &Path {
        &self.dir
    }

    fn path_for(&self, r: &ObjectRef) -> PathBuf {
        let h = r.to_hex();
        self.dir.join("loose").join(&h[..2]).join(&h[2..])
    }

    pub fn has(&self, r: &ObjectRef) -> bool {
        self.path_for(r).is_file() || self.find_packed(r).is_some()
    }

    /// Store `body` as an object of type `ty`; idempotent.
    pub fn put(&self, ty: ObjectType, body: &[u8]) -> Result<ObjectRef> {
        let r = hash_object(ty, body);
        let path = self.path_for(&r);
        if path.is_file() || self.find_packed(&r).is_some() {
            return Ok(r);
        }
        fs::create_dir_all(path.parent().unwrap())?;
        let mut nonce = [0u8; 8];
        rand::rngs::OsRng.fill_bytes(&mut nonce);
        let tmp = self.dir.join("loose").join(format!("tmp-{}", hex::encode(nonce)));
        {
            let mut f = fs::File::create(&tmp)?;
            f.write_all(&header(ty, body.len()))?;
            f.write_all(body)?;
            f.sync_all()?;
        }
        fs::rename(&tmp, &path)?;
        Ok(r)
    }

    /// Read and verify an object, loose or packed.
    pub fn get(&self, r: &ObjectRef) -> Result<(ObjectType, Vec<u8>)> {
        let path = self.path_for(r);
        let data = match fs::read(&path) {
            Ok(d) => d,
            Err(e) if e.kind() == io::ErrorKind::NotFound => {
                return match self.find_packed(r) {
                    Some((i, e)) => self.packs.borrow()[i].read(r, e),
                    None => Err(StoreError::NotFound(*r)),
                };
            }
            Err(e) => return Err(e.into()),
        };
        let nul = data.iter().position(|&b| b == 0).ok_or_else(|| StoreError::Corrupt(*r, "no header".into()))?;
        let head = std::str::from_utf8(&data[..nul]).map_err(|_| StoreError::Corrupt(*r, "header not UTF-8".into()))?;
        let mut parts = head.split(' ');
        let (Some("oxy1"), Some(ty), Some(len), None) = (parts.next(), parts.next(), parts.next(), parts.next()) else {
            return Err(StoreError::Corrupt(*r, format!("bad header {head:?}")));
        };
        let ty = ObjectType::from_name(ty).ok_or_else(|| StoreError::Corrupt(*r, format!("unknown type {ty}")))?;
        let len: usize = len.parse().map_err(|_| StoreError::Corrupt(*r, "bad length".into()))?;
        let body = &data[nul + 1..];
        if body.len() != len {
            return Err(StoreError::Corrupt(*r, format!("length {} != header {len}", body.len())));
        }
        if hash_object(ty, body) != *r {
            return Err(StoreError::Corrupt(*r, "hash mismatch".into()));
        }
        Ok((ty, body.to_vec()))
    }

    pub fn get_typed(&self, r: &ObjectRef, want: ObjectType) -> Result<Vec<u8>> {
        let (ty, body) = self.get(r)?;
        if ty != want {
            return Err(StoreError::WrongType(*r, ty.name(), want.name()));
        }
        Ok(body)
    }

    pub fn put_change(&self, c: &Change) -> Result<ObjectRef> {
        self.put(ObjectType::Change, &c.encode())
    }

    pub fn get_change(&self, r: &ObjectRef) -> Result<Change> {
        Ok(Change::decode(&self.get_typed(r, ObjectType::Change)?)?)
    }

    pub fn put_chunk(&self, bytes: &[u8]) -> Result<ObjectRef> {
        self.put(ObjectType::Chunk, bytes)
    }

    pub fn put_state(&self, st: &StateObject) -> Result<ObjectRef> {
        self.put(ObjectType::State, &st.encode())
    }

    pub fn get_state(&self, r: &ObjectRef) -> Result<StateObject> {
        Ok(StateObject::decode(&self.get_typed(r, ObjectType::State)?)?)
    }

    pub fn put_attest(&self, a: &Attest) -> Result<ObjectRef> {
        self.put(ObjectType::Attest, &a.encode())
    }

    pub fn get_attest(&self, r: &ObjectRef) -> Result<Attest> {
        Ok(Attest::decode(&self.get_typed(r, ObjectType::Attest)?)?)
    }

    fn seal_index_path(&self, change: &ObjectRef) -> PathBuf {
        let h = change.to_hex();
        self.dir.join("seals").join(&h[..2]).join(&h[2..])
    }

    /// Record that seal attestation `attest` covers `change`.
    pub fn index_seal(&self, change: &ObjectRef, attest: &ObjectRef) -> Result<()> {
        let mut seals = self.seals_for(change)?;
        if seals.contains(attest) {
            return Ok(());
        }
        seals.push(*attest);
        seals.sort();
        let path = self.seal_index_path(change);
        fs::create_dir_all(path.parent().unwrap())?;
        let text: String = seals.iter().map(|s| format!("{}\n", s.to_hex())).collect();
        let tmp = path.with_extension("tmp");
        fs::write(&tmp, text)?;
        fs::rename(&tmp, &path)?;
        Ok(())
    }

    /// Seal attestations known to cover `change`.
    pub fn seals_for(&self, change: &ObjectRef) -> Result<Vec<ObjectRef>> {
        match fs::read_to_string(self.seal_index_path(change)) {
            Ok(t) => Ok(t.lines().filter_map(ObjectRef::from_hex).collect()),
            Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(vec![]),
            Err(e) => Err(e.into()),
        }
    }

    pub fn put_identity(&self, id: &Identity) -> Result<ObjectRef> {
        self.put(ObjectType::Identity, &id.encode())
    }

    pub fn get_identity(&self, r: &ObjectRef) -> Result<Identity> {
        Ok(Identity::decode(&self.get_typed(r, ObjectType::Identity)?)?)
    }

    fn sig_path(&self, r: &ObjectRef) -> PathBuf {
        let h = r.to_hex();
        self.dir.join("sigs").join(&h[..2]).join(&h[2..])
    }

    /// Signatures stored alongside object `r` (spec §5.4). Empty if none.
    pub fn signatures(&self, r: &ObjectRef) -> Result<Vec<Signature>> {
        match fs::read(self.sig_path(r)) {
            Ok(b) => Ok(Signature::decode_list(&b)?),
            Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(vec![]),
            Err(e) => Err(e.into()),
        }
    }

    /// Attach a signature to object `r`, replacing any earlier one by the same key.
    pub fn put_signature(&self, r: &ObjectRef, sig: &Signature) -> Result<()> {
        let mut sigs = self.signatures(r)?;
        sigs.retain(|s| s.key != sig.key);
        sigs.push(sig.clone());
        sigs.sort_by_key(|s| s.key);
        let path = self.sig_path(r);
        fs::create_dir_all(path.parent().unwrap())?;
        let mut nonce = [0u8; 8];
        rand::rngs::OsRng.fill_bytes(&mut nonce);
        let tmp = self.dir.join(format!("sigs-tmp-{}", hex::encode(nonce)));
        {
            let mut f = fs::File::create(&tmp)?;
            f.write_all(&Signature::encode_list(&sigs))?;
            f.sync_all()?;
        }
        fs::rename(&tmp, &path)?;
        Ok(())
    }

    /// Every object reference in the store, loose and packed (slow; for gc and tests).
    pub fn list(&self) -> io::Result<Vec<ObjectRef>> {
        let mut out: Vec<ObjectRef> = self.loose_objects()?.into_iter().map(|(r, _)| r).collect();
        for p in self.packs.borrow().iter() {
            out.extend(p.refs.iter().copied());
        }
        out.sort();
        out.dedup();
        Ok(out)
    }
}

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

    #[test]
    fn put_get_verify() {
        let dir = tempfile::tempdir().unwrap();
        let s = Store::open(dir.path().join("objects")).unwrap();
        let r = s.put(ObjectType::Chunk, b"hello").unwrap();
        assert!(s.has(&r));
        assert_eq!(s.get(&r).unwrap(), (ObjectType::Chunk, b"hello".to_vec()));
        assert_eq!(s.put(ObjectType::Chunk, b"hello").unwrap(), r);
        assert_eq!(s.list().unwrap(), vec![r]);
        assert!(matches!(s.get_typed(&r, ObjectType::Blob), Err(StoreError::WrongType(..))));
        assert!(matches!(s.get(&ObjectRef([9; 32])), Err(StoreError::NotFound(_))));

        // Tamper: flip a byte, read must fail.
        let path = s.path_for(&r);
        let mut data = fs::read(&path).unwrap();
        let last = data.len() - 1;
        data[last] ^= 1;
        fs::write(&path, data).unwrap();
        assert!(matches!(s.get(&r), Err(StoreError::Corrupt(..))));
    }

    #[test]
    fn packing_keeps_every_object_readable() {
        let dir = tempfile::tempdir().unwrap();
        let s = Store::open(dir.path()).unwrap();
        let mut refs = Vec::new();
        for i in 0..300u32 {
            let body = format!("object {i} ").repeat(1 + (i % 7) as usize);
            refs.push((s.put(ObjectType::Chunk, body.as_bytes()).unwrap(), body));
        }
        let c = crate::objects::tests::sample_change();
        let cr = s.put_change(&c).unwrap();
        assert_eq!(s.loose_count().unwrap(), 301);
        let report = s.pack_loose().unwrap();
        assert_eq!(report.objects, 301);
        assert!(report.packed_bytes < report.raw_bytes, "{report:?}");
        assert_eq!(s.loose_count().unwrap(), 0);
        assert_eq!(s.pack_count(), 1);
        for (r, body) in &refs {
            assert!(s.has(r));
            assert_eq!(s.get(r).unwrap(), (ObjectType::Chunk, body.as_bytes().to_vec()));
        }
        assert_eq!(s.get_change(&cr).unwrap(), c);
        assert_eq!(s.list().unwrap().len(), 301);
        // Re-putting a packed object is a no-op; new objects go loose and a
        // second pack coexists with the first.
        assert_eq!(s.put(ObjectType::Chunk, refs[0].1.as_bytes()).unwrap(), refs[0].0);
        assert_eq!(s.loose_count().unwrap(), 0);
        let extra = s.put(ObjectType::Chunk, b"after packing").unwrap();
        s.pack_loose().unwrap();
        assert_eq!(s.pack_count(), 2);
        assert_eq!(s.get(&extra).unwrap().1, b"after packing");
        // A reopened store sees the packs.
        let s2 = Store::open(dir.path()).unwrap();
        assert_eq!(s2.pack_count(), 2);
        assert_eq!(s2.get(&refs[5].0).unwrap().1, refs[5].1.as_bytes());
        assert!(matches!(s2.get(&ObjectRef([9; 32])), Err(StoreError::NotFound(_))));
        // Corrupting an object's body inside the pack is detected on read.
        let target = refs[5].0;
        let (pack_path, offset) = {
            let packs = s2.packs.borrow();
            let (i, e) = s2.find_packed(&target).unwrap();
            (packs[i].path.clone(), e.offset as usize)
        };
        let mut data = fs::read(&pack_path).unwrap();
        data[offset] ^= 0xff;
        fs::write(&pack_path, data).unwrap();
        let s3 = Store::open(dir.path()).unwrap();
        assert!(matches!(s3.get(&target), Err(StoreError::Corrupt(..))));
        assert_eq!(s3.get(&refs[6].0).unwrap().1, refs[6].1.as_bytes(), "other objects are unaffected");
    }

    #[test]
    fn signatures_are_stored_beside_objects() {
        use crate::identity::{Domain, SecretKey};
        let dir = tempfile::tempdir().unwrap();
        let s = Store::open(dir.path()).unwrap();
        let r = s.put(ObjectType::Chunk, b"x").unwrap();
        assert!(s.signatures(&r).unwrap().is_empty());
        let sk = SecretKey { ed_seed: [1; 32], ml_seed: [2; 32] };
        let id = sk.identity("a", "b", 0);
        let key = s.put_identity(&id).unwrap();
        assert_eq!(s.get_identity(&key).unwrap(), id);
        let sig = sk.sign(Domain::Change, &key, &r);
        s.put_signature(&r, &sig).unwrap();
        s.put_signature(&r, &sig).unwrap();
        assert_eq!(s.signatures(&r).unwrap(), vec![sig]);
        // The object itself is untouched by its signatures.
        assert_eq!(s.get(&r).unwrap().1, b"x");
    }

    #[test]
    fn change_round_trip_through_store() {
        let dir = tempfile::tempdir().unwrap();
        let s = Store::open(dir.path()).unwrap();
        let c = crate::objects::tests::sample_change();
        let r = s.put_change(&c).unwrap();
        assert_eq!(r, c.hash());
        assert_eq!(s.get_change(&r).unwrap(), c);
    }
}