#![no_std]
extern crate alloc;

mod types;
use alloc::string::String;
use alloc::vec;
use alloc::vec::Vec;
use core::convert::Infallible;
use core::error::Error;
pub use types::*;

#[derive(Clone, Default)]
pub struct ZipArchive {
    entries: Vec<FileEntry>,
    decompressors: Vec<(u16, DecompressFn)>,
}

#[cfg(feature = "deflate")]
fn deflate_decompress(data: &[u8]) -> Result<Vec<u8>, DecompressError> {
    miniz_oxide::inflate::decompress_to_vec(data)
        .map_err(|_| DecompressError("deflate decompression failed"))
}

pub trait Crc32 {
    fn hash(data: &[u8]) -> u32;
}

#[cfg(feature = "crc32")]
impl Crc32 for () {
    fn hash(data: &[u8]) -> u32 {
        crc32fast::hash(data)
    }
}

impl ZipArchive {
    pub fn new() -> Self {
        let mut this = Self {
            entries: Vec::new(),
            decompressors: Vec::new(),
        };

        #[cfg(feature = "deflate")]
        this.register_decompressor(8, deflate_decompress);

        this
    }

    pub fn entries(&self) -> &[FileEntry] {
        &self.entries
    }

    pub fn find_by_name(&self, name: &str) -> Option<usize> {
        self.entries.iter().position(|entry| entry.name == name)
    }

    #[cfg(feature = "crc32")]
    pub fn add_file(&mut self, name: impl Into<String>, data: impl Into<Vec<u8>>) {
        self.add_file_with_crc::<()>(name, data)
    }

    pub fn add_file_with_crc<C: Crc32>(
        &mut self,
        name: impl Into<String>,
        data: impl Into<Vec<u8>>,
    ) {
        let data = data.into();
        let crc32 = C::hash(&data);
        let size = data.len() as u64;

        self.entries.push(FileEntry {
            compression_method: 0,
            flags: 0,
            dos_mod_time: 0,
            dos_mod_date: 0,
            crc32,
            compressed_size: size,
            uncompressed_size: size,
            local_header_offset: 0,
            name: name.into(),
            data: FileData::Memory(data),
        });
    }
    pub fn register_decompressor(&mut self, method: u16, f: DecompressFn) -> &mut Self {
        self.decompressors.retain(|(m, _)| *m != method);
        self.decompressors.push((method, f));
        self
    }

    fn decompress<E: Error>(&self, method: u16, data: &[u8]) -> Result<Vec<u8>, ZipError<E>> {
        match method {
            0 => Ok(data.to_vec()),
            8 => deflate_decompress(data).map_err(|e| ZipError::FormatError(e.0)),
            _ => self
                .decompressors
                .iter()
                .find(|(m, _)| *m == method)
                .map(|(_, f)| f(data).map_err(|e| ZipError::FormatError(e.0)))
                .unwrap_or(Err(ZipError::UnsupportedCompression(method))),
        }
    }

    pub fn read_file<R: ReadAt>(&self, index: usize, r: &R) -> Result<Vec<u8>, ZipError<R::Error>> {
        self.read_file_with_strategy(index, r, ZipReadStrategy::Merged)
    }

    pub fn read_file_with_strategy<R: ReadAt>(
        &self,
        index: usize,
        r: &R,
        strategy: ZipReadStrategy,
    ) -> Result<Vec<u8>, ZipError<R::Error>> {
        let entry = &self.entries[index];

        if let FileData::Memory(data) = &entry.data {
            return Ok(data.clone());
        }

        let compressed_size = to_usize(entry.compressed_size)?;
        if compressed_size == 0 {
            return Ok(Vec::new());
        }

        match strategy {
            ZipReadStrategy::Strict => {
                let lfh_offset = entry.local_header_offset;
                let mut lfh_base = [0u8; SIZEOF_LFH_BASE];
                r.read_at(lfh_offset, &mut lfh_base)
                    .map_err(ZipError::ReadError)?;

                if &lfh_base[0..4] != SIG_LOCAL_FILE_HEADER {
                    return Err(ZipError::InvalidSignature("LFH signature mismatch"));
                }

                let lfh_name_len = u16::from_le_bytes([lfh_base[26], lfh_base[27]]) as usize;
                let lfh_extra_len = u16::from_le_bytes([lfh_base[28], lfh_base[29]]) as usize;

                let data_offset = lfh_offset
                    + SIZEOF_LFH_BASE as u64
                    + lfh_name_len as u64
                    + lfh_extra_len as u64;

                let mut compressed = vec![0u8; compressed_size];
                r.read_at(data_offset, &mut compressed)
                    .map_err(ZipError::ReadError)?;

                self.decompress(entry.compression_method, &compressed)
            }

            ZipReadStrategy::Optimistic => {
                let FileData::DataOffset { offset, size } = entry.data else {
                    unreachable!()
                };

                let mut compressed = vec![0u8; size];
                r.read_at(offset, &mut compressed)
                    .map_err(ZipError::ReadError)?;

                self.decompress(entry.compression_method, &compressed)
            }

            ZipReadStrategy::Merged => {
                let estimated_lfh_padding = 128;
                let name_len = entry.name.as_bytes().len();
                let chunk_size =
                    SIZEOF_LFH_BASE + name_len + estimated_lfh_padding + compressed_size;

                let file_size = r.size();
                let chunk_size = to_usize(core::cmp::min(
                    chunk_size as u64,
                    file_size - entry.local_header_offset,
                ))?;

                let mut buffer = vec![0u8; chunk_size];
                r.read_at(entry.local_header_offset, &mut buffer)
                    .map_err(ZipError::ReadError)?;

                if buffer.len() < SIZEOF_LFH_BASE
                    || buffer[..SIZEOF_SIGNATURE] != *SIG_LOCAL_FILE_HEADER
                {
                    return Err(ZipError::InvalidSignature("LFH signature mismatch"));
                }

                let lfh_name_len =
                    u16::from_le_bytes([buffer[LFH_OFF_NAME_LEN], buffer[LFH_OFF_NAME_LEN + 1]])
                        as usize;
                let lfh_extra_len =
                    u16::from_le_bytes([buffer[LFH_OFF_EXTRA_LEN], buffer[LFH_OFF_EXTRA_LEN + 1]])
                        as usize;

                let data_start = SIZEOF_LFH_BASE + lfh_name_len + lfh_extra_len;
                let data_end = data_start + compressed_size;

                if data_end <= buffer.len() {
                    self.decompress(entry.compression_method, &buffer[data_start..data_end])
                } else {
                    let exact_data_offset = entry.local_header_offset + data_start as u64;
                    let mut exact_compressed = vec![0u8; compressed_size];
                    r.read_at(exact_data_offset, &mut exact_compressed)
                        .map_err(ZipError::ReadError)?;

                    self.decompress(entry.compression_method, &exact_compressed)
                }
            }
        }
    }

