atlas-archive-core 1.1.0

#![cfg_attr(not(feature = "std"), no_std)]
#![cfg_attr(feature = "simd", feature(portable_simd))]
extern crate alloc;

pub mod ans;
pub mod huffman;
pub mod lzma;
pub mod mixer;
pub mod poor_compress;
pub mod predictor;
pub mod transform;

pub mod loom;

#[cfg(feature = "std")]
extern crate std;
#[cfg(feature = "std")]
use rayon::prelude::*;
#[cfg(feature = "std")]
use std::sync::Once;

#[cfg(feature = "std")]
static THREAD_POOL_INIT: Once = Once::new();

/// Global shared runtime for background PagedLoom writes.
#[cfg(all(feature = "std", feature = "async"))]
pub(crate) fn get_background_runtime() -> Option<std::sync::Arc<tokio::runtime::Runtime>> {
    use std::sync::OnceLock;
    static RUNTIME: OnceLock<Option<std::sync::Arc<tokio::runtime::Runtime>>> = OnceLock::new();
    RUNTIME
        .get_or_init(|| {
            tokio::runtime::Builder::new_multi_thread()
                .worker_threads(4) // Increased for faster async I/O
                .enable_all()
                .build()
                .ok()
                .map(std::sync::Arc::new)
        })
        .clone()
}

/// Initializes the Rayon thread pool to use at most 80% of available CPU cores.
#[cfg(feature = "std")]
pub fn init_thread_pool() {
    THREAD_POOL_INIT.call_once(|| {
        let cpus = num_cpus::get();
        let threads = (cpus * 8 / 10).max(1);
        let _ = rayon::ThreadPoolBuilder::new()
            .num_threads(threads)
            .build_global();
    });
}

use crate::ans::{RansDecoder, RansEncoder};
use crate::loom::{LoomEnum, LoomPredictor, LoomWeaver};
use crate::transform::{Transform, TransformChain};
use alloc::vec::Vec;

const CHUNK_SIZE: usize = 1024 * 1024; // 1MB chunks for better global context and ratio

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum LoomMode {
    Off,
    Standard,
    Paged,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum LoomAggression {
    Low,
    Med,
    High,
    Ultra,
}

impl LoomAggression {
    pub fn throttle_rate(&self) -> usize {
        match self {
            LoomAggression::Low => 8,
            LoomAggression::Med => 2,
            LoomAggression::High => 1,
            LoomAggression::Ultra => 1,
        }
    }
}

/// Pruning aggressiveness - controls how much of the cache is pruned per cycle.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum PruneAggression {
    /// Prune 2% - gradual, keeps more high-relevance nodes (for Safe Ultra)
    Gradual,
    /// Prune 5% - balanced (default)
    Balanced,
    /// Prune 10% - aggressive (for memory-constrained environments)
    Aggressive,
}

impl PruneAggression {
    /// Returns the percentage of nodes to prune (as denominator: 50 = 2%, 20 = 5%, 10 = 10%)
    pub fn prune_divisor(&self) -> usize {
        match self {
            PruneAggression::Gradual => 50,    // 2%
            PruneAggression::Balanced => 20,   // 5%
            PruneAggression::Aggressive => 10, // 10%
        }
    }
}

