#!/usr/bin/env python3
"""
Efficient JSON extraction from JPEG images using EXIF metadata.
Supports large JSON payloads through chunking and decompression.

Copyright (c) 2024-2025 Overview Corporation. All rights reserved.
"""

import logging
import struct
import zlib
from dataclasses import dataclass
from io import BytesIO
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union, cast

import orjson

# Configure logging
logger = logging.getLogger(__name__)


# Custom exceptions
class NoMetadataError(Exception):
    """Raised when the image contains no JSON metadata."""


class MetadataDecodingError(Exception):
    """Raised when metadata exists but could not be decoded properly."""


# JPEG markers
JPEG_SOI = b'\xff\xd8'  # Start of Image
JPEG_APP1 = b'\xff\xe1'  # EXIF/XMP marker
JPEG_SOS = b'\xff\xda'  # Start of Scan (image data begins)
JPEG_EOI = b'\xff\xd9'  # End of Image

# Custom identifier for our JSON data
JSON_EXIF_HEADER = b'JSON\x00\x00'


@dataclass
class ChunkInfo:
    """Information about a data chunk."""

    index: int
    total: int
    data: bytes


class JpegJsonReader:
    """
    Efficient JSON extraction from JPEG images using EXIF metadata.
    Supports large JSON payloads through chunking and decompression.
    """

    def __init__(self, source: Union[str, Path, bytes, BytesIO]):
        """
        Initialize with a JPEG image from file path or binary data.

        Args:
            source: File path (str/Path) or binary JPEG data (bytes/BytesIO)

        Raises:
            ValueError: If not a valid JPEG image
            FileNotFoundError: If file path doesn't exist
        """
        self._embedded_json: Optional[Dict[str, Any]] = None
        self._source_path: Optional[Path] = None

        if isinstance(source, (str, Path)):
            self._source_path = Path(source)
            if not self._source_path.exists():
                raise FileNotFoundError(f'Image not found: {self._source_path}')

            logger.info(f'Loading JPEG from file: {self._source_path}')
            with open(self._source_path, 'rb') as f:
                self._image_data = f.read()
            logger.info(f'Loaded {len(self._image_data)} bytes from file')
        elif isinstance(source, bytes):
            logger.debug(f'Loading JPEG from bytes: {len(source)} bytes')
            self._image_data = source
        elif isinstance(source, BytesIO):
            logger.debug('Loading JPEG from BytesIO')
            self._image_data = source.getvalue()
            logger.debug(f'Loaded {len(self._image_data)} bytes from BytesIO')
        else:
            raise ValueError('Source must be file path (str/Path) or binary data (bytes/BytesIO)')

        self._validate_jpeg()

    @property
    def image_data(self) -> bytes:
        """Get the current image data."""
        return self._image_data

    @property
    def source_path(self) -> Optional[Path]:
        """Get the original source path if provided."""
        return self._source_path

    def _validate_jpeg(self) -> None:
        """Validate that the image data is a valid JPEG."""
        if len(self._image_data) < 2 or self._image_data[:2] != JPEG_SOI:
            raise ValueError('Not a valid JPEG image')

    def _decompress_json(self, data: bytes) -> Dict[str, Any]:
        """
        Decompress and parse JSON data.

        Args:
            data: Compressed JSON bytes

        Returns:
            Parsed JSON dictionary

        Raises:
            ValueError: If decompression or JSON parsing fails
        """
        logger.debug(f'Decompressing {len(data)} bytes')
        try:
            json_bytes = zlib.decompress(data)
            logger.debug(f'Decompressed to {len(json_bytes)} bytes')
            return cast(Dict[str, Any], orjson.loads(json_bytes))
        except zlib.error as e:
            logger.error(f'Zlib decompression failed: {e}')
            raise ValueError(f'Failed to decompress data: {e}') from e
        except orjson.JSONDecodeError as e:
            logger.error(f'JSON parsing failed: {e}')
            raise ValueError(f'Failed to parse JSON: {e}') from e

    def _find_json_segments(self) -> List[Tuple[int, bytes]]:  # noqa: C901
        """
        Find all JSON segments in the JPEG image.

        Returns:
            List of (position, segment_data) tuples
        """
        segments = []
        file_obj = BytesIO(self._image_data)

        # Read SOI
        if file_obj.read(2) != JPEG_SOI:
            raise ValueError('Not a valid JPEG image')

        while True:
            marker = file_obj.read(2)
            if len(marker) < 2:
                break

            if marker == JPEG_APP1:
                size_bytes = file_obj.read(2)
                if len(size_bytes) < 2:
                    break

                segment_size = struct.unpack('>H', size_bytes)[0]

                # Validate segment size is reasonable
                if segment_size < 2:
                    logger.warning(f'Invalid segment size: {segment_size}')
                    break

                # Check if we have enough data remaining
                remaining_data = len(self._image_data) - file_obj.tell()
                actual_data_size = segment_size - 2  # Subtract size field itself

                if actual_data_size > remaining_data:
                    logger.warning(f'Segment claims {actual_data_size} bytes but only {remaining_data} available')
                    # Try to read what's available
                    actual_data_size = remaining_data

                segment_data = file_obj.read(actual_data_size)

                # Only process if it's our JSON segment
                if segment_data.startswith(JSON_EXIF_HEADER):
                    segments.append((file_obj.tell() - segment_size - 2, segment_data))

            elif marker == JPEG_SOS or marker == JPEG_EOI:
                break
            elif marker[0:1] == b'\xff' and marker[1:2] >= b'\xe0':
                # Other APP markers, skip them
                size_bytes = file_obj.read(2)
                if len(size_bytes) < 2:
                    break
                try:
                    segment_size = struct.unpack('>H', size_bytes)[0]
                    if segment_size >= 2:
                        skip_size = segment_size - 2
                        # Ensure we don't seek beyond the data
                        remaining = len(self._image_data) - file_obj.tell()
                        skip_size = min(skip_size, remaining)
                        file_obj.seek(skip_size, 1)  # Skip segment data
                except struct.error:
                    logger.warning('Failed to parse segment size')
                    break
            else:
                # Unknown marker, try to continue
                file_obj.seek(-1, 1)  # Back up one byte

        logger.debug(f'Found {len(segments)} JSON segment(s)')
        return segments

    def _reassemble_chunks(self, segments: List[Tuple[int, bytes]]) -> Optional[bytes]:
        """
        Reassemble data from chunk segments.

        Args:
            segments: List of (position, segment_data) tuples

        Returns:
            Reassembled data or None if incomplete
        """
        chunks: Dict[int, bytes] = {}
        total_chunks = 0

        for _, segment_data in segments:
            try:
                header_len = len(JSON_EXIF_HEADER)

                # Ensure segment has enough data for chunk header
                if len(segment_data) < header_len + 4:
                    logger.warning(f'Segment too short: {len(segment_data)} bytes')
                    continue

                # Read chunk info
                chunk_index, chunk_total = struct.unpack('>HH', segment_data[header_len : header_len + 4])

                # Validate chunk indices
                if chunk_index >= chunk_total:
                    logger.warning(f'Invalid chunk index {chunk_index} >= total {chunk_total}')
                    continue

                if chunk_total == 0:
                    logger.warning('Invalid total chunks: 0')
                    continue

                # Extract chunk data
                chunk_data = segment_data[header_len + 4 :]
                chunks[chunk_index] = chunk_data
                total_chunks = max(total_chunks, chunk_total)

            except (struct.error, IndexError) as e:
                logger.warning(f'Failed to parse chunk header: {e}')
                continue

        # Check if we found all chunks
        if not chunks or len(chunks) != total_chunks:
            logger.warning(f'Incomplete chunks: found {len(chunks)}, expected {total_chunks}')
            return None

        # Validate all expected chunk indices are present
        for i in range(total_chunks):
            if i not in chunks:
                logger.warning(f'Missing chunk {i}')
                return None

        # Reassemble data efficiently using join
        try:
            return b''.join(chunks[i] for i in range(total_chunks))
        except KeyError as e:
            logger.warning(f'Missing chunk during reassembly: {e}')
            return None

    def extract_json(self) -> Dict[str, Any]:
        """
        Extract JSON data from the JPEG image.

        Returns:
            Extracted JSON dictionary

        Raises:
            NoMetadataError: If the image contains no JSON metadata
            MetadataDecodingError: If metadata exists but could not be decoded
        """
        # Return cached data if available
        if self._embedded_json is not None:
            logger.debug('Returning cached JSON data')
            return self._embedded_json

        segments = self._find_json_segments()
        if not segments:
            logger.debug('No JSON segments found')
            raise NoMetadataError('Image contains no JSON metadata')

        # Reassemble chunks
        full_data = self._reassemble_chunks(segments)
        if full_data is None:
            logger.warning('Failed to reassemble chunks - incomplete data')
            raise MetadataDecodingError('Failed to reassemble metadata chunks - incomplete data')

        # Decompress and parse
        try:
            self._embedded_json = self._decompress_json(full_data)
            logger.info(f'Successfully extracted JSON with {len(self._embedded_json)} top-level keys')
            return self._embedded_json
        except Exception as e:
            logger.error(f'Failed to decompress/parse JSON: {e}')
            raise MetadataDecodingError(f'Failed to decompress or parse JSON metadata: {e}') from e

    def has_embedded_json(self) -> bool:
        """
        Check if the image has embedded JSON data.

        Returns:
            True if JSON data is found
        """
        if self._embedded_json is not None:
            logger.debug('JSON data already cached - returning True')
            return True

        segments = self._find_json_segments()
        result = len(segments) > 0
        logger.debug(f'JSON detection result: {result}')
        return result

    def to_bytes(self) -> bytes:
        """
        Get the image data as bytes.

        Returns:
            Image data as bytes
        """
        return self._image_data
