path: root/src/protocol/jpeg.rs

//! JPEG parsing utilities for the UPP protocol

use crate::error::{JpegError, Result};
use tracing::{debug, trace, warn};

/// JPEG marker constants
const JPEG_SOI: u8 = 0xD8; // Start of Image
const JPEG_EOI: u8 = 0xD9; // End of Image
const JPEG_SOS: u8 = 0xDA; // Start of Scan
const JPEG_SOF0: u8 = 0xC0; // Start of Frame (Baseline DCT)
const JPEG_SOF1: u8 = 0xC1; // Start of Frame (Extended sequential DCT)
const JPEG_SOF2: u8 = 0xC2; // Start of Frame (Progressive DCT)

/// JPEG image dimensions
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct JpegDimensions {
    pub width: u16,
    pub height: u16,
}

impl JpegDimensions {
    /// Create new dimensions
    pub fn new(width: u16, height: u16) -> Self {
        Self { width, height }
    }

    /// Check if dimensions are valid
    pub fn is_valid(&self) -> bool {
        self.width > 0 && self.height > 0
    }

    /// Get aspect ratio
    pub fn aspect_ratio(&self) -> f64 {
        if self.height > 0 {
            self.width as f64 / self.height as f64
        } else {
            0.0
        }
    }
}

/// JPEG parser for extracting metadata
pub struct JpegParser {
    enable_debug: bool,
}

impl JpegParser {
    /// Create a new JPEG parser
    pub fn new() -> Self {
        Self {
            enable_debug: false,
        }
    }

    /// Create a new JPEG parser with debug enabled
    pub fn new_with_debug(enable_debug: bool) -> Self {
        Self { enable_debug }
    }

    /// Parse JPEG dimensions from raw data
    pub fn parse_dimensions(&self, data: &[u8]) -> Result<JpegDimensions> {
        if data.len() < 4 {
            return Err(JpegError::InvalidHeader.into());
        }

        // Check JPEG SOI marker
        if data[0] != 0xFF || data[1] != JPEG_SOI {
            return Err(JpegError::InvalidHeader.into());
        }

        trace!("Parsing JPEG dimensions from {} bytes", data.len());

        let mut i = 2;
        while i + 3 < data.len() {
            // Look for marker
            if data[i] != 0xFF {
                i += 1;
                continue;
            }

            let marker = data[i + 1];
            i += 2;

            // Check for end markers
            if marker == JPEG_EOI || marker == JPEG_SOS {
                break;
            }

            // Check if we have enough data for segment length
            if i + 1 >= data.len() {
                break;
            }

            // Read segment length (big-endian)
            let segment_length = ((data[i] as u16) << 8) | (data[i + 1] as u16);
            if segment_length < 2 || i + segment_length as usize > data.len() {
                warn!("Invalid segment length: {}", segment_length);
                break;
            }

            // Check for SOF markers (Start of Frame)
            if self.is_sof_marker(marker) {
                if segment_length < 7 {
                    return Err(JpegError::InvalidHeader.into());
                }

                // Height and width are in big-endian format
                let height = ((data[i + 3] as u16) << 8) | (data[i + 4] as u16);
                let width = ((data[i + 5] as u16) << 8) | (data[i + 6] as u16);

                if self.enable_debug {
                    debug!(
                        "Found SOF marker 0x{:02X}, dimensions: {}x{}",
                        marker, width, height
                    );
                }

                if width > 0 && height > 0 {
                    return Ok(JpegDimensions::new(width, height));
                } else {
                    return Err(JpegError::DimensionsNotFound.into());
                }
            }

            // Move to next segment
            i += segment_length as usize;
        }

        Err(JpegError::DimensionsNotFound.into())
    }

    /// Check if a marker is a Start of Frame marker
    fn is_sof_marker(&self, marker: u8) -> bool {
        matches!(marker, JPEG_SOF0 | JPEG_SOF1 | JPEG_SOF2)
    }

    /// Extract JPEG metadata
    pub fn parse_metadata(&self, data: &[u8]) -> Result<JpegMetadata> {
        let dimensions = self.parse_dimensions(data)?;

        Ok(JpegMetadata {
            dimensions,
            file_size: data.len(),
            is_valid: true,
        })
    }

    /// Validate JPEG data
    pub fn validate_jpeg(&self, data: &[u8]) -> Result<bool> {
        if data.len() < 4 {
            return Ok(false);
        }

        // Check SOI marker
        if data[0] != 0xFF || data[1] != JPEG_SOI {
            return Ok(false);
        }

        // Check EOI marker (should be near the end)
        if data.len() >= 2 {
            let end = data.len() - 2;
            if data[end] == 0xFF && data[end + 1] == JPEG_EOI {
                return Ok(true);
            }
        }

        // If we can't find EOI, check if we can parse dimensions
        Ok(self.parse_dimensions(data).is_ok())
    }

    /// Check if data represents a complete JPEG frame
    pub fn is_complete_jpeg(&self, data: &[u8]) -> bool {
        if data.len() < 4 {
            return false;
        }

        // Must start with SOI marker
        if data[0] != 0xFF || data[1] != JPEG_SOI {
            return false;
        }

        // Must end with EOI marker
        if data.len() >= 2 {
            let end = data.len() - 2;
            if data[end] == 0xFF && data[end + 1] == JPEG_EOI {
                return true;
            }
        }

        false
    }
}

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

/// JPEG metadata
#[derive(Debug, Clone)]
pub struct JpegMetadata {
    pub dimensions: JpegDimensions,
    pub file_size: usize,
    pub is_valid: bool,
}

impl JpegMetadata {
    /// Create new metadata
    pub fn new(dimensions: JpegDimensions, file_size: usize) -> Self {
        Self {
            dimensions,
            file_size,
            is_valid: true,
        }
    }

    /// Get estimated bit depth (assume 8-bit for most JPEGs)
    pub fn estimated_bit_depth(&self) -> u8 {
        8
    }

    /// Get estimated color space (assume YUV for most JPEGs)
    pub fn estimated_color_space(&self) -> &'static str {
        "YUV"
    }
}

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

    // Sample JPEG data for testing (minimal valid JPEG)
    const MINIMAL_JPEG: &[u8] = &[
        0xFF, 0xD8, // SOI
        0xFF, 0xE0, // APP0
        0x00, 0x10, // Length
        0x4A, 0x46, 0x49, 0x46, 0x00, // "JFIF\0"
        0x01, 0x01, // Version
        0x00, // Units
        0x00, 0x01, // Density
        0x00, 0x01, 0x00, 0x00, // No thumbnail
        0xFF, 0xC0, // SOF0
        0x00, 0x0B, // Length
        0x08, // Precision
        0x00, 0x40, // Height (64)
        0x00, 0x40, // Width (64)
        0x03, // Components
        0x01, 0x11, 0x00, // Y component
        0x02, 0x11, 0x01, // U component
        0x03, 0x11, 0x01, // V component
        0xFF, 0xD9, // EOI
    ];

    #[test]
    fn test_jpeg_dimensions_creation() {
        let dims = JpegDimensions::new(640, 480);
        assert_eq!(dims.width, 640);
        assert_eq!(dims.height, 480);
        assert!(dims.is_valid());
        assert_eq!(dims.aspect_ratio(), 640.0 / 480.0);
    }

    #[test]
    fn test_jpeg_dimensions_validation() {
        let dims = JpegDimensions::new(0, 480);
        assert!(!dims.is_valid());

        let dims = JpegDimensions::new(640, 0);
        assert!(!dims.is_valid());

        let dims = JpegDimensions::new(0, 0);
        assert!(!dims.is_valid());
    }

    #[test]
    fn test_jpeg_parser_creation() {
        let parser = JpegParser::new();
        assert!(!parser.enable_debug);

        let parser = JpegParser::new_with_debug(true);
        assert!(parser.enable_debug);
    }

    #[test]
    fn test_jpeg_parser_parse_dimensions() {
        let parser = JpegParser::new();
        let dimensions = parser.parse_dimensions(MINIMAL_JPEG).unwrap();
        assert_eq!(dimensions.width, 64);
        assert_eq!(dimensions.height, 64);
    }

    #[test]
    fn test_jpeg_parser_invalid_header() {
        let parser = JpegParser::new();
        let invalid_data = &[0x00, 0x01, 0x02, 0x03];

        let result = parser.parse_dimensions(invalid_data);
        assert!(result.is_err());
        assert!(matches!(
            result.unwrap_err(),
            Error::Jpeg(JpegError::InvalidHeader)
        ));
    }

    #[test]
    fn test_jpeg_parser_short_data() {
        let parser = JpegParser::new();
        let short_data = &[0xFF, 0xD8];

        let result = parser.parse_dimensions(short_data);
        assert!(result.is_err());
        assert!(matches!(
            result.unwrap_err(),
            Error::Jpeg(JpegError::InvalidHeader)
        ));
    }

    #[test]
    fn test_jpeg_parser_validate_jpeg() {
        let parser = JpegParser::new();
        assert!(parser.validate_jpeg(MINIMAL_JPEG).unwrap());

        let invalid_data = &[0x00, 0x01, 0x02, 0x03];
        assert!(!parser.validate_jpeg(invalid_data).unwrap());
    }

    #[test]
    fn test_jpeg_metadata_creation() {
        let dimensions = JpegDimensions::new(640, 480);
        let metadata = JpegMetadata::new(dimensions, 1024);

        assert_eq!(metadata.dimensions.width, 640);
        assert_eq!(metadata.dimensions.height, 480);
        assert_eq!(metadata.file_size, 1024);
        assert!(metadata.is_valid);
        assert_eq!(metadata.estimated_bit_depth(), 8);
        assert_eq!(metadata.estimated_color_space(), "YUV");
    }

    #[test]
    fn test_sof_marker_detection() {
        let parser = JpegParser::new();
        assert!(parser.is_sof_marker(JPEG_SOF0));
        assert!(parser.is_sof_marker(JPEG_SOF1));
        assert!(parser.is_sof_marker(JPEG_SOF2));
        assert!(!parser.is_sof_marker(0x00));
        assert!(!parser.is_sof_marker(JPEG_SOI));
    }
}