/// Configuration for the Predictive Loom Compression (PLC) engine.
#[derive(Clone, Debug)]
pub struct PlcConfig {
    /// Dimension of the internal predictor embeddings.
    pub model_dim: usize,
    /// Number of bytes to use as history for prediction.
    pub window_size: usize,
    /// Which Loom implementation to use.
    pub loom_mode: LoomMode,
    /// How aggressively to train the Loom.
    pub aggression: LoomAggression,
    /// Enable high-order context mixing (PAQ-style).
    pub use_mixer: bool,
    /// Context orders for the mixer (e.g., [1, 2, 4, 8]).
    pub mixer_orders: Vec<usize>,
    /// Enable LZMA-style sliding window compression.
    pub use_lzma: bool,
    /// Slider window size (e.g., 64KB).
    pub lzma_window_size: usize,
    /// Preprocessing transforms to apply.
    pub transforms: Vec<Transform>,
    /// Max RAM usage for Loom Cache (if paging enabled).
    pub paged_loom_max_ram: usize,
    /// Hint for target page size in bytes (50MB-500MB recommended).
    pub page_size_hint: usize,
    /// Maximum nodes per page before splitting (default 4096).
    pub page_max_nodes: usize,
    /// Minimum size in bytes to attempt Loom compression (default 1024).
    pub min_loom_size: usize,
    /// Minimum size in bytes to attempt any compression (default 32).
    pub min_compress_size: usize,
    /// Maximum nodes per mixer model (Ultra mode: ~1M, Default: 1024).
    pub max_nodes: usize,
    /// Pre-trained dictionary for priming the engine.
    pub dictionary: Option<alloc::vec::Vec<u8>>,
    /// Enable verbose logging for Loom operations.
    pub verbose: bool,
    /// Pruning aggressiveness (how much of cache is pruned per cycle).
    pub prune_aggression: PruneAggression,
    /// Path to a persistent paged dictionary (folder). If set, pages are loaded/saved here.
    pub persistent_dict_path: Option<alloc::string::String>,
}

impl Default for PlcConfig {
    fn default() -> Self {
        Self {
            model_dim: 32,
            window_size: 16,
            loom_mode: LoomMode::Standard,
            aggression: LoomAggression::Med,
            use_mixer: true,
            mixer_orders: alloc::vec![1, 2, 4, 8],
            use_lzma: true,
            lzma_window_size: 64 * 1024,
            transforms: alloc::vec![Transform::Bwt, Transform::Mtf],
            paged_loom_max_ram: 1024 * 1024 * 1024, // 1GB default
            page_size_hint: 128 * 1024 * 1024,      // 128MB default page
            page_max_nodes: 4096,                   // 4096 nodes per page
            min_loom_size: 0,                       // 0B min for Loom (always attempt)
            min_compress_size: 32,                  // 32B min for compression
            max_nodes: 1024,                        // 1024 nodes default
            dictionary: None,
            verbose: false,
            prune_aggression: PruneAggression::Balanced,
            persistent_dict_path: None,
        }
    }
}

impl PlcConfig {
    /// Recommended settings for maximum compression ratio.
    /// Optimized for ratio <0.0066 with adaptive caching and deeper prediction.
    pub fn ultra() -> Self {
        Self {
            model_dim: 96,
            window_size: 64,
            loom_mode: LoomMode::Paged,
            aggression: LoomAggression::Ultra,
            use_mixer: true,
            mixer_orders: alloc::vec![1, 2, 4, 8, 12, 16, 24, 32, 40, 48, 64],
            use_lzma: true,
            lzma_window_size: 1024 * 1024,
            transforms: alloc::vec![Transform::Auto],
            paged_loom_max_ram: 512 * 1024 * 1024, // 512MB default (Safe for multi-thread)
            page_size_hint: 32 * 1024 * 1024,
            page_max_nodes: 16384, // 16k nodes (~16MB per page)
            min_loom_size: 0,
            min_compress_size: 32,
            max_nodes: 500_000,
            dictionary: None,
            verbose: false,
            prune_aggression: PruneAggression::Balanced,
            persistent_dict_path: None,
        }
    }

    pub fn safe_ultra() -> Self {
        Self {
            model_dim: 128,
            window_size: 64,
            loom_mode: LoomMode::Paged,
            aggression: LoomAggression::Ultra,
            use_mixer: true,
            // Extended orders: include 64 for deep pattern matching
            mixer_orders: alloc::vec![1, 2, 4, 8, 12, 16, 24, 32, 40, 48, 56, 64],
            use_lzma: true,
            lzma_window_size: 1024 * 1024, // 1MB LZMA window
            transforms: alloc::vec![Transform::Auto],
            paged_loom_max_ram: 1024 * 1024 * 1024, // 1GB strict cap
            page_size_hint: 32 * 1024 * 1024,       // 32MB pages (larger for better context)
            page_max_nodes: 32768,                  // 32k nodes per page
            min_loom_size: 0,
            min_compress_size: 32,
            max_nodes: 1_000_000, // 1M nodes (~1GB budget)
            dictionary: None,
            verbose: false,
            prune_aggression: PruneAggression::Gradual,
            persistent_dict_path: None,
        }
    }
}

