alec 1.3.6

// ALEC - Adaptive Lazy Evolving Compression
// Copyright (c) 2025 David Martin Venti
//
// Dual-licensed under AGPL-3.0 and Commercial License.
// See LICENSE file for details.

//! Preload file support for ALEC contexts
//!
//! This module provides the ability to save and load pre-trained context files
//! (preloads) that allow optimal compression from the first byte.
//!
//! # File Format
//!
//! The `.alec-context` file format uses little-endian binary encoding:
//!
//! ```text
//! ┌─────────────────────────────────────┐
//! │ Header (64 bytes)                   │
//! ├─────────────────────────────────────┤
//! │ Dictionary Section                  │
//! ├─────────────────────────────────────┤
//! │ Statistics Section                  │
//! ├─────────────────────────────────────┤
//! │ Prediction Model Section            │
//! └─────────────────────────────────────┘
//! ```

#[cfg(not(feature = "std"))]
use alloc::{
    string::{String, ToString},
    vec::Vec,
};

#[cfg(feature = "std")]
use crate::error::ContextError;
use crate::error::{AlecError, DecodeError};
#[cfg(feature = "std")]
use std::io::{Read, Write};
#[cfg(feature = "std")]
use std::path::Path;

/// Magic bytes for ALEC preload files
pub const PRELOAD_MAGIC: [u8; 4] = *b"ALEC";

/// Current preload file format version
pub const PRELOAD_FORMAT_VERSION: u32 = 1;

/// Header size in bytes
pub const PRELOAD_HEADER_SIZE: usize = 64;

/// Maximum sensor type string length
pub const MAX_SENSOR_TYPE_LEN: usize = 32;

/// Prediction model type stored in preload files
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
#[repr(u8)]
pub enum PreloadPredictionType {
    /// No prediction model
    #[default]
    None = 0,
    /// Use last observed value
    LastValue = 1,
    /// Linear prediction (a*x + b)
    Linear = 2,
    /// Moving average of last N values
    MovingAverage = 3,
    /// Periodic pattern
    Periodic = 4,
}

impl PreloadPredictionType {
    /// Convert from u8
    pub fn from_u8(value: u8) -> Option<Self> {
        match value {
            0 => Some(Self::None),
            1 => Some(Self::LastValue),
            2 => Some(Self::Linear),
            3 => Some(Self::MovingAverage),
            4 => Some(Self::Periodic),
            _ => None,
        }
    }
}

impl From<super::PredictionModel> for PreloadPredictionType {
    fn from(model: super::PredictionModel) -> Self {
        match model {
            super::PredictionModel::LastValue => Self::LastValue,
            super::PredictionModel::MovingAverage => Self::MovingAverage,
            super::PredictionModel::LinearRegression => Self::Linear,
            super::PredictionModel::Periodic => Self::Periodic,
        }
    }
}

impl From<PreloadPredictionType> for super::PredictionModel {
    fn from(ptype: PreloadPredictionType) -> Self {
        match ptype {
            PreloadPredictionType::None | PreloadPredictionType::LastValue => Self::LastValue,
            PreloadPredictionType::MovingAverage => Self::MovingAverage,
            PreloadPredictionType::Linear => Self::LinearRegression,
            PreloadPredictionType::Periodic => Self::Periodic,
        }
    }
}

/// Dictionary entry in preload file
#[derive(Debug, Clone, PartialEq)]
pub struct PreloadDictEntry {
    /// The byte pattern (1-255 bytes)
    pub pattern: Vec<u8>,
    /// Short code assigned to this pattern
    pub code: u16,
    /// Frequency count from training
    pub frequency: u32,
}

impl PreloadDictEntry {
    /// Serialize to bytes
    pub fn to_bytes(&self) -> Vec<u8> {
        let mut bytes = Vec::with_capacity(7 + self.pattern.len());
        bytes.push(self.pattern.len() as u8);
        bytes.extend_from_slice(&self.pattern);
        bytes.extend_from_slice(&self.code.to_le_bytes());
        bytes.extend_from_slice(&self.frequency.to_le_bytes());
        bytes
    }

