//! Compression functionality for TACO format
//!
//! This module provides compression settings and methods for compressing
//! and decompressing trajectory data, supporting both lossless and lossy compression.

use ndarray::Array2;
use serde::{Deserialize, Serialize};
use std::io::{Read, Write};

use crate::{Error, Result};

/// Precision modes for compression
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub enum PrecisionMode {
    /// Full precision (32-bit float)
    Full,

    /// Half precision (16-bit float)
    Half,
}

/// Settings for compression of trajectory data
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CompressionSettings {
    /// Precision mode for all data (positions, velocities, and forces)
    pub precision: PrecisionMode,

    /// Compression level for Zstandard (1-22, higher = better compression but slower)
    pub zstd_level: i32,
}

impl Default for CompressionSettings {
    fn default() -> Self {
        CompressionSettings {
            precision: PrecisionMode::Half, // Half precision for all data
            zstd_level: 1,                  // Default compression level
        }
    }
}

impl CompressionSettings {
    /// Create a new compression settings instance with default values
    pub fn new() -> Self {
        Default::default()
    }

    /// Create settings for maximum compression (slower but smaller files)
    pub fn max_compression() -> Self {
        CompressionSettings {
            zstd_level: 19,
            ..Default::default()
        }
    }

    /// Create settings for fast compression (faster but larger files)
    pub fn fast_compression() -> Self {
        CompressionSettings {
            zstd_level: -7,
            ..Default::default()
        }
    }

    /// Create settings for lossless compression
    pub fn lossless() -> Self {
        CompressionSettings {
            precision: PrecisionMode::Full,
            zstd_level: 3,
        }
    }
}

/// Compress data using the specified precision mode
pub fn compress_tensor(
    data: &Array2<f32>,
    precision: PrecisionMode,
    zstd_level: i32,
) -> Result<Vec<u8>> {
    // First serialize data with precision mode
    let mut buffer = Vec::new();

    // Write dimensions
    let shape = data.shape();
    buffer.extend_from_slice(&(shape[0] as u32).to_le_bytes());
    buffer.extend_from_slice(&(shape[1] as u32).to_le_bytes());

    // Write data based on precision
    match precision {
        PrecisionMode::Full => {
            // Full precision (32-bit float)
            for value in data.iter() {
                buffer.extend_from_slice(&value.to_le_bytes());
            }
        }
        PrecisionMode::Half => {
            // Half precision (16-bit float)
            for value in data.iter() {
                let half = half::f16::from_f32(*value);
                buffer.extend_from_slice(&half.to_bits().to_le_bytes());
            }
        }
    }

    // Compress with zstd
    let mut compressed = Vec::new();
    let mut encoder = zstd::stream::Encoder::new(&mut compressed, zstd_level)?;
    encoder.write_all(&buffer)?;
    encoder.finish()?;

    Ok(compressed)
}