/// Core trait for Atlas compressors.
pub trait Compressor {
    type Error;
    fn compress(&mut self, data: &[u8]) -> Result<Vec<u8>, Self::Error>;
}

/// Core trait for Atlas decompressors.
pub trait Decompressor {
    type Error;
    fn decompress(&mut self, data: &[u8], original_len: usize) -> Result<Vec<u8>, Self::Error>;
}

/// Predictive Loom Compressor (PLC) implementation (Enhanced with High-Order Context).
pub struct LoomCompressor {
    loom: Option<LoomEnum>,
    config: PlcConfig,
}

impl LoomCompressor {
    pub fn new(config: PlcConfig) -> Self {
        // Multi-thread RAM guard: divide budget by threads to avoid OOM
        // For Paged mode, use gentler scaling since pages are disk-backed
        // Multi-thread RAM guard: divide budget by threads to avoid OOM
        // For Paged mode, use gentler scaling since pages are disk-backed
        #[cfg(feature = "std")]
        {
            // DISABLED: Automatic scaling breaks determinism between Compressor and Decompressor!
            // Decompressor doesn't know we scaled down, so it uses full RAM/Nodes.
            // This causes divergence in Pruning (Compressor prunes earlier).
            // User must configure RAM appropriately for their thread count.
            /*
            let threads = (num_cpus::get() * 8 / 10).max(1);
            if config.paged_loom_max_ram > 64 * 1024 * 1024 {
                // Paged mode: pages spill to disk, so less aggressive scaling
                if config.loom_mode == LoomMode::Paged {
                    // Scale by sqrt(threads) instead of threads
                    let scale = (threads as f64).sqrt().ceil() as usize;
                    config.paged_loom_max_ram /= scale.max(1);
                } else {
                    config.paged_loom_max_ram /= threads;
                }
            }
            */
        }

        let estimated_ram_bytes = config.max_nodes * 1024; // ~1KB per node overhead
        if config.aggression == LoomAggression::Ultra
            && estimated_ram_bytes > config.paged_loom_max_ram
        {
            #[cfg(feature = "std")]
            if config.verbose {
                std::println!(
                    "[Atlas-Archive] Ultra mode RAM budget per thread exceeded, capping nodes."
                );
            }
            // config.max_nodes = config.paged_loom_max_ram / 1024; // Also disabled to prevent divergence
        }

        let loom = if config.loom_mode == LoomMode::Off {
            None
        } else {
            // Step 7: Fallback: Loom off if error
            match LoomEnum::new(config.clone()) {
                Ok(mut l) => {
                    // Warm cache on page load: Prime Loom with dictionary if provided
                    if let Some(dict) = &config.dictionary {
                        let mut history = Vec::new();
                        for &sym in dict {
                            l.weave(&history, sym);
                            history.push(sym);
                            if history.len() > config.window_size {
                                history.remove(0);
                            }
                        }
                    }
                    Some(l)
                }
                Err(e) => {
                    #[cfg(feature = "std")]
                    if config.verbose {
                        std::println!(
                            "[Atlas-Archive] Loom initialization failed, falling back to Off: {}",
                            e
                        );
                    }
                    None
                }
            }
        };
        Self { loom, config }
    }
}

#[derive(Debug)]
pub enum LoomError {
    PackingError,
    AnsError,
    LoomInternalError,
    InvalidArchive,
}

impl core::fmt::Display for LoomError {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            LoomError::PackingError => write!(f, "Packing error: internal buffer mismatch"),
            LoomError::AnsError => write!(f, "ANS entropy coding error"),
            LoomError::LoomInternalError => write!(f, "Loom adaptive model error"),
            LoomError::InvalidArchive => write!(f, "Invalid NYX archive magic or format"),
        }
    }
}

#[cfg(feature = "std")]
impl std::error::Error for LoomError {}