    /// Deserialize from bytes, returns (entry, bytes_consumed)
    pub fn from_bytes(data: &[u8]) -> Result<(Self, usize), AlecError> {
        if data.is_empty() {
            return Err(DecodeError::BufferTooShort {
                needed: 1,
                available: 0,
            }
            .into());
        }

        let pattern_len = data[0] as usize;
        let total_len = 1 + pattern_len + 2 + 4; // len + pattern + code + frequency

        if data.len() < total_len {
            return Err(DecodeError::BufferTooShort {
                needed: total_len,
                available: data.len(),
            }
            .into());
        }

        let pattern = data[1..1 + pattern_len].to_vec();
        let code = u16::from_le_bytes([data[1 + pattern_len], data[2 + pattern_len]]);
        let frequency = u32::from_le_bytes([
            data[3 + pattern_len],
            data[4 + pattern_len],
            data[5 + pattern_len],
            data[6 + pattern_len],
        ]);

        Ok((
            Self {
                pattern,
                code,
                frequency,
            },
            total_len,
        ))
    }
}

/// Source statistics stored in preload files
#[derive(Debug, Clone, PartialEq)]
pub struct PreloadStatistics {
    /// Running mean of values
    pub mean: f64,
    /// Variance
    pub variance: f64,
    /// Minimum value observed in training
    pub min_observed: f64,
    /// Maximum value observed in training
    pub max_observed: f64,
    /// Expected minimum for anomaly detection
    pub min_expected: f64,
    /// Expected maximum for anomaly detection
    pub max_expected: f64,
    /// Recent values (circular buffer)
    pub recent_values: Vec<f64>,
}

impl Default for PreloadStatistics {
    fn default() -> Self {
        Self {
            mean: 0.0,
            variance: 0.0,
            min_observed: f64::MAX,
            max_observed: f64::MIN,
            min_expected: f64::MIN,
            max_expected: f64::MAX,
            recent_values: Vec::new(),
        }
    }
}

impl PreloadStatistics {
    /// Serialize to bytes
    pub fn to_bytes(&self) -> Vec<u8> {
        let mut bytes = Vec::with_capacity(49 + self.recent_values.len() * 8);

        bytes.extend_from_slice(&self.mean.to_le_bytes());
        bytes.extend_from_slice(&self.variance.to_le_bytes());
        bytes.extend_from_slice(&self.min_observed.to_le_bytes());
        bytes.extend_from_slice(&self.max_observed.to_le_bytes());
        bytes.extend_from_slice(&self.min_expected.to_le_bytes());
        bytes.extend_from_slice(&self.max_expected.to_le_bytes());

        bytes.push(self.recent_values.len() as u8);
        for val in &self.recent_values {
            bytes.extend_from_slice(&val.to_le_bytes());
        }

        bytes
    }

    /// Deserialize from bytes
    pub fn from_bytes(data: &[u8]) -> Result<(Self, usize), AlecError> {
        if data.len() < 49 {
            return Err(DecodeError::BufferTooShort {
                needed: 49,
                available: data.len(),
            }
            .into());
        }

        let mean = f64::from_le_bytes(data[0..8].try_into().unwrap());
        let variance = f64::from_le_bytes(data[8..16].try_into().unwrap());
        let min_observed = f64::from_le_bytes(data[16..24].try_into().unwrap());
        let max_observed = f64::from_le_bytes(data[24..32].try_into().unwrap());
        let min_expected = f64::from_le_bytes(data[32..40].try_into().unwrap());
        let max_expected = f64::from_le_bytes(data[40..48].try_into().unwrap());

        let history_len = data[48] as usize;
        let total_len = 49 + history_len * 8;

        if data.len() < total_len {
            return Err(DecodeError::BufferTooShort {
                needed: total_len,
                available: data.len(),
            }
            .into());
        }

        let mut recent_values = Vec::with_capacity(history_len);
        for i in 0..history_len {
            let offset = 49 + i * 8;
            let val = f64::from_le_bytes(data[offset..offset + 8].try_into().unwrap());
            recent_values.push(val);
        }

        Ok((
            Self {
                mean,
                variance,
                min_observed,
                max_observed,
                min_expected,
                max_expected,
                recent_values,
            },
            total_len,
        ))
    }
}

