packdir 0.1.0

use anyhow::{Result, anyhow, bail};
use serde::{Deserialize, Serialize};
use std::borrow::Cow;
use std::collections::HashMap;

use crate::packer::{Packer, SIG};

#[macro_export]
macro_rules! unpack {
    ($pack_name:literal) => {
        #[allow(unexpected_cfgs)]
        {
            #[cfg(not(rust_analyzer))]
            let packed_bytes: &'static [u8] =
                include_bytes!(concat!(env!("OUT_DIR"), "/PACKDIR_", $pack_name));

            #[cfg(rust_analyzer)]
            let packed_bytes: &'static [u8] = &[];

            $crate::PackedDir::new(packed_bytes).expect("Failed to load archive")
        }
    };
}

#[macro_export]
macro_rules! unpack_lazy {
    ($pack_name:literal => $vis:vis static $name:ident) => {
        #[allow(unexpected_cfgs)]
        $vis static $name: $crate::LazyDir = $crate::LazyDir::new(|| {
            #[cfg(rust_analyzer)]
            let packed_bytes: &'static [u8] = unreachable!();
            #[cfg(not(rust_analyzer))]
            let packed_bytes: &'static [u8] =
                include_bytes!(concat!(env!("OUT_DIR"), "/PACKDIR_", $pack_name));
            $crate::PackedDir::new(packed_bytes).expect("Failed to initialize Archive")
        });
    };
}

pub(crate) fn read_le<T>(bytes: &[u8], offset: &mut usize) -> Option<T>
where
    T: FromLeBytes,
{
    let size = T::size();
    let end = offset.checked_add(size)?;
    let window = bytes.get(*offset..end)?;
    *offset += size;
    T::from_le_bytes(window)
}

pub(crate) fn try_read_utf8_string(
    name: &str,
    bytes: &[u8],
    offset: &mut usize,
    len: usize,
) -> Result<String> {
    let end = offset
        .checked_add(len)
        .ok_or_else(|| anyhow!("Failed to read {name} at offset {offset}"))?;
    let string_bytes = bytes
        .get(*offset..end)
        .ok_or_else(|| anyhow!("Failed to read {name} at offset {offset}"))?;
    let string = String::from_utf8_lossy(string_bytes).to_string();
    *offset += len;
    Ok(string)
}

pub(crate) fn try_read_le<T>(name: &str, bytes: &[u8], offset: &mut usize) -> Result<T>
where
    T: FromLeBytes,
{
    read_le(bytes, offset).ok_or_else(|| anyhow!("Failed to read {name} at offset {offset}"))
}

pub(crate) trait FromLeBytes: Sized {
    fn from_le_bytes(bytes: &[u8]) -> Option<Self>;
    fn size() -> usize;
}

impl FromLeBytes for u16 {
    fn from_le_bytes(b: &[u8]) -> Option<Self> {
        Some(u16::from_le_bytes(TryInto::<[u8; 2]>::try_into(b).ok()?))
    }
    fn size() -> usize {
        2
    }
}

impl FromLeBytes for u32 {
    fn from_le_bytes(b: &[u8]) -> Option<Self> {
        Some(u32::from_le_bytes(TryInto::<[u8; 4]>::try_into(b).ok()?))
    }
    fn size() -> usize {
        4
    }
}

impl FromLeBytes for u64 {
    fn from_le_bytes(b: &[u8]) -> Option<Self> {
        Some(u64::from_le_bytes(TryInto::<[u8; 8]>::try_into(b).ok()?))
    }
    fn size() -> usize {
        8
    }
}

/// Archive data that can be either a static reference or owned bytes
#[derive(Clone, Debug)]
#[cfg_attr(feature = "enc-dec", derive(bincode::Encode, bincode::Decode))]
pub enum PackedData {
    /// Static reference to included bytes (from include_bytes!)
    Static(&'static [u8]),
    /// Owned bytes (from deserialization or manual creation)
    Owned(Vec<u8>),
}

impl Serialize for PackedData {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        self.as_bytes().serialize(serializer)
    }
}

impl<'de> Deserialize<'de> for PackedData {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        let bytes = Vec::<u8>::deserialize(deserializer)?;
        Ok(Self::Owned(bytes))
    }
}

impl PackedData {
    /// Create from static bytes (typically from include_bytes!)
    pub fn from_static(data: &'static [u8]) -> Self {
        Self::Static(data)
    }

