# core/receiver.py
"""
THIS SOFTWARE IS PROVIDED “AS IS” AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING,
BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
PARTICULAR PURPOSE ARE DISCLAIMED.

Acts as the application-layer processor for image-specific payloads.

This module is responsible for interpreting fully reassembled byte payloads,
parsing them according to the defined image format (SAR or MFD), extracting
metadata and pixel data, and invoking the appropriate application callbacks
for further processing (e.g., display, recording).
"""

# Standard library imports
import logging
import math
from typing import Optional, Dict, Any, Tuple, Collection, Callable 

# Third-party imports
import numpy as np
import cv2

# Local application imports
from controlpanel import config
from controlpanel.utils.image_processing import normalize_image
from controlpanel.utils.utils import format_ctypes_structure, put_queue
from controlpanel.core.sfp_structures import ImageLeaderData
from controlpanel.core.image_recorder import ImageRecorder


class ImagePayloadProcessor:
    """
    Parses reassembled SFP payloads that conform to the image data format.
    """

    def __init__(
        self,
        app,
        set_new_sar_image_callback: Callable,
        set_new_mfd_indices_image_callback: Callable,
        image_recorder: Optional[ImageRecorder] = None,
    ):
        """
        Initializes the Image Payload Processor.

        Args:
            app (ControlPanelApp): The main application instance for context (state, queues).
            set_new_sar_image_callback (Callable): Callback for new SAR images.
            set_new_mfd_indices_image_callback (Callable): Callback for new MFD images.
            image_recorder (Optional[ImageRecorder]): Instance for saving images.
        """
        self._log_prefix = "[ImageProcessor]"
        logging.debug(f"{self._log_prefix} Initializing...")
        self.app = app
        self.set_new_sar_image_callback = set_new_sar_image_callback
        self.set_new_mfd_image_callback = set_new_mfd_indices_image_callback
        self.image_recorder = image_recorder
        logging.info(f"{self._log_prefix} Initialization complete.")

    def process_sar_payload(self, payload: bytearray):
        """
        Processes a complete SAR image payload from a bytearray.

        Args:
            payload (bytearray): The fully reassembled payload for a SAR image.
        """
        log_prefix = f"{self._log_prefix} SAR"
        if self.app.state.shutting_down:
            return

        try:
            image_leader = ImageLeaderData.from_buffer(payload)
            fcounter = image_leader.HEADER_DATA.FCOUNTER
            log_prefix = f"{self._log_prefix} SAR(FCNT={fcounter})"

            reassembly_result = self._reassemble_sar_image(image_leader, payload, log_prefix)

            if self.app.state.shutting_down:
                return

            if reassembly_result:
                raw_sar_data, normalized_sar_uint8, geo_info_radians = reassembly_result

                # --- Callback for Display ---
                logging.debug(f"{log_prefix} Invoking app callback for display...")
                self.set_new_sar_image_callback(normalized_sar_uint8, geo_info_radians)

                # --- Callback for Recording ---
                if self.image_recorder:
                    logging.debug(f"{log_prefix} Invoking image recorder...")
                    self.image_recorder.record_sar_image(raw_sar_data, geo_info_radians)

                # --- Metadata Formatting and Queueing ---
                if self.app.state.display_sar_metadata:
                    metadata_str = format_ctypes_structure(image_leader)
                    put_queue(
                        self.app.tkinter_queue,
                        ("SAR_METADATA_UPDATE", metadata_str),
                        "tkinter",
                        self.app,
                    )
            else:
                logging.error(f"{log_prefix} SAR image reassembly failed.")
                self.app.state.increment_incomplete_rx_count("sar")

        except Exception:
            logging.exception(f"{log_prefix} Unexpected error processing SAR payload.")
            self.app.state.increment_incomplete_rx_count("sar")

    def process_mfd_payload(self, payload: bytearray):
        """
        Processes a complete MFD image payload from a bytearray.

        Args:
            payload (bytearray): The fully reassembled payload for an MFD image.
        """
        log_prefix = f"{self._log_prefix} MFD"
        if self.app.state.shutting_down:
            return

        try:
            image_leader = ImageLeaderData.from_buffer(payload)
            fcounter = image_leader.HEADER_DATA.FCOUNTER
            log_prefix = f"{self._log_prefix} MFD(FCNT={fcounter})"

            mfd_indices = self._reassemble_mfd_image(image_leader, payload, log_prefix)

            if self.app.state.shutting_down:
                return

            if mfd_indices is not None:
                logging.debug(f"{log_prefix} MFD indices reassembled. Invoking app callback...")
                self.set_new_mfd_image_callback(mfd_indices)
            else:
                logging.error(f"{log_prefix} MFD image reassembly failed.")
                self.app.state.increment_incomplete_rx_count("mfd")

        except Exception:
            logging.exception(f"{log_prefix} Unexpected error processing MFD payload.")
            self.app.state.increment_incomplete_rx_count("mfd")

    def _calculate_pixel_data_offset(self, image_leader: ImageLeaderData) -> int:
        """Calculates the expected byte offset to the start of the pixel data."""
        # This is based on the static structure sizes defined in sfp_structures
        offset = (
            image_leader.HEADER_TAG.size() + image_leader.HEADER_DATA.size()
            + image_leader.GEO_TAG.size() + image_leader.GEO_DATA.size()
            + image_leader.RESERVED_TAG.size() + image_leader.get_reserved_data_size()
            + image_leader.CM_TAG.size() + image_leader.get_colour_map_size()
            + image_leader.PIXEL_TAG.size()
        )
        return offset

    def _reassemble_sar_image(
        self, image_leader: ImageLeaderData, image_data: bytearray, log_prefix: str
    ) -> Optional[Tuple[np.ndarray, np.ndarray, Dict[str, Any]]]:
        """
        Extracts SAR metadata and pixel data from a complete buffer.

        Returns:
            A tuple of (raw_image, normalized_image, geo_info) or None on error.
        """
        try:
            hdr_d = image_leader.HEADER_DATA
            dx, dy, bpp = int(hdr_d.DX), int(hdr_d.DY), int(hdr_d.BPP)
            stride_pixels, pal_type = int(hdr_d.STRIDE), int(hdr_d.PALTYPE)

            if dx <= 0 or dy <= 0 or bpp not in [1, 2] or stride_pixels < dx or pal_type != 0:
                logging.error(f"{log_prefix} Invalid SAR metadata in header. Cannot parse.")
                return None

            pixel_dtype = np.uint16 if bpp == 2 else np.uint8
            pixel_data_offset = self._calculate_pixel_data_offset(image_leader)
            
            min_required_bytes = dy * dx * bpp
            if (pixel_data_offset + min_required_bytes) > len(image_data):
                logging.error(f"{log_prefix} Incomplete data buffer for SAR image.")
                return None

            raw_image_data = np.ndarray(
                shape=(dy, dx),
                dtype=pixel_dtype,
                buffer=image_data,
                offset=pixel_data_offset,
                strides=(stride_pixels * bpp, bpp),
            )

            if self.app.state.shutting_down: return None

            normalized_image_uint8 = normalize_image(raw_image_data, target_type=np.uint8)
            if normalized_image_uint8 is None:
                logging.error(f"{log_prefix} SAR normalization to uint8 failed.")
                return None

            # Extract and validate Geo Info
            geo_d = image_leader.GEO_DATA
            lat_rad, lon_rad, orient_rad = float(geo_d.LATITUDE), float(geo_d.LONGITUDE), float(geo_d.ORIENTATION)
            
            is_valid = (
                float(geo_d.SCALE_X) > 0 and float(geo_d.SCALE_Y) > 0 and
                -math.pi / 2 <= lat_rad <= math.pi / 2 and
                -math.pi <= lon_rad <= math.pi and math.isfinite(orient_rad)
            )

            geo_info = {
                "lat": lat_rad,
                "lon": lon_rad,
                "orientation": orient_rad,
                "ref_x": int(geo_d.REF_X),
                "ref_y": int(geo_d.REF_Y),
                "scale_x": float(geo_d.SCALE_X),
                "scale_y": float(geo_d.SCALE_Y),
                "width_px": dx,
                "height_px": dy,
                "valid": is_valid,
            }

            if not is_valid:
                logging.warning(f"{log_prefix} GeoData values are invalid. Marking as not valid.")

            return raw_image_data.copy(), normalized_image_uint8, geo_info

        except Exception:
            logging.exception(f"{log_prefix} Unexpected error parsing SAR image from buffer.")
            return None

    def _reassemble_mfd_image(
        self, image_leader: ImageLeaderData, image_data: bytearray, log_prefix: str
    ) -> Optional[np.ndarray]:
        """
        Extracts MFD pixel indices (uint8) from a complete buffer.

        Returns:
            A NumPy array of the indices, or None on error.
        """
        try:
            hdr_d = image_leader.HEADER_DATA
            dx, dy, bpp = int(hdr_d.DX), int(hdr_d.DY), int(hdr_d.BPP)
            stride_pixels, pal_type = int(hdr_d.STRIDE), int(hdr_d.PALTYPE)

            if dx <= 0 or dy <= 0 or bpp != 1 or stride_pixels < dx or pal_type != 1:
                logging.error(f"{log_prefix} Invalid MFD metadata in header. Cannot parse.")
                return None

            pixel_data_offset = self._calculate_pixel_data_offset(image_leader)
            
            min_required_bytes = dy * dx * bpp
            if (pixel_data_offset + min_required_bytes) > len(image_data):
                logging.error(f"{log_prefix} Incomplete data buffer for MFD image.")
                return None

            mfd_index_view = np.ndarray(
                shape=(dy, dx),
                dtype=np.uint8,
                buffer=image_data,
                offset=pixel_data_offset,
                strides=(stride_pixels * bpp, bpp),
            )
            
            return mfd_index_view.copy()

        except Exception:
            logging.exception(f"{log_prefix} Unexpected error parsing MFD image from buffer.")
            return None