/// Prediction model parameters stored in preload files
#[derive(Debug, Clone, PartialEq)]
pub struct PreloadPredictionModel {
    /// Type of prediction model
    pub model_type: PreloadPredictionType,
    /// Coefficients for the model
    pub coefficients: Vec<f64>,
    /// Period length in samples (for periodic model)
    pub period_samples: u32,
}

impl Default for PreloadPredictionModel {
    fn default() -> Self {
        Self {
            model_type: PreloadPredictionType::LastValue,
            coefficients: Vec::new(),
            period_samples: 0,
        }
    }
}

impl PreloadPredictionModel {
    /// Serialize to bytes
    pub fn to_bytes(&self) -> Vec<u8> {
        let mut bytes = Vec::with_capacity(6 + self.coefficients.len() * 8);

        bytes.push(self.model_type as u8);
        bytes.push(self.coefficients.len() as u8);

        for coef in &self.coefficients {
            bytes.extend_from_slice(&coef.to_le_bytes());
        }

        bytes.extend_from_slice(&self.period_samples.to_le_bytes());

        bytes
    }

    /// Deserialize from bytes
    pub fn from_bytes(data: &[u8]) -> Result<(Self, usize), AlecError> {
        if data.len() < 6 {
            return Err(DecodeError::BufferTooShort {
                needed: 6,
                available: data.len(),
            }
            .into());
        }

        let model_type = PreloadPredictionType::from_u8(data[0]).unwrap_or_default();
        let coef_count = data[1] as usize;
        let total_len = 2 + coef_count * 8 + 4;

        if data.len() < total_len {
            return Err(DecodeError::BufferTooShort {
                needed: total_len,
                available: data.len(),
            }
            .into());
        }

        let mut coefficients = Vec::with_capacity(coef_count);
        for i in 0..coef_count {
            let offset = 2 + i * 8;
            let val = f64::from_le_bytes(data[offset..offset + 8].try_into().unwrap());
            coefficients.push(val);
        }

        let period_offset = 2 + coef_count * 8;
        let period_samples =
            u32::from_le_bytes(data[period_offset..period_offset + 4].try_into().unwrap());

        Ok((
            Self {
                model_type,
                coefficients,
                period_samples,
            },
            total_len,
        ))
    }
}

/// ALEC Preload Context File
///
/// A preload file contains a pre-trained context that can be loaded
/// at startup to achieve optimal compression from the first byte.
#[derive(Debug, Clone, PartialEq)]
pub struct PreloadFile {
    /// File format version
    pub format_version: u32,
    /// Context version for sync checking
    pub context_version: u32,
    /// Sensor type identifier
    pub sensor_type: String,
    /// Unix timestamp when the preload was created
    pub created_timestamp: u64,
    /// Number of samples used to train this context
    pub training_samples: u64,
    /// Dictionary entries
    pub dictionary: Vec<PreloadDictEntry>,
    /// Source statistics
    pub statistics: PreloadStatistics,
    /// Prediction model
    pub prediction: PreloadPredictionModel,
}

