thin-fetch 0.1.0

use std::{collections::HashMap, io::BufRead};

use gix_object::Kind;
use gix_pack::data::{entry::Header, input::BytesToEntriesIter};

use crate::error::{Error, WithError as _, bail, error};
use crate::object::{Blob, Commit, Object, Tag, Tree};
use crate::utils::as_hex;
use crate::{HashKind, delta};

#[derive(Clone, Debug)]
struct RawObject {
    kind: Kind,
    data: Vec<u8>,
}

#[derive(Clone)]
struct ObjectStore {
    objects: HashMap<String, RawObject>,
}

impl ObjectStore {
    pub fn new() -> Self {
        Self {
            objects: HashMap::new(),
        }
    }

    pub fn get(&self, hash: &str) -> Option<&RawObject> {
        self.objects.get(hash)
    }

    pub fn insert(&mut self, hash: String, kind: Kind, data: Vec<u8>) {
        let obj = RawObject { kind, data };

        self.objects.entry(hash).or_insert(obj);
    }
}

pub struct PackDecoder<R: BufRead> {
    reader: BytesToEntriesIter<R>,
    objects: ObjectStore,
    hash_kind: HashKind,
}

fn new_object(kind: Kind, data: &[u8], hash_kind: HashKind) -> Result<Object, Error> {
    let obj = match kind {
        Kind::Commit => {
            let commit = Commit::parse(data, hash_kind)?;
            Object::Commit(Box::new(commit))
        }
        Kind::Tree => {
            let tree = Tree::parse(data, hash_kind)?;
            Object::Tree(tree)
        }
        Kind::Tag => {
            let tag = Tag::parse(data, hash_kind)?;
            Object::Tag(tag)
        }
        Kind::Blob => {
            let blob = Blob::new(data.to_vec(), hash_kind);
            Object::Blob(blob)
        }
    };

    Ok(obj)
}

/// Pack file decoder.
///
/// Implements an iterator over [`Object`], resolving deltas as it reads data.
impl<R: BufRead> PackDecoder<R> {
    /// Initialize a decoder of pack files.
    ///
    /// # Errors
    ///
    /// Returns an error when the pack header is invalid.
    pub fn new(reader: R, hash_kind: HashKind) -> Result<Self, Error> {
        let file = BytesToEntriesIter::new_from_header(
            reader,
            gix_pack::data::input::Mode::AsIs,
            gix_pack::data::input::EntryDataMode::KeepAndCrc32,
            hash_kind.into(),
        )
        .map_err(|err| error!("could not read header: {err}"))?;

        let store = ObjectStore::new();

        Ok(Self {
            reader: file,
            objects: store,
            hash_kind,
        })
    }

    fn get_next(&mut self) -> Result<Option<Object>, Error> {
        let Some(hdr) = self.reader.next() else {
            return Ok(None);
        };

        let entry = hdr.map_err(|err| error!("could not read header: {err}"))?;
        let data = entry.compressed.with_err_msg("compressed data not found")?;

        let size = usize::try_from(entry.decompressed_size).with_err(|| {
            error!(
                "compressed data size exceeds max usize: {}",
                entry.decompressed_size
            )
        })?;

        let mut raw_data = vec![0; size];

        let mut inf = zlib_rs::Inflate::new(true, 15);
        inf.decompress(&data, &mut raw_data, zlib_rs::InflateFlush::Finish)
            .map_err(|err| error!("could not decompress: {err:?}"))?;

        // eprintln!("Received object: {:?}", entry.header);

        let obj = match entry.header {
            Header::Commit => new_object(Kind::Commit, &raw_data, self.hash_kind)?,
            Header::Tree => new_object(Kind::Tree, &raw_data, self.hash_kind)?,
            Header::Tag => new_object(Kind::Tag, &raw_data, self.hash_kind)?,
            Header::Blob => new_object(Kind::Blob, &raw_data, self.hash_kind)?,
            Header::RefDelta { base_id } => {
                let hash = as_hex(base_id.as_bytes());

                let apply_to = self
                    .objects
                    .get(&hash)
                    .with_err(|| error!("object not found: {hash}"))?;

                let mut delta_data = raw_data.as_slice();

                let (base_size, read) = delta::decode_header_size(delta_data)
                    .with_err_msg("could not get delta base size")?;
                delta_data = &delta_data[read..];

                let (target_size, read) = delta::decode_header_size(delta_data)
                    .with_err_msg("could not get delta target size")?;
                delta_data = &delta_data[read..];

                if base_size != apply_to.data.len() {
                    bail!(
                        "delta base size mismatch: expected {base_size}, got {}",
                        apply_to.data.len()
                    );
                }

                let mut target = vec![0; target_size];

                delta::apply(&apply_to.data, &mut target, delta_data)
                    .with_err_msg("could not apply delta")?;

                raw_data.clone_from(&target);

                new_object(apply_to.kind, &target, self.hash_kind)?
            }
            Header::OfsDelta { base_distance: _ } => {
                bail!("ofs-delta not supported")
            }
        };

        let hash = as_hex(obj.hash().as_ref());
        self.objects.insert(hash, obj.kind(), raw_data);

        Ok(Some(obj))
    }
}

impl<R: BufRead> Iterator for PackDecoder<R> {
    type Item = Result<Object, Error>;

    fn next(&mut self) -> Option<Self::Item> {
        self.get_next().transpose()
    }
}