    /// Create from owned bytes
    pub fn from_owned(data: Vec<u8>) -> Self {
        Self::Owned(data)
    }

    /// Get the bytes as a slice
    pub fn as_bytes(&self) -> &[u8] {
        match self {
            Self::Static(data) => data,
            Self::Owned(data) => data,
        }
    }

    /// Convert to owned variant (useful for ensuring we can serialize)
    pub fn into_owned(self) -> Self {
        match self {
            Self::Static(data) => Self::Owned(data.to_vec()),
            Self::Owned(_) => self,
        }
    }

    /// Get the length of the data
    pub fn len(&self) -> usize {
        self.as_bytes().len()
    }

    /// Check if the data is empty
    pub fn is_empty(&self) -> bool {
        self.as_bytes().is_empty()
    }
}

// Convenient From implementations
impl From<&'static [u8]> for PackedData {
    fn from(data: &'static [u8]) -> Self {
        Self::Static(data)
    }
}

impl From<Vec<u8>> for PackedData {
    fn from(data: Vec<u8>) -> Self {
        Self::Owned(data)
    }
}

#[derive(Clone, Serialize, Deserialize, Debug)]
#[cfg_attr(feature = "enc-dec", derive(bincode::Encode, bincode::Decode))]
pub struct PackedFileRecord {
    pub unpacked_size: u64,
    pub packed_size: u64,
    pub data_start: u64,
    pub packer: Packer,
}
impl PackedFileRecord {
    pub(crate) fn try_read(bytes: &[u8], offset: &mut usize) -> Result<Self> {
        Ok(PackedFileRecord {
            unpacked_size: try_read_le::<u64>("unpacked_size", bytes, offset)?,
            packed_size: try_read_le::<u64>("packed_size", bytes, offset)?,
            data_start: try_read_le::<u64>("data_start", bytes, offset)?,
            packer: Packer::try_read_u8(bytes, offset)?,
        })
    }
}

pub type LazyDir = std::sync::LazyLock<PackedDir>;

#[derive(Clone, Serialize, Deserialize, Debug)]
#[cfg_attr(feature = "enc-dec", derive(bincode::Encode, bincode::Decode))]
pub struct PackedDir {
    pub(crate) records: HashMap<String, PackedFileRecord>,
    pub(crate) packed_data: Option<PackedData>,
    #[serde(skip)]
    #[cfg_attr(feature = "enc-dec", bincode(skip))]
    pub(crate) unpack_cache: HashMap<String, Vec<u8>>,
}

impl AsRef<PackedDir> for LazyDir {
    fn as_ref(&self) -> &PackedDir {
        self
    }
}

impl AsRef<PackedDir> for PackedDir {
    fn as_ref(&self) -> &PackedDir {
        self
    }
}

impl PackedDir {
    pub fn new(packed_bytes: &'static [u8]) -> Result<Self> {
        Self::from_data(PackedData::Static(packed_bytes))
    }

    pub fn from_data(archive_data: PackedData) -> Result<Self> {
        let packed_bytes = archive_data.as_bytes();

        if packed_bytes.len() < 9 || &packed_bytes[0..8] != SIG {
            bail!("Invalid archive - header bytes don't match")
        }

        let version = packed_bytes[8];
        match version {
            1 => Self::new_v1(archive_data, 9),
            _ => bail!("Unsupported archive version: {}", version),
        }
    }

    fn new_v1(data: PackedData, offset_start: usize) -> Result<Self> {
        let packed = data.as_bytes();
        let mut offset = offset_start;
        let num_files = try_read_le::<u32>("num_files", packed, &mut offset)? as usize;
        if num_files == 0 {
            return Ok(Self {
                records: Default::default(),
                packed_data: None,
                unpack_cache: Default::default(),
            });
        }

        const MIN_ENTRY_SIZE: usize =
        2  // path_len
        + 8 //unpacked
        + 8 //packed
        + 8 //offset
        + 1 //packer
        ;
        let remaining = packed.len().saturating_sub(offset);

        // Bail out if the header cannot possibly contain that many entries
        if remaining / MIN_ENTRY_SIZE < num_files {
            bail!("Corrupted archive: header too small for declared file count ({num_files})");
        }

        let mut records = HashMap::with_capacity(num_files);

        for _ in 0..num_files {
            let path_len = try_read_le::<u16>("path_len", packed, &mut offset)? as usize;
            let path = try_read_utf8_string("path", packed, &mut offset, path_len)?;
            let record = PackedFileRecord::try_read(packed, &mut offset)?;
            records.insert(path, record);
        }

        Ok(PackedDir {
            records,
            packed_data: Some(data),
            unpack_cache: HashMap::new(),
        })
    }