impl PreloadFile {
    /// Create a new preload file from a context
    pub fn from_context(ctx: &super::Context, sensor_type: &str) -> Self {
        // Collect dictionary entries from context
        let mut dictionary = Vec::new();
        for (&code, pattern) in ctx.patterns_iter() {
            if code <= u16::MAX as u32 {
                dictionary.push(PreloadDictEntry {
                    pattern: pattern.data.clone(),
                    code: code as u16,
                    frequency: pattern.frequency as u32,
                });
            }
        }

        // Sort by code for deterministic output
        dictionary.sort_by_key(|e| e.code);

        // Build statistics from context
        // Note: Context doesn't expose all stats, so we use defaults for some
        let statistics = PreloadStatistics::default();

        // Build prediction model
        let prediction = PreloadPredictionModel {
            model_type: ctx.model_type().into(),
            coefficients: Vec::new(),
            period_samples: 0,
        };

        Self {
            format_version: PRELOAD_FORMAT_VERSION,
            context_version: ctx.version(),
            sensor_type: sensor_type.to_string(),
            created_timestamp: {
                #[cfg(feature = "std")]
                {
                    std::time::SystemTime::now()
                        .duration_since(std::time::UNIX_EPOCH)
                        .map(|d| d.as_secs())
                        .unwrap_or(0)
                }
                #[cfg(not(feature = "std"))]
                {
                    0
                }
            },
            training_samples: ctx.observation_count(),
            dictionary,
            statistics,
            prediction,
        }
    }

    /// Serialize the preload file to bytes
    pub fn to_bytes(&self) -> Vec<u8> {
        let mut bytes = Vec::with_capacity(PRELOAD_HEADER_SIZE + 1024);

        // === HEADER (64 bytes) ===

        // Magic (4 bytes)
        bytes.extend_from_slice(&PRELOAD_MAGIC);

        // Format version (4 bytes)
        bytes.extend_from_slice(&self.format_version.to_le_bytes());

        // Context version (4 bytes)
        bytes.extend_from_slice(&self.context_version.to_le_bytes());

        // Sensor type length (2 bytes)
        let sensor_type_bytes = self.sensor_type.as_bytes();
        let sensor_type_len = sensor_type_bytes.len().min(MAX_SENSOR_TYPE_LEN) as u16;
        bytes.extend_from_slice(&sensor_type_len.to_le_bytes());

        // Sensor type (up to 32 bytes, padded)
        let mut sensor_type_buf = [0u8; MAX_SENSOR_TYPE_LEN];
        sensor_type_buf[..sensor_type_len as usize]
            .copy_from_slice(&sensor_type_bytes[..sensor_type_len as usize]);
        bytes.extend_from_slice(&sensor_type_buf);

        // Created timestamp (8 bytes)
        bytes.extend_from_slice(&self.created_timestamp.to_le_bytes());

        // Training samples (8 bytes)
        bytes.extend_from_slice(&self.training_samples.to_le_bytes());

        // Placeholder for checksum (4 bytes) - will be filled at the end
        let checksum_offset = bytes.len();
        bytes.extend_from_slice(&[0u8; 4]);

        // Reserved (remaining bytes to reach 64)
        // 4 + 4 + 4 + 2 + 32 + 8 + 8 + 4 = 66 - but header is 64, so we need 0 reserved
        // Actually let's recalculate: 4+4+4+2+32+8+8+4 = 66, so we're 2 bytes over
        // Let's adjust: remove 2 bytes from reserved or truncate sensor type
        // For simplicity, let's just continue (header can be slightly larger)

        // === DICTIONARY SECTION ===
        bytes.extend_from_slice(&(self.dictionary.len() as u32).to_le_bytes());
        for entry in &self.dictionary {
            bytes.extend_from_slice(&entry.to_bytes());
        }

        // === STATISTICS SECTION ===
        bytes.extend_from_slice(&self.statistics.to_bytes());

        // === PREDICTION MODEL SECTION ===
        bytes.extend_from_slice(&self.prediction.to_bytes());

        // Calculate and insert checksum (CRC32 of everything except the checksum field)
        use crc::{Crc, CRC_32_ISO_HDLC};
        const CRC32: Crc<u32> = Crc::<u32>::new(&CRC_32_ISO_HDLC);
        let mut digest = CRC32.digest();
        digest.update(&bytes[..checksum_offset]);
        digest.update(&bytes[checksum_offset + 4..]);
        let checksum = digest.finalize();
        bytes[checksum_offset..checksum_offset + 4].copy_from_slice(&checksum.to_le_bytes());

        bytes
    }