    pub fn with_binary<R: ReadAt>(mut self, r: &R) -> Result<Self, ZipError<R::Error>> {
        let eocd_info = EndOfCentralDirectory::read(r)?;

        let mut cd_buf = Vec::new();
        cd_buf.resize(to_usize(eocd_info.cd_size)?, 0);
        r.read_at(eocd_info.cd_offset, &mut cd_buf)
            .map_err(ZipError::ReadError)?;

        let mut cursor = 0usize;
        while cursor + SIZEOF_CD_BASE <= cd_buf.len() {
            if &cd_buf[cursor..cursor + 4] != SIG_CENTRAL_DIRECTORY {
                break;
            }
            let entry = FileEntry::read_cd(&cd_buf, &mut cursor)?;
            self.entries.push(entry);
        }

        Ok(self)
    }

    fn build_output_size(&self) -> usize {
        let mut lfh_data_total = 0usize;
        let mut offsets = Vec::with_capacity(self.entries.len());

        for entry in &self.entries {
            let current_offset = lfh_data_total as u64;
            offsets.push(current_offset);
            lfh_data_total += entry.lfh_len(current_offset) + entry.data_len();
        }

        let cd_offset = lfh_data_total as u64;
        let mut cd_size = 0usize;
        for (i, entry) in self.entries.iter().enumerate() {
            cd_size += entry.cd_len(offsets[i]);
        }

        let eocd_size =
            EndOfCentralDirectory::len(self.entries.len() as u64, cd_size as u64, cd_offset);

        lfh_data_total + cd_size + eocd_size
    }

    pub fn build(&self) -> Result<Vec<u8>, ZipError<Infallible>> {
        let mut buf = Vec::with_capacity(self.build_output_size());
        self.build_impl::<NoReader>(None, &mut buf)?;
        Ok(buf)
    }

    pub fn build_with_readat<R: ReadAt>(&self, ra: &R) -> Result<Vec<u8>, ZipError<R::Error>> {
        let mut buf = Vec::with_capacity(self.build_output_size());
        self.build_impl(Some(ra), &mut buf)?;
        Ok(buf)
    }

    fn build_impl<R: ReadAt>(
        &self,
        r: Option<&R>,
        buf: &mut Vec<u8>,
    ) -> Result<(), ZipError<R::Error>> {
        let mut cd_entries = Vec::with_capacity(self.entries.len());

        for entry in &self.entries {
            let local_offset = buf.len() as u64;
            cd_entries.push((local_offset, entry));

            entry.write_lfh(local_offset, buf);

            match &entry.data {
                FileData::Memory(v) => {
                    buf.extend_from_slice(v);
                }
                FileData::DataOffset { offset, size } => {
                    if let Some(reader) = r {
                        let start = buf.len();
                        buf.resize(start + *size, 0);
                        reader
                            .read_at(*offset, &mut buf[start..])
                            .map_err(ZipError::ReadError)?;
                    } else {
                        return Err(ZipError::UnsupportedFileData);
                    }
                }
            }
        }

        let cd_offset = buf.len() as u64;
        let mut cd_size = 0u64;

        for (local_offset, entry) in &cd_entries {
            let cd_start = buf.len();
            entry.write_cd(*local_offset, buf);
            cd_size += (buf.len() - cd_start) as u64;
        }

        EndOfCentralDirectory::write(buf, cd_entries.len() as u64, cd_size, cd_offset);
        Ok(())
    }
}

struct NoReader;
impl ReadAt for NoReader {
    type Error = Infallible;
    fn read_at(&self, _offset: u64, _buf: &mut [u8]) -> Result<(), Self::Error> {
        Ok(())
    }
    fn size(&self) -> u64 {
        0
    }
}
