tx2-iff 0.1.0

//! IFF file format structures and serialization
//!
//! The PPF-IFF format consists of:
//!
//! ```text
//! [Header]
//! [Layer 1: Wavelet Skeleton]
//! [Layer 2: Texture Synthesis Regions]
//! [Layer 3: Warp Field]
//! [Residual Data (sparse)]
//! ```
//!
//! ## File Layout
//!
//! ```text
//! Offset  Size  Description
//! 0       4     Magic ("PPFI")
//! 4       2     Version (major.minor)
//! 6       4     Width
//! 10      4     Height
//! 14      1     Wavelet levels
//! 15      1     Flags
//! 16      4     Layer 1 size (bytes)
//! 20      4     Layer 2 size (bytes)
//! 24      4     Layer 3 size (bytes)
//! 28      4     Residual size (bytes)
//! 32      ...   Layer data
//! ```

use crate::error::{IffError, Result};
use crate::texture::Region;
use crate::warp::WarpField;
use crate::wavelet::WaveletDecomposition;
use crate::compression::{compress_rle, decompress_rle};
use crate::MAGIC;
use serde::{Deserialize, Serialize};

/// IFF file version
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct Version {
    pub major: u8,
    pub minor: u8,
}

impl Version {
    /// Current version
    pub const CURRENT: Version = Version { major: 0, minor: 2 };

    /// Check if version is compatible
    pub fn is_compatible(&self, other: &Version) -> bool {
        self.major == other.major
    }
}

/// IFF file flags
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub struct Flags {
    /// Has alpha channel
    pub has_alpha: bool,
    /// Is lossless
    pub lossless: bool,
    /// Uses GPU-optimized format
    pub gpu_optimized: bool,
    /// Uses YCoCg color space with 4:2:0 subsampling
    pub ycocg_420: bool,
}

impl Default for Flags {
    fn default() -> Self {
        Flags {
            has_alpha: false,
            lossless: false,
            gpu_optimized: true,
            ycocg_420: true,
        }
    }
}

impl Flags {
    /// Convert to u8
    pub fn to_byte(&self) -> u8 {
        let mut byte = 0u8;
        if self.has_alpha {
            byte |= 0x01;
        }
        if self.lossless {
            byte |= 0x02;
        }
        if self.gpu_optimized {
            byte |= 0x04;
        }
        if self.ycocg_420 {
            byte |= 0x08;
        }
        byte
    }

    /// Create from u8
    pub fn from_byte(byte: u8) -> Self {
        Flags {
            has_alpha: (byte & 0x01) != 0,
            lossless: (byte & 0x02) != 0,
            gpu_optimized: (byte & 0x04) != 0,
            ycocg_420: (byte & 0x08) != 0,
        }
    }
}

/// IFF file header
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Header {
    /// Magic number (0x50504649 = "PPFI")
    pub magic: u32,
    /// File version
    pub version: Version,
    /// Image width
    pub width: u32,
    /// Image height
    pub height: u32,
    /// Number of wavelet decomposition levels
    pub wavelet_levels: u8,
    /// Flags
    pub flags: Flags,
    /// Layer 1 size in bytes
    pub layer1_size: u32,
    /// Layer 2 size in bytes
    pub layer2_size: u32,
    /// Layer 3 size in bytes
    pub layer3_size: u32,
    /// Residual size in bytes
    pub residual_size: u32,
}

impl Header {
    /// Header size in bytes (fixed)
    pub const SIZE: usize = 32;

    /// Create a new header
    pub fn new(width: u32, height: u32, wavelet_levels: u8) -> Self {
        Header {
            magic: MAGIC,
            version: Version::CURRENT,
            width,
            height,
            wavelet_levels,
            flags: Flags::default(),
            layer1_size: 0,
            layer2_size: 0,
            layer3_size: 0,
            residual_size: 0,
        }
    }

    /// Serialize header to bytes
    pub fn to_bytes(&self) -> Vec<u8> {
        let mut bytes = Vec::with_capacity(Self::SIZE);

        // Magic (4 bytes)
        bytes.extend_from_slice(&self.magic.to_be_bytes());

        // Version (2 bytes)
        bytes.push(self.version.major);
        bytes.push(self.version.minor);

        // Dimensions (8 bytes)
        bytes.extend_from_slice(&self.width.to_be_bytes());
        bytes.extend_from_slice(&self.height.to_be_bytes());

        // Wavelet levels (1 byte)
        bytes.push(self.wavelet_levels);

        // Flags (1 byte)
        bytes.push(self.flags.to_byte());

        // Layer sizes (16 bytes)
        bytes.extend_from_slice(&self.layer1_size.to_be_bytes());
        bytes.extend_from_slice(&self.layer2_size.to_be_bytes());
        bytes.extend_from_slice(&self.layer3_size.to_be_bytes());
        bytes.extend_from_slice(&self.residual_size.to_be_bytes());

        bytes
    }

    /// Deserialize header from bytes
    pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
        if bytes.len() < Self::SIZE {
            return Err(IffError::InsufficientData {
                expected: Self::SIZE,
                got: bytes.len(),
            });
        }

        // Parse magic
        let magic = u32::from_be_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]);
        if magic != MAGIC {
            return Err(IffError::InvalidMagic {
                expected: MAGIC,
                got: magic,
            });
        }

        // Parse version
        let version = Version {
            major: bytes[4],
            minor: bytes[5],
        };

        // Parse dimensions
        let width = u32::from_be_bytes([bytes[6], bytes[7], bytes[8], bytes[9]]);
        let height = u32::from_be_bytes([bytes[10], bytes[11], bytes[12], bytes[13]]);

        // Validate dimensions
        if width == 0 || height == 0 {
            return Err(IffError::Other("Invalid dimensions".to_string()));
        }

        // Parse wavelet levels
        let wavelet_levels = bytes[14];

        // Parse flags
        let flags = Flags::from_byte(bytes[15]);

        // Parse layer sizes
        let layer1_size = u32::from_be_bytes([bytes[16], bytes[17], bytes[18], bytes[19]]);
        let layer2_size = u32::from_be_bytes([bytes[20], bytes[21], bytes[22], bytes[23]]);
        let layer3_size = u32::from_be_bytes([bytes[24], bytes[25], bytes[26], bytes[27]]);
        let residual_size = u32::from_be_bytes([bytes[28], bytes[29], bytes[30], bytes[31]]);

        Ok(Header {
            magic,
            version,
            width,
            height,
            wavelet_levels,
            flags,
            layer1_size,
            layer2_size,
            layer3_size,
            residual_size,
        })
    }
}

/// Layer 1: Wavelet skeleton
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Layer1 {
    pub y: WaveletDecomposition,
    pub co: WaveletDecomposition,
    pub cg: WaveletDecomposition,
}

impl Layer1 {
    pub fn new(width: u32, height: u32, levels: usize) -> Self {
        Layer1 {
            y: WaveletDecomposition::new(width, height, levels, 1),
            co: WaveletDecomposition::new(width / 2, height / 2, levels, 1),
            cg: WaveletDecomposition::new(width / 2, height / 2, levels, 1),
        }
    }
}

/// Layer 2: Texture synthesis regions
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Layer2 {
    /// Texture regions
    pub regions: Vec<Region>,
}

impl Layer2 {
    /// Create empty layer
    pub fn new() -> Self {
        Layer2 {
            regions: Vec::new(),
        }
    }

    /// Add a region
    pub fn add_region(&mut self, region: Region) {
        self.regions.push(region);
    }
}

impl Default for Layer2 {
    fn default() -> Self {
        Self::new()
    }
}

/// Layer 3: Symplectic warp field
pub type Layer3 = WarpField;

/// Residual data (Compressed)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Residual {
    pub width: u32,
    pub height: u32,
    pub channels: u8,
    /// Compressed residual data (RLE encoded)
    /// Layout: [Channel 0] [Channel 1] [Channel 2]
    pub data: Vec<u8>,
}

impl Residual {
    /// Create empty residual
    pub fn new(width: u32, height: u32) -> Self {
        Residual {
            width,
            height,
            channels: 3,
            data: Vec::new(),
        }
    }

    /// Create from dense residual image
    pub fn from_dense(width: u32, height: u32, pixels: &[[i16; 3]]) -> Result<Self> {
        let mut channels = vec![Vec::with_capacity(pixels.len()); 3];
        
        for p in pixels {
            channels[0].push(p[0] as i32);
            channels[1].push(p[1] as i32);
            channels[2].push(p[2] as i32);
        }

        let mut all_data = Vec::new();
        all_data.extend_from_slice(&channels[0]);
        all_data.extend_from_slice(&channels[1]);
        all_data.extend_from_slice(&channels[2]);

        // First pass: RLE
        let rle_compressed = compress_rle(&all_data)?;
        
        // Second pass: Zstd
        // Level 3 is default, good balance. Level 0 uses default.
        let final_compressed = zstd::stream::encode_all(std::io::Cursor::new(rle_compressed), 3)
            .map_err(|e| IffError::Other(format!("Zstd compression failed: {}", e)))?;

        Ok(Residual {
            width,
            height,
            channels: 3,
            data: final_compressed,
        })
    }

    /// Decompress to dense residual image
    pub fn to_dense(&self) -> Result<Vec<[i16; 3]>> {
        let total_pixels = (self.width * self.height) as usize;
        let expected_len = total_pixels * 3;

        // First pass: Zstd decompression
        let rle_compressed = zstd::stream::decode_all(std::io::Cursor::new(&self.data))
            .map_err(|e| IffError::Other(format!("Zstd decompression failed: {}", e)))?;

        // Second pass: RLE decompression
        let all_data = decompress_rle(&rle_compressed, Some(expected_len))?;
        
        if all_data.len() < expected_len {
             return Err(IffError::Other("Insufficient residual data".to_string()));
        }

        let mut pixels = Vec::with_capacity(total_pixels);
        for i in 0..total_pixels {
            pixels.push([
                all_data[i] as i16,
                all_data[total_pixels + i] as i16,
                all_data[2 * total_pixels + i] as i16,
            ]);
        }

        Ok(pixels)
    }
}

/// Complete IFF image
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IffImage {
    /// File header
    pub header: Header,
    /// Layer 1: Wavelet skeleton
    pub layer1: Layer1,
    /// Layer 2: Texture synthesis
    pub layer2: Layer2,
    /// Layer 3: Warp field
    pub layer3: Layer3,
    /// Residual correction
    pub residual: Residual,
}

impl IffImage {
    /// Create a new IFF image
    pub fn new(width: u32, height: u32, wavelet_levels: u8) -> Self {
        IffImage {
            header: Header::new(width, height, wavelet_levels),
            layer1: Layer1::new(width, height, wavelet_levels as usize),
            layer2: Layer2::new(),
            layer3: WarpField::new(width, height),
            residual: Residual::new(width, height),
        }
    }

    /// Serialize to bytes
    pub fn to_bytes(&self) -> Result<Vec<u8>> {
        // Serialize layers
        let layer1_bytes = bincode::serialize(&self.layer1)?;
        let layer2_bytes = bincode::serialize(&self.layer2)?;
        let layer3_bytes = bincode::serialize(&self.layer3)?;
        let residual_bytes = bincode::serialize(&self.residual)?;

        // Update header with sizes
        let mut header = self.header.clone();
        header.layer1_size = layer1_bytes.len() as u32;
        header.layer2_size = layer2_bytes.len() as u32;
        header.layer3_size = layer3_bytes.len() as u32;
        header.residual_size = residual_bytes.len() as u32;

        // Combine all parts
        let mut bytes = header.to_bytes();
        bytes.extend_from_slice(&layer1_bytes);
        bytes.extend_from_slice(&layer2_bytes);
        bytes.extend_from_slice(&layer3_bytes);
        bytes.extend_from_slice(&residual_bytes);

        Ok(bytes)
    }

    /// Deserialize from bytes
    pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
        // Parse header
        let header = Header::from_bytes(bytes)?;

        // Check version compatibility
        if !header.version.is_compatible(&Version::CURRENT) {
            return Err(IffError::UnsupportedVersion(
                (header.version.major as u32) << 8 | header.version.minor as u32,
            ));
        }

        // Calculate offsets
        let offset1 = Header::SIZE;
        let offset2 = offset1 + header.layer1_size as usize;
        let offset3 = offset2 + header.layer2_size as usize;
        let offset_residual = offset3 + header.layer3_size as usize;
        let total_size = offset_residual + header.residual_size as usize;

        // Validate total size
        if bytes.len() < total_size {
            return Err(IffError::InsufficientData {
                expected: total_size,
                got: bytes.len(),
            });
        }

        // Deserialize layers
        let layer1: Layer1 = bincode::deserialize(&bytes[offset1..offset2])?;
        let layer2: Layer2 = bincode::deserialize(&bytes[offset2..offset3])?;
        let layer3: Layer3 = bincode::deserialize(&bytes[offset3..offset_residual])?;
        let residual: Residual = bincode::deserialize(&bytes[offset_residual..total_size])?;

        Ok(IffImage {
            header,
            layer1,
            layer2,
            layer3,
            residual,
        })
    }

    /// Get file size estimate
    pub fn estimated_size(&self) -> usize {
        Header::SIZE
            + self.layer1.y.data.len() + self.layer1.co.data.len() + self.layer1.cg.data.len()
            + self.layer2.regions.len() * 32
            + self.layer3.vortices.len() * 16
            + self.residual.data.len()
    }
}

/// Layer type identifier
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Layer {
    /// Wavelet skeleton
    Skeleton,
    /// Texture synthesis
    Texture,
    /// Warp field
    Warp,
    /// Residual correction
    Residual,
}

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

    #[test]
    fn test_version() {
        let v1 = Version { major: 0, minor: 2 };
        let v2 = Version { major: 0, minor: 2 };
        let v3 = Version { major: 1, minor: 0 };

        assert!(v1.is_compatible(&v2)); // Same major version
        assert!(!v1.is_compatible(&v3)); // Different major version
    }

    #[test]
    fn test_flags() {
        let flags = Flags {
            has_alpha: true,
            lossless: false,
            gpu_optimized: true,
            ycocg_420: true,
        };

        let byte = flags.to_byte();
        let restored = Flags::from_byte(byte);

        assert_eq!(flags.has_alpha, restored.has_alpha);
        assert_eq!(flags.lossless, restored.lossless);
        assert_eq!(flags.gpu_optimized, restored.gpu_optimized);
        assert_eq!(flags.ycocg_420, restored.ycocg_420);
    }
}