/// An entry in an Atlas archive.
#[derive(Clone, Debug, Default)]
pub struct ArchiveEntry {
    pub name: Vec<u8>,
    pub data: Vec<u8>,
}

/// A multi-file Atlas archive (.nyx format).
#[derive(Clone, Debug, Default)]
pub struct AtlasArchive {
    pub entries: Vec<ArchiveEntry>,
}

impl AtlasArchive {
    pub const MAGIC: &'static [u8; 4] = b"NYX\x01";

    /// Pack the archive into a binary stream.
    pub fn pack(&self) -> Vec<u8> {
        let mut buf = Vec::new();
        buf.extend_from_slice(Self::MAGIC);
        buf.extend_from_slice(&(self.entries.len() as u32).to_le_bytes());

        for entry in &self.entries {
            buf.extend_from_slice(&(entry.name.len() as u32).to_le_bytes());
            buf.extend_from_slice(&entry.name);
            buf.extend_from_slice(&(entry.data.len() as u64).to_le_bytes());
            buf.extend_from_slice(&entry.data);
        }
        buf
    }

    /// Unpack an archive from a binary stream.
    pub fn unpack(data: &[u8]) -> Result<Self, LoomError> {
        if data.len() < 8 || &data[0..4] != Self::MAGIC {
            return Err(LoomError::InvalidArchive);
        }

        let mut entries = Vec::new();
        let entry_count = u32::from_le_bytes(
            data[4..8]
                .try_into()
                .map_err(|_| LoomError::InvalidArchive)?,
        ) as usize;
        let mut offset = 8;

        for _ in 0..entry_count {
            if offset + 4 > data.len() {
                return Err(LoomError::InvalidArchive);
            }
            let name_len = u32::from_le_bytes(
                data[offset..offset + 4]
                    .try_into()
                    .map_err(|_| LoomError::InvalidArchive)?,
            ) as usize;
            offset += 4;

            if offset + name_len > data.len() {
                return Err(LoomError::InvalidArchive);
            }
            let name = data[offset..offset + name_len].to_vec();
            offset += name_len;

            if offset + 8 > data.len() {
                return Err(LoomError::InvalidArchive);
            }
            let data_len = u64::from_le_bytes(
                data[offset..offset + 8]
                    .try_into()
                    .map_err(|_| LoomError::InvalidArchive)?,
            ) as usize;
            offset += 8;

            if offset + data_len > data.len() {
                return Err(LoomError::InvalidArchive);
            }
            let entry_data = data[offset..offset + data_len].to_vec();
            offset += data_len;

            entries.push(ArchiveEntry {
                name,
                data: entry_data,
            });
        }

        Ok(Self { entries })
    }
}

impl Compressor for LoomCompressor {
    type Error = LoomError;

    fn compress(&mut self, data: &[u8]) -> Result<Vec<u8>, Self::Error> {
        if data.is_empty() {
            return Ok(Vec::new());
        }

        // --- Multi-File Archive Handling ---
        // For simplicity, LoomCompressor here handles single continuous byte stream.
        // Parallel chunking happens inside.

        self.compress_parallel(data)
    }
}

impl LoomCompressor {
    fn compress_parallel(&mut self, data: &[u8]) -> Result<Vec<u8>, LoomError> {
        let chunk_size = CHUNK_SIZE;

        #[cfg(feature = "std")]
        {
            // Use Rayon for parallel chunk compression with balanced workload
            let chunks: Vec<_> = data.chunks(chunk_size).collect();

            // Throttle: ensure each thread gets at least 2 chunks to amortize overhead
            let min_chunks_per_thread = 2;

            let results: Vec<Result<Vec<u8>, LoomError>> = if chunks.len() == 1 {
                // Optimization: Avoid Rayon overhead for single chunk + Better Debugging
                chunks
                    .into_iter()
                    .map(|chunk| {
                        let mut inst = LoomCompressor::new(self.config.clone());
                        inst.compress_chunk(chunk)
                    })
                    .collect()
            } else {
                chunks
                    .into_par_iter()
                    .with_min_len(min_chunks_per_thread)
                    .map(|chunk| {
                        let mut inst = LoomCompressor::new(self.config.clone());
                        inst.compress_chunk(chunk)
                    })
                    .collect()
            };

            let mut packed = Vec::new();
            let mut total_compressed = 0;
            for res in results {
                let chunk_data = res?;
                total_compressed += chunk_data.len();
                packed.extend_from_slice(&chunk_data);
            }

            if self.config.verbose {
                let ratio = total_compressed as f64 / data.len() as f64;
                std::println!(
                    "[Atlas] Parallel compression complete. Ratio: {:.4}, Peak RAM ~{}MB",
                    ratio,
                    (total_compressed * 1) / (1024 * 1024) // Placeholder for better estimate if needed
                );
            }
            Ok(packed)
        }

        #[cfg(not(feature = "std"))]
        {
            // Sequential fallback for no_std
            let mut packed = Vec::new();
            for chunk in data.chunks(chunk_size) {
                let chunk_data = self.compress_chunk(chunk)?;
                packed.extend_from_slice(&chunk_data);
            }
            Ok(packed)
        }
    }

    fn compress_chunk(&mut self, data: &[u8]) -> Result<Vec<u8>, LoomError> {
        #[cfg(feature = "std")]
        let chunk_start = std::time::Instant::now();

        // Step 1: Handle Minimum Size for Compression
        if data.len() < self.config.min_compress_size {
            return self.pack_raw_chunk(data);
        }

        // Step 2: Apply Preprocessing Transforms
        let file_type = crate::poor_compress::DetectedType::detect(data);

        let mut transform_list = self.config.transforms.clone();

        // Auto-Heuristic: If Auto is set, adjust based on detected type
        if transform_list.contains(&Transform::Auto) {
            match file_type {
                crate::poor_compress::DetectedType::Jpeg => {
                    // JPEGs benefit from YUV before Auto chain
                    transform_list.insert(0, Transform::Yuv);
                }
                crate::poor_compress::DetectedType::Png
                | crate::poor_compress::DetectedType::Gif => {
                    // Already compressed. Avoid heavy BWT/MTF, use Delta+Rle
                    transform_list = vec![Transform::Delta, Transform::Rle];
                }
                crate::poor_compress::DetectedType::Mp4
                | crate::poor_compress::DetectedType::Mp3 => {
                    // Media benefit from ChannelSeparation, but skip heavy BWT
                    transform_list = vec![
                        Transform::ChannelSeparation(2),
                        Transform::Delta,
                        Transform::Rle,
                    ];
                }
                crate::poor_compress::DetectedType::Encrypted => {
                    // Don't try complex transforms on already encrypted/random data
                    transform_list = vec![];
                }
                crate::poor_compress::DetectedType::Zip => {
                    // Archives are already compressed
                    transform_list = vec![Transform::Rle];
                }
                _ => {}
            }
        }

        let chain = TransformChain::new(transform_list);
        let (transformed, bwt_indices) = chain.apply(data.to_vec());

        if self.config.verbose {
            let est_ram = transformed.len() + (bwt_indices.len() * 8);
            std::println!(
                "[Atlas] Detected Type: {:?}, Final Ratio: {:.4}, Est. Transform RAM: {:.2}MB",
                file_type,
                transformed.len() as f64 / data.len() as f64,
                est_ram as f64 / (1024.0 * 1024.0)
            );
        }

        // Step 3: Handle Minimum Size for Loom
        if transformed.len() < self.config.min_loom_size || self.loom.is_none() {
            #[cfg(feature = "std")]
            if self.config.verbose {
                std::println!(
                    "[Atlas] Chunk {}B compressed (no Loom) in {:?}",
                    data.len(),
                    chunk_start.elapsed()
                );
            }
            return self.pack_ans_only_chunk(&transformed, bwt_indices, data.len(), file_type);
        }

        let loom = self.loom.as_mut().unwrap();

        // Step 4: Two-Pass Loom Compression
        // Pass 1: Forward weave to sync state and collect probability models
        let mut models = Vec::with_capacity(transformed.len());
        let mut history = Vec::with_capacity(transformed.len());

        for &sym in transformed.iter() {
            // Predict NEXT symbol probabilities
            let model = loom.predict(&history);
            models.push(model);

            // Weave current symbol into Loom state
            loom.weave(&history, sym);
            history.push(sym);
        }

        // Pass 2: Reverse ANS encoding
        let mut encoder = RansEncoder::new();
        let mut compressed_u16 = Vec::new();
        for (sym, model) in transformed.iter().zip(models.into_iter()).rev() {
            encoder.encode(&model, *sym, &mut compressed_u16);
        }
        encoder.finish(&mut compressed_u16);

        // Convert u16 buffer to bytes
        let mut compressed_ans = Vec::with_capacity(compressed_u16.len() * 2);
        for &val in &compressed_u16 {
            compressed_ans.extend_from_slice(&val.to_le_bytes());
        }

        // Step 5: Pack Chunk
        let packed = self.pack_loom_chunk(
            &compressed_ans,
            bwt_indices,
            data.len(),
            transformed.len(),
            file_type,
        );

        // --- No Expansion Guarantee ---
        if packed.len() >= data.len() {
            #[cfg(feature = "std")]
            if self.config.verbose {
                std::println!(
                    "[Loom] Chunk expanded ({} -> {}), falling back to Raw",
                    data.len(),
                    packed.len()
                );
            }
            return self.pack_raw_chunk(data);
        }

        #[cfg(feature = "std")]
        if self.config.verbose {
            std::println!(
                "[Loom] Chunk compressed: {} -> {} bytes (Ratio: {:.4}) in {:?}",
                data.len(),
                packed.len(),
                packed.len() as f64 / data.len() as f64,
                chunk_start.elapsed()
            );
        }

        Ok(packed)
    }