    pub fn has_entry(&self, path: &str) -> bool {
        self.records.contains_key(path)
    }

    /// Check if all files have been unpacked to cache
    pub fn is_fully_cached(&self) -> bool {
        self.unpack_cache.len() == self.records.len()
    }

    /// Free the packed data if everything is cached
    fn maybe_free_packed_data(&mut self) {
        if self.is_fully_cached() {
            self.packed_data = None;
        }
    }

    /// Common function to unpack a file from the archive data
    fn unpack_file(&self, path: &str) -> Result<Vec<u8>> {
        let entry = self
            .records
            .get(path)
            .ok_or_else(|| anyhow!("File not found: {}", path))?;

        let packed_bytes = self
            .packed_data
            .as_ref()
            .ok_or_else(|| {
                anyhow!(
                    "Packed data has been freed and file '{}' is not cached",
                    path
                )
            })?
            .as_bytes();

        let start = entry.data_start as usize;
        let end = start + entry.packed_size as usize;
        let packed_data = &packed_bytes[start..end];

        entry.packer.unpack(packed_data)
    }

    /// Unpack a file and store it in the cache. Returns a copy-on-write reference.
    /// This method takes &mut self and will populate the cache if the file isn't cached.
    pub fn get(&mut self, path: &str) -> Result<Cow<[u8]>> {
        // Check if already cached first
        if self.unpack_cache.contains_key(path) {
            return Ok(Cow::Borrowed(self.unpack_cache.get(path).expect("cached")));
        }

        // Not in cache, need to unpack
        let unpacked_data = self.unpack_file(path)?;

        self.unpack_cache.insert(path.to_string(), unpacked_data);

        self.maybe_free_packed_data();

        Ok(Cow::Borrowed(self.unpack_cache.get(path).expect("cached")))
    }

    /// Get a cached file if it exists in the cache, otherwise return None.
    /// This method is non-blocking and won't unpack files.
    pub fn get_cached(&self, path: &str) -> Option<Cow<[u8]>> {
        self.unpack_cache
            .get(path)
            .map(|data| Cow::Borrowed(data.as_slice()))
    }

    /// Always get the file. First tries cache, then unpacks on-demand without caching.
    /// This method takes &self and won't modify the cache.
    pub fn get_always(&self, path: &str) -> Result<Cow<[u8]>> {
        // First check if we have it in cache
        if let Some(cached) = self.unpack_cache.get(path) {
            return Ok(Cow::Borrowed(cached));
        }

        // Not in cache, unpack on-demand (without caching)
        let unpacked_data = self.unpack_file(path)?;
        Ok(Cow::Owned(unpacked_data))
    }

    pub fn inflated(&self) -> Result<Self> {
        self.clone().inflate()
    }

    /// Inflate (unpack) all files into the cache and free the packed data.
    /// This is useful when you want to trade memory for speed.
    pub fn inflate(mut self) -> Result<Self> {
        // Get all the paths we need to unpack
        let paths: Vec<String> = self
            .records
            .keys()
            .filter(|path| !self.unpack_cache.contains_key(*path))
            .cloned()
            .collect();

        // Unpack all missing files
        for path in paths {
            let unpacked_data = self.unpack_file(&path)?;
            self.unpack_cache.insert(path, unpacked_data);
        }

        // Free packed data since everything is now cached
        self.maybe_free_packed_data();

        Ok(self)
    }

    pub fn record(&self, path: &str) -> Option<&PackedFileRecord> {
        self.records.get(path)
    }

    pub fn entries(&self) -> Vec<&str> {
        self.records.keys().map(|s| s.as_str()).collect()
    }

    /// Check if the packed data has been freed
    pub fn is_packed_data_freed(&self) -> bool {
        self.packed_data.is_none()
    }

    /// Get cache len
    pub fn cache_len(&self) -> usize {
        self.unpack_cache.len()
    }
    /// Get records len
    pub fn records_len(&self) -> usize {
        self.records.len()
    }

    pub(crate) fn is_valid_packed_name(name: &str) -> bool {
        !name.is_empty()
            && name
                .chars()
                .all(|c| c.is_alphanumeric() || c == '_' || c == '-')
            && !name.starts_with(|c: char| c.is_ascii_digit())
    }
}