    /// Deserialize a preload file from bytes
    pub fn from_bytes(data: &[u8]) -> Result<Self, AlecError> {
        if data.len() < 66 {
            return Err(DecodeError::BufferTooShort {
                needed: 66,
                available: data.len(),
            }
            .into());
        }

        // Verify magic
        if data[0..4] != PRELOAD_MAGIC {
            return Err(DecodeError::MalformedMessage {
                offset: 0,
                reason: "Invalid magic bytes".to_string(),
            }
            .into());
        }

        // Parse header
        let format_version = u32::from_le_bytes(data[4..8].try_into().unwrap());
        let context_version = u32::from_le_bytes(data[8..12].try_into().unwrap());
        let sensor_type_len = u16::from_le_bytes(data[12..14].try_into().unwrap()) as usize;

        let sensor_type_end = (14 + sensor_type_len).min(14 + MAX_SENSOR_TYPE_LEN);
        let sensor_type = String::from_utf8_lossy(&data[14..sensor_type_end])
            .trim_end_matches('\0')
            .to_string();

        let created_timestamp = u64::from_le_bytes(data[46..54].try_into().unwrap());
        let training_samples = u64::from_le_bytes(data[54..62].try_into().unwrap());
        let stored_checksum = u32::from_le_bytes(data[62..66].try_into().unwrap());

        // Verify checksum
        use crc::{Crc, CRC_32_ISO_HDLC};
        const CRC32: Crc<u32> = Crc::<u32>::new(&CRC_32_ISO_HDLC);
        let mut digest = CRC32.digest();
        digest.update(&data[..62]);
        digest.update(&data[66..]);
        let computed_checksum = digest.finalize();

        if stored_checksum != computed_checksum {
            return Err(DecodeError::InvalidChecksum {
                expected: stored_checksum,
                actual: computed_checksum,
            }
            .into());
        }

        let mut offset = 66;

        // === DICTIONARY SECTION ===
        if data.len() < offset + 4 {
            return Err(DecodeError::BufferTooShort {
                needed: offset + 4,
                available: data.len(),
            }
            .into());
        }

        let dict_count = u32::from_le_bytes(data[offset..offset + 4].try_into().unwrap()) as usize;
        offset += 4;

        let mut dictionary = Vec::with_capacity(dict_count);
        for _ in 0..dict_count {
            let (entry, consumed) = PreloadDictEntry::from_bytes(&data[offset..])?;
            dictionary.push(entry);
            offset += consumed;
        }

        // === STATISTICS SECTION ===
        let (statistics, consumed) = PreloadStatistics::from_bytes(&data[offset..])?;
        offset += consumed;

        // === PREDICTION MODEL SECTION ===
        let (prediction, _) = PreloadPredictionModel::from_bytes(&data[offset..])?;

        Ok(Self {
            format_version,
            context_version,
            sensor_type,
            created_timestamp,
            training_samples,
            dictionary,
            statistics,
            prediction,
        })
    }

    /// Save preload to a file
    #[cfg(feature = "std")]
    pub fn save_to_file(&self, path: &Path) -> Result<(), AlecError> {
        let bytes = self.to_bytes();
        let mut file = std::fs::File::create(path).map_err(|e| ContextError::SyncFailed {
            reason: format!("Failed to create preload file: {}", e),
        })?;
        file.write_all(&bytes)
            .map_err(|e| ContextError::SyncFailed {
                reason: format!("Failed to write preload file: {}", e),
            })?;
        Ok(())
    }

    /// Load preload from a file
    #[cfg(feature = "std")]
    pub fn load_from_file(path: &Path) -> Result<Self, AlecError> {
        let mut file = std::fs::File::open(path).map_err(|e| ContextError::SyncFailed {
            reason: format!("Failed to open preload file: {}", e),
        })?;
        let mut bytes = Vec::new();
        file.read_to_end(&mut bytes)
            .map_err(|e| ContextError::SyncFailed {
                reason: format!("Failed to read preload file: {}", e),
            })?;
        Self::from_bytes(&bytes)
    }
}

/// Version synchronization result
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum VersionCheckResult {
    /// Versions match, proceed normally
    Match,
    /// Version mismatch detected
    Mismatch {
        /// Expected version (decoder's context version)
        expected: u32,
        /// Actual version from message
        actual: u32,
    },
}

impl VersionCheckResult {
    /// Check if versions match
    pub fn is_match(&self) -> bool {
        matches!(self, Self::Match)
    }
}

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

    #[test]
    fn test_dict_entry_roundtrip() {
        let entry = PreloadDictEntry {
            pattern: vec![0x42, 0x43, 0x44],
            code: 123,
            frequency: 456789,
        };

        let bytes = entry.to_bytes();
        let (restored, consumed) = PreloadDictEntry::from_bytes(&bytes).unwrap();

        assert_eq!(entry, restored);
        assert_eq!(consumed, bytes.len());
    }

    #[test]
    fn test_statistics_roundtrip() {
        let stats = PreloadStatistics {
            mean: 25.5,
            variance: 0.1,
            min_observed: 20.0,
            max_observed: 30.0,
            min_expected: 15.0,
            max_expected: 35.0,
            recent_values: vec![24.0, 25.0, 26.0],
        };

        let bytes = stats.to_bytes();
        let (restored, consumed) = PreloadStatistics::from_bytes(&bytes).unwrap();

        assert_eq!(stats, restored);
        assert_eq!(consumed, bytes.len());
    }

    #[test]
    fn test_prediction_model_roundtrip() {
        let model = PreloadPredictionModel {
            model_type: PreloadPredictionType::Linear,
            coefficients: vec![0.1, 2.5],
            period_samples: 100,
        };

        let bytes = model.to_bytes();
        let (restored, consumed) = PreloadPredictionModel::from_bytes(&bytes).unwrap();

        assert_eq!(model, restored);
        assert_eq!(consumed, bytes.len());
    }

    #[test]
    fn test_preload_file_roundtrip() {
        let preload = PreloadFile {
            format_version: PRELOAD_FORMAT_VERSION,
            context_version: 42,
            sensor_type: "temperature".to_string(),
            created_timestamp: 1234567890,
            training_samples: 10000,
            dictionary: vec![
                PreloadDictEntry {
                    pattern: vec![0x00, 0x01],
                    code: 0,
                    frequency: 100,
                },
                PreloadDictEntry {
                    pattern: vec![0x02, 0x03, 0x04],
                    code: 1,
                    frequency: 50,
                },
            ],
            statistics: PreloadStatistics {
                mean: 22.5,
                variance: 0.5,
                min_observed: 18.0,
                max_observed: 28.0,
                min_expected: 15.0,
                max_expected: 35.0,
                recent_values: vec![22.0, 22.5, 23.0],
            },
            prediction: PreloadPredictionModel {
                model_type: PreloadPredictionType::MovingAverage,
                coefficients: vec![5.0],
                period_samples: 0,
            },
        };

        let bytes = preload.to_bytes();
        let restored = PreloadFile::from_bytes(&bytes).unwrap();

        assert_eq!(preload.format_version, restored.format_version);
        assert_eq!(preload.context_version, restored.context_version);
        assert_eq!(preload.sensor_type, restored.sensor_type);
        assert_eq!(preload.created_timestamp, restored.created_timestamp);
        assert_eq!(preload.training_samples, restored.training_samples);
        assert_eq!(preload.dictionary.len(), restored.dictionary.len());
    }

    #[test]
    fn test_invalid_magic() {
        let mut data = vec![0u8; 100];
        data[0..4].copy_from_slice(b"BADM");

        let result = PreloadFile::from_bytes(&data);
        assert!(result.is_err());
    }

    #[test]
    fn test_version_check_result() {
        assert!(VersionCheckResult::Match.is_match());
        assert!(!VersionCheckResult::Mismatch {
            expected: 1,
            actual: 2
        }
        .is_match());
    }
}