    fn pack_raw_chunk(&self, data: &[u8]) -> Result<Vec<u8>, LoomError> {
        let mut packed = Vec::with_capacity(data.len() + 8);
        // Header: [raw_len: 24 bit | type: 7 bit | flag: 1 (MSB) ] [compressed_len: 32 bit]
        let file_type = crate::poor_compress::DetectedType::detect(data);
        let type_val = (file_type as u32) & 0x7F;
        let len = data.len() as u32;
        let header_flag = (len & 0x00FF_FFFF) | (type_val << 24) | 0x8000_0000;

        packed.extend_from_slice(&header_flag.to_le_bytes());
        packed.extend_from_slice(&len.to_le_bytes()); // compressed_len is original for raw
        packed.extend_from_slice(data);
        Ok(packed)
    }

    fn pack_ans_only_chunk(
        &self,
        transformed: &[u8],
        bwt_indices: Vec<usize>,
        original_len: usize,
        file_type: crate::poor_compress::DetectedType,
    ) -> Result<Vec<u8>, LoomError> {
        // If no transforms at all, safe to use raw
        if bwt_indices.is_empty() && self.config.transforms.is_empty() {
            return self.pack_raw_chunk(transformed);
        }

        // Pack as a Loom chunk with compressed_len == transformed_len to signal "No ANS"
        Ok(self.pack_loom_chunk(
            transformed,
            bwt_indices,
            original_len,
            transformed.len(),
            file_type,
        ))
    }

    fn pack_loom_chunk(
        &self,
        data: &[u8],
        bwt_indices: Vec<usize>,
        original_len: usize,
        transformed_len: usize,
        file_type: crate::poor_compress::DetectedType,
    ) -> Vec<u8> {
        let mut packed = Vec::new();
        // Header: [orig_len: 24 bit | type: 7 bit | flag: 0 (MSB)] [compressed_len: u32]...
        let type_val = (file_type as u32) & 0x7F;
        let meta = (original_len as u32 & 0x00FF_FFFF) | (type_val << 24);

        packed.extend_from_slice(&meta.to_le_bytes());
        packed.extend_from_slice(&(data.len() as u32).to_le_bytes());
        packed.extend_from_slice(&(transformed_len as u32).to_le_bytes());
        packed.extend_from_slice(&(bwt_indices.len() as u32).to_le_bytes());
        for idx in bwt_indices {
            packed.extend_from_slice(&(idx as u32).to_le_bytes());
        }
        packed.extend_from_slice(data);
        packed
    }
}

/// Predictive Loom Decompressor.
pub struct LoomDecompressor {
    loom: Option<LoomEnum>,
    config: PlcConfig,
}

impl LoomDecompressor {
    pub fn new(config: PlcConfig) -> Self {
        let loom = if config.loom_mode == LoomMode::Off {
            None
        } else {
            LoomEnum::new(config.clone()).ok()
        };
        Self { loom, config }
    }
}

impl Decompressor for LoomDecompressor {
    type Error = LoomError;

    fn decompress(&mut self, data: &[u8], original_len: usize) -> Result<Vec<u8>, Self::Error> {
        if original_len == 0 {
            return Ok(Vec::new());
        }

        let mut output = Vec::with_capacity(original_len);
        let mut cursor = 0;

        #[cfg(feature = "std")]
        if self.config.verbose {
            std::println!(
                "[Decompress] Starting: data_len={}, original_len={}",
                data.len(),
                original_len
            );
        }

        while cursor < data.len() {
            // Header: [meta: u32][len_or_indices...]
            if cursor + 4 > data.len() {
                #[cfg(feature = "std")]
                if self.config.verbose {
                    std::println!(
                        "[Decompress] Breaking: cursor={} + 4 > data_len={}",
                        cursor,
                        data.len()
                    );
                }
                break;
            }
            let meta = u32::from_le_bytes(data[cursor..cursor + 4].try_into().unwrap());
            cursor += 4;

            #[cfg(feature = "std")]
            if self.config.verbose {
                std::println!(
                    "[Decompress] Chunk meta={:#010x}, cursor={}, is_raw={}",
                    meta,
                    cursor,
                    (meta & 0x8000_0000) != 0
                );
            }

            // --- No Expansion Detection ---
            if (meta & 0x8000_0000) != 0 {
                // Raw Chunk
                let _raw_len = (meta & 0x00FF_FFFF) as usize;
                let _file_type_val = (meta >> 24) & 0x7F;
                if cursor + 4 > data.len() {
                    return Err(LoomError::PackingError);
                }
                let compressed_len_field =
                    u32::from_le_bytes(data[cursor..cursor + 4].try_into().unwrap()) as usize;
                cursor += 4;

                if cursor + compressed_len_field > data.len() {
                    return Err(LoomError::PackingError);
                }
                output.extend_from_slice(&data[cursor..cursor + compressed_len_field]);
                cursor += compressed_len_field;
                continue;
            }

            // Loom Chunk
            let chunk_orig_len = (meta & 0x00FF_FFFF) as usize;
            let file_type_val = (meta >> 24) & 0x7F;
            let file_type: crate::poor_compress::DetectedType =
                unsafe { core::mem::transmute(file_type_val as u8) };
            if cursor + 12 > data.len() {
                return Err(LoomError::PackingError);
            }
            let compressed_len =
                u32::from_le_bytes(data[cursor..cursor + 4].try_into().unwrap()) as usize;
            let transformed_len =
                u32::from_le_bytes(data[cursor + 4..cursor + 8].try_into().unwrap()) as usize;
            let bwt_count =
                u32::from_le_bytes(data[cursor + 8..cursor + 12].try_into().unwrap()) as usize;
            cursor += 12;

            #[cfg(feature = "std")]
            if self.config.verbose {
                std::println!(
                    "[Decompress] Chunk: orig={}, compressed={}, transformed={}, bwt_count={}",
                    chunk_orig_len,
                    compressed_len,
                    transformed_len,
                    bwt_count
                );
            }

            let mut bwt_indices = Vec::with_capacity(bwt_count);
            for _ in 0..bwt_count {
                if cursor + 4 > data.len() {
                    return Err(LoomError::PackingError);
                }
                bwt_indices.push(
                    u32::from_le_bytes(data[cursor..cursor + 4].try_into().unwrap()) as usize,
                );
                cursor += 4;
            }

            if cursor + compressed_len > data.len() {
                return Err(LoomError::PackingError);
            }
            let ans_data = &data[cursor..cursor + compressed_len];
            let chunk_out = self.decompress_chunk(
                ans_data,
                chunk_orig_len,
                transformed_len,
                bwt_indices,
                file_type,
            )?;

            #[cfg(feature = "std")]
            if self.config.verbose {
                std::println!("[Decompress] Chunk output: {} bytes", chunk_out.len());
            }

            output.extend_from_slice(&chunk_out);
            cursor += compressed_len;
        }

        #[cfg(feature = "std")]
        if self.config.verbose {
            std::println!("[Decompress] Finished: output_len={}", output.len());
        }

        Ok(output)
    }
}

impl LoomDecompressor {
    fn decompress_chunk(
        &mut self,
        data: &[u8],
        _original_len: usize,
        transformed_len: usize,
        bwt_indices: Vec<usize>,
        file_type: crate::poor_compress::DetectedType,
    ) -> Result<Vec<u8>, LoomError> {
        if self.loom.is_none() || data.len() == transformed_len {
            // Revert transforms only (data is raw transformed bytes, not ANS)
            let mut transform_list = self.config.transforms.clone();

            // Re-apply same logic as compressor
            if transform_list.contains(&Transform::Auto) {
                match file_type {
                    crate::poor_compress::DetectedType::Jpeg => {
                        transform_list.insert(0, Transform::Yuv);
                    }
                    crate::poor_compress::DetectedType::Png
                    | crate::poor_compress::DetectedType::Gif => {
                        transform_list = vec![Transform::Delta, Transform::Rle];
                    }
                    crate::poor_compress::DetectedType::Mp4
                    | crate::poor_compress::DetectedType::Mp3 => {
                        transform_list = vec![
                            Transform::ChannelSeparation(2),
                            Transform::Delta,
                            Transform::Rle,
                        ];
                    }
                    crate::poor_compress::DetectedType::Encrypted => {
                        transform_list = vec![];
                    }
                    crate::poor_compress::DetectedType::Zip => {
                        transform_list = vec![Transform::Rle];
                    }
                    _ => {}
                }
            }

            let chain = TransformChain::new(transform_list);
            return Ok(chain.inverse(data.to_vec(), bwt_indices));
        }

        let mut transform_list = self.config.transforms.clone();
        if transform_list.contains(&Transform::Auto) {
            match file_type {
                crate::poor_compress::DetectedType::Jpeg => {
                    transform_list.insert(0, Transform::Yuv);
                }
                crate::poor_compress::DetectedType::Png
                | crate::poor_compress::DetectedType::Gif => {
                    transform_list = vec![Transform::Delta, Transform::Rle];
                }
                crate::poor_compress::DetectedType::Mp4
                | crate::poor_compress::DetectedType::Mp3 => {
                    transform_list = vec![
                        Transform::ChannelSeparation(2),
                        Transform::Delta,
                        Transform::Rle,
                    ];
                }
                crate::poor_compress::DetectedType::Encrypted => {
                    transform_list = vec![];
                }
                crate::poor_compress::DetectedType::Zip => {
                    transform_list = vec![Transform::Rle];
                }
                _ => {}
            }
        }
        let chain = TransformChain::new(transform_list);

        // Convert input bytes back to u16
        let mut input_u16 = Vec::with_capacity(data.len() / 2);
        for chunk in data.chunks_exact(2) {
            input_u16.push(u16::from_le_bytes(chunk.try_into().unwrap()));
        }

        let mut decoder = RansDecoder::new(&mut input_u16);
        let mut transformed = Vec::with_capacity(transformed_len);
        let loom = self.loom.as_mut().unwrap();

        for _ in 0..transformed_len {
            let model = loom.predict(&transformed);
            let sym = decoder.decode(&model, &mut input_u16);

            // Update Loom state
            loom.weave(&transformed, sym);
            transformed.push(sym);
        }

        Ok(chain.inverse(transformed, bwt_indices))
    }
}