/// Decompress data
pub fn decompress_tensor(data: &[u8], precision: PrecisionMode) -> Result<Array2<f32>> {
    // Decompress with zstd
    let mut decoder = zstd::stream::Decoder::new(data)?;
    let mut decompressed = Vec::new();
    decoder.read_to_end(&mut decompressed)?;

    // Read dimensions
    if decompressed.len() < 8 {
        return Err(Error::Decompression(
            "Data too short to contain dimensions".to_string(),
        ));
    }

    let rows = u32::from_le_bytes([
        decompressed[0],
        decompressed[1],
        decompressed[2],
        decompressed[3],
    ]) as usize;
    let cols = u32::from_le_bytes([
        decompressed[4],
        decompressed[5],
        decompressed[6],
        decompressed[7],
    ]) as usize;

    let mut result = Array2::<f32>::zeros((rows, cols));
    let offset = 8; // Start after dimensions

    // Read data based on precision
    match precision {
        PrecisionMode::Full => {
            // Full precision (32-bit float)
            let bytes_per_value: usize = 4;
            let expected = (rows as u64)
                .checked_mul(cols as u64)
                .and_then(|n| n.checked_mul(bytes_per_value as u64))
                .ok_or_else(|| {
                    Error::Decompression("Full precision, dimension overflow".to_string())
                })?;
            if (decompressed.len() as u64) < (offset as u64) + expected {
                return Err(Error::Decompression(
                    "Full precision, data too short for specified dimensions".to_string(),
                ));
            }

            for i in 0..rows {
                for j in 0..cols {
                    let start = offset + (i * cols + j) * bytes_per_value;
                    let value = f32::from_le_bytes([
                        decompressed[start],
                        decompressed[start + 1],
                        decompressed[start + 2],
                        decompressed[start + 3],
                    ]);
                    result[(i, j)] = value;
                }
            }
        }
        PrecisionMode::Half => {
            // Half precision (16-bit float)
            let bytes_per_value: usize = 2;
            let expected = (rows as u64)
                .checked_mul(cols as u64)
                .and_then(|n| n.checked_mul(bytes_per_value as u64))
                .ok_or_else(|| {
                    Error::Decompression("Half precision, dimension overflow".to_string())
                })?;
            if (decompressed.len() as u64) < (offset as u64) + expected {
                return Err(Error::Decompression(
                    "Half precision, data too short for specified dimensions".to_string(),
                ));
            }

            for i in 0..rows {
                for j in 0..cols {
                    let start = offset + (i * cols + j) * bytes_per_value;
                    let bits = u16::from_le_bytes([decompressed[start], decompressed[start + 1]]);
                    let half = half::f16::from_bits(bits);
                    result[(i, j)] = half.to_f32();
                }
            }
        }
    }

    Ok(result)
}

/// Compress f64 tensor data (always full precision)
pub fn compress_tensor_f64(data: &Array2<f64>, zstd_level: i32) -> Result<Vec<u8>> {
    let mut buffer = Vec::new();
    let shape = data.shape();
    buffer.extend_from_slice(&(shape[0] as u32).to_le_bytes());
    buffer.extend_from_slice(&(shape[1] as u32).to_le_bytes());

    for value in data.iter() {
        buffer.extend_from_slice(&value.to_le_bytes());
    }

    let mut compressed = Vec::new();
    let mut encoder = zstd::stream::Encoder::new(&mut compressed, zstd_level)?;
    encoder.write_all(&buffer)?;
    encoder.finish()?;

    Ok(compressed)
}

/// Decompress f64 tensor data
pub fn decompress_tensor_f64(data: &[u8]) -> Result<Array2<f64>> {
    let mut decoder = zstd::stream::Decoder::new(data)?;
    let mut decompressed = Vec::new();
    decoder.read_to_end(&mut decompressed)?;

    if decompressed.len() < 8 {
        return Err(Error::Decompression(
            "Data too short to contain dimensions".to_string(),
        ));
    }

    let rows = u32::from_le_bytes([
        decompressed[0],
        decompressed[1],
        decompressed[2],
        decompressed[3],
    ]) as usize;
    let cols = u32::from_le_bytes([
        decompressed[4],
        decompressed[5],
        decompressed[6],
        decompressed[7],
    ]) as usize;

    let mut result = Array2::<f64>::zeros((rows, cols));
    let offset: usize = 8;
    let bytes_per_value: usize = 8;
    let expected = (rows as u64)
        .checked_mul(cols as u64)
        .and_then(|n| n.checked_mul(bytes_per_value as u64))
        .ok_or_else(|| Error::Decompression("f64 dimension overflow".to_string()))?;
    if (decompressed.len() as u64) < (offset as u64) + expected {
        return Err(Error::Decompression(
            "Full precision, data too short for specified dimensions".to_string(),
        ));
    }

    for i in 0..rows {
        for j in 0..cols {
            let start = offset + (i * cols + j) * bytes_per_value;
            let value = f64::from_le_bytes([
                decompressed[start],
                decompressed[start + 1],
                decompressed[start + 2],
                decompressed[start + 3],
                decompressed[start + 4],
                decompressed[start + 5],
                decompressed[start + 6],
                decompressed[start + 7],
            ]);
            result[(i, j)] = value;
        }
    }

    Ok(result)
}
