"""
Hypfer Image Handler Class.
It returns the PIL PNG image frame relative to the Map Data extrapolated from the vacuum json.
It also returns calibration, rooms data to the card and other images information to the camera.
Version: 0.1.13
"""

from __future__ import annotations

import asyncio

import numpy as np
from mvcrender.autocrop import AutoCrop
from PIL import Image

from .config.async_utils import AsyncPIL
from .config.drawable_elements import DrawableElement
from .config.shared import CameraShared
from .config.types import (
    LOGGER,
    CalibrationPoints,
    Colors,
    JsonType,
    RoomsProperties,
    RoomStore,
)
from .config.utils import (
    BaseHandler,
    initialize_drawing_config,
)
from .const import COLORS
from .hypfer_draw import ImageDraw as ImDraw
from .map_data import ImageData
from .rooms_handler import RoomsHandler


class HypferMapImageHandler(BaseHandler, AutoCrop):
    """Map Image Handler Class.
    This class is used to handle the image data and the drawing of the map."""

    def __init__(self, shared_data: CameraShared):
        """Initialize the Map Image Handler."""
        BaseHandler.__init__(self)
        self.shared = shared_data  # camera shared data
        AutoCrop.__init__(self, self)
        self.calibration_data = None  # camera shared data.
        self.data = ImageData  # imported Image Data Module.
        # Initialize drawing configuration using the shared utility function
        self.drawing_config, self.draw = initialize_drawing_config(self)

        self.go_to = None  # vacuum go to data
        self.img_hash = None  # hash of the image calculated to check differences.
        self.img_base_layer = None  # numpy array store the map base layer.
        self.img_work_layer = (
            None  # persistent working buffer to avoid per-frame allocations
        )
        self.active_zones = []  # vacuum active zones.
        self.imd = ImDraw(self)  # Image Draw class.
        self.color_grey = (128, 128, 128, 255)
        self.file_name = self.shared.file_name  # file name of the vacuum.
        self.rooms_handler = RoomsHandler(
            self.file_name, self.drawing_config
        )  # Room data handler

    @staticmethod
    def get_corners(x_max, x_min, y_max, y_min):
        """Get the corners of the room."""
        return [(x_min, y_min), (x_max, y_min), (x_max, y_max), (x_min, y_max)]

    async def async_extract_room_properties(self, json_data) -> RoomsProperties:
        """Extract room properties from the JSON data."""
        room_properties = await self.rooms_handler.async_extract_room_properties(
            json_data
        )
        if room_properties:
            _ = RoomStore(self.file_name, room_properties)
            # Convert room_properties to the format expected by async_get_robot_in_room
            self.rooms_pos = []
            for room_id, room_data in room_properties.items():
                self.rooms_pos.append(
                    {
                        "id": room_id,
                        "name": room_data["name"],
                        "outline": room_data["outline"],
                    }
                )
        else:
            self.rooms_pos = None
        return room_properties

    def _identify_disabled_rooms(self) -> set:
        """Identify which rooms are disabled in the drawing configuration."""
        disabled_rooms = set()
        room_id = 0
        for layer_type, _ in self.json_data.layers.items():
            if layer_type == "segment":
                current_room_id = room_id + 1
                if 1 <= current_room_id <= 15:
                    room_element = getattr(
                        DrawableElement, f"ROOM_{current_room_id}", None
                    )
                    if room_element and not self.drawing_config.is_enabled(
                        room_element
                    ):
                        disabled_rooms.add(room_id)
                room_id = (room_id + 1) % 16
        return disabled_rooms

    async def _draw_layer_if_enabled(
        self,
        img_np_array,
        layer_type,
        compressed_pixels_list,
        colors,
        pixel_size,
        disabled_rooms,
        room_id,
    ):
        """Draw a layer if it's enabled in the drawing configuration."""
        is_room_layer = layer_type == "segment"

        if is_room_layer:
            current_room_id = room_id + 1
            if 1 <= current_room_id <= 15:
                room_element = getattr(DrawableElement, f"ROOM_{current_room_id}", None)
                if not self.drawing_config.is_enabled(room_element):
                    return room_id + 1, img_np_array  # Skip this room

        is_wall_layer = layer_type == "wall"
        if is_wall_layer and not self.drawing_config.is_enabled(DrawableElement.WALL):
            return room_id, img_np_array  # Skip walls

        # Draw the layer (floor layers are always drawn when present)
        room_id, img_np_array = await self.imd.async_draw_base_layer(
            img_np_array,
            compressed_pixels_list,
            layer_type,
            colors["wall"],
            colors["zone_clean"],
            pixel_size,
            disabled_rooms if layer_type == "wall" else None,
        )
        return room_id, img_np_array

    async def _draw_base_layers(self, img_np_array, colors, pixel_size):
        """Draw all base layers (rooms, walls, floors)."""
        disabled_rooms = self._identify_disabled_rooms()
        room_id = 0

        for layer_type, compressed_pixels_list in self.json_data.layers.items():
            room_id, img_np_array = await self._draw_layer_if_enabled(
                img_np_array,
                layer_type,
                compressed_pixels_list,
                colors,
                pixel_size,
                disabled_rooms,
                room_id,
            )

        return img_np_array, room_id

    async def _draw_additional_elements(
        self, img_np_array, m_json, entity_dict, colors
    ):
        """Draw additional elements like walls, charger, and obstacles."""
        if self.drawing_config.is_enabled(DrawableElement.VIRTUAL_WALL):
            img_np_array = await self.imd.async_draw_virtual_walls(
                m_json, img_np_array, colors["no_go"]
            )

        if self.drawing_config.is_enabled(DrawableElement.CHARGER):
            img_np_array = await self.imd.async_draw_charger(
                img_np_array, entity_dict, colors["charger"]
            )

        if self.drawing_config.is_enabled(DrawableElement.OBSTACLE):
            self.shared.obstacles_pos = self.data.get_obstacles(entity_dict)
            if self.shared.obstacles_pos:
                img_np_array = await self.imd.async_draw_obstacle(
                    img_np_array, self.shared.obstacles_pos, colors["no_go"]
                )

        return img_np_array

    async def _setup_room_and_robot_data(
        self, room_id, robot_position, robot_position_angle
    ):
        """Setup room properties and robot position data."""
        if (room_id > 0) and not self.room_propriety:
            self.room_propriety = await self.async_extract_room_properties(
                self.json_data.json_data
            )

        if not self.rooms_pos and not self.room_propriety:
            self.room_propriety = await self.async_extract_room_properties(
                self.json_data.json_data
            )

        if self.rooms_pos and robot_position and robot_position_angle:
            self.robot_pos = await self.imd.async_get_robot_in_room(
                robot_x=(robot_position[0]),
                robot_y=(robot_position[1]),
                angle=robot_position_angle,
            )

    async def _prepare_data_tasks(self, m_json, entity_dict):
        """Prepare and execute data extraction tasks in parallel."""
        data_tasks = []

        if self.drawing_config.is_enabled(DrawableElement.RESTRICTED_AREA):
            data_tasks.append(self._prepare_zone_data(m_json))

        if self.drawing_config.is_enabled(DrawableElement.GO_TO_TARGET):
            data_tasks.append(self._prepare_goto_data(entity_dict))

        path_enabled = self.drawing_config.is_enabled(DrawableElement.PATH)
        LOGGER.debug("%s: PATH element enabled: %s", self.file_name, path_enabled)
        if path_enabled:
            LOGGER.debug("%s: Drawing path", self.file_name)
            data_tasks.append(self._prepare_path_data(m_json))

        if data_tasks:
            await asyncio.gather(*data_tasks)

        return path_enabled

    async def _draw_dynamic_elements(
        self, img_np_array, m_json, entity_dict, colors, path_enabled
    ):
        """Draw dynamic elements like zones, paths, and go-to targets."""
        if self.drawing_config.is_enabled(DrawableElement.RESTRICTED_AREA):
            img_np_array = await self.imd.async_draw_zones(
                m_json,
                img_np_array,
                colors["zone_clean"],
                colors["no_go"],
                colors.get("carpet"),
            )

        if self.drawing_config.is_enabled(DrawableElement.GO_TO_TARGET):
            img_np_array = await self.imd.draw_go_to_flag(
                img_np_array, entity_dict, colors["go_to"]
            )

        if path_enabled:
            img_np_array = await self.imd.async_draw_paths(
                img_np_array,
                m_json,
                colors["move"],
                self.color_grey,
                colors["mop_move"],
            )
        else:
            LOGGER.debug("%s: Skipping path drawing", self.file_name)

        return img_np_array

    async def _draw_robot_if_enabled(
        self, img_np_array, robot_pos, robot_position, robot_position_angle, colors
    ):
        """Draw the robot on the map if enabled."""
        if self.shared.vacuum_state == "docked":
            robot_position_angle -= 180

        if robot_pos and self.drawing_config.is_enabled(DrawableElement.ROBOT):
            robot_color = self.drawing_config.get_property(
                DrawableElement.ROBOT, "color", colors["robot"]
            )
            img_np_array = await self.draw.robot(
                layers=img_np_array,
                x=robot_position[0],
                y=robot_position[1],
                angle=robot_position_angle,
                fill=robot_color,
                radius=self.imd.robot_size,
                robot_state=self.shared.vacuum_state,
            )

        return img_np_array

    # noinspection PyUnresolvedReferences,PyUnboundLocalVariable
    async def async_get_image_from_json(
        self,
        m_json: JsonType | None,
    ) -> Image.Image | None:
        """Get the image from the JSON data.
        It uses the ImageDraw class to draw some of the elements of the image.
        The robot itself will be drawn in this function as per some of the values are needed for other tasks.
        @param m_json: The JSON data to use to draw the image.
        @return Image.Image: PIL Image.
        """
        # Initialize the colors.
        colors: Colors = {
            name: self.shared.user_colors[idx] for idx, name in enumerate(COLORS)
        }
        # Initialize img_np_array to None before conditional
        img_np_array = None
        # Check if the JSON data is not None else process the image.
        try:
            if m_json is not None:
                # Get the image size from the JSON data
                self.img_size = self.json_data.image_size
                # Get the JSON ID from the JSON data.
                self.json_id = self.json_data.json_id
                # Check entity data.
                entity_dict = self.json_data.entity_dict
                # Update the Robot position.
                (
                    robot_pos,
                    robot_position,
                    robot_position_angle,
                ) = await self.imd.async_get_robot_position(entity_dict)

                # Get the pixels size and layers from the JSON data
                pixel_size = int(m_json["pixelSize"])
                self.active_zones = self.json_data.active_zones

                new_frame_hash = await self.calculate_array_hash(
                    self.json_data.layers, self.active_zones
                )
                if self.frame_number == 0:
                    self.img_hash = new_frame_hash
                    # Create empty image
                    img_np_array = await self.draw.create_empty_image(
                        self.img_size["x"], self.img_size["y"], colors["background"]
                    )
                    room_id = 0

                    if self.drawing_config.is_enabled(DrawableElement.FLOOR):
                        img_np_array, room_id = await self._draw_base_layers(
                            img_np_array, colors, pixel_size
                        )

                    img_np_array = await self._draw_additional_elements(
                        img_np_array, m_json, entity_dict, colors
                    )

                    await self._setup_room_and_robot_data(
                        room_id, robot_position, robot_position_angle
                    )

                    LOGGER.debug("%s: Completed base Layers", self.file_name)
                    # Copy the new array in base layer.
                    # Delete old base layer before creating new one to free memory
                    if self.img_base_layer is not None:
                        del self.img_base_layer
                    self.img_base_layer = await self.async_copy_array(img_np_array)
                    # Delete source array after copying to free memory
                    del img_np_array

                self.shared.frame_number = self.frame_number
                self.frame_number += 1
                if (self.frame_number >= self.max_frames) or (
                    new_frame_hash != self.img_hash
                ):
                    self.frame_number = 0
                # Ensure persistent working buffer exists and matches base (allocate only when needed)
                if (
                    self.img_work_layer is None
                    or self.img_work_layer.shape != self.img_base_layer.shape
                    or self.img_work_layer.dtype != self.img_base_layer.dtype
                ):
                    # Delete old buffer before creating new one to free memory
                    if self.img_work_layer is not None:
                        del self.img_work_layer
                    self.img_work_layer = np.empty_like(self.img_base_layer)

                # Copy the base layer into the persistent working buffer (no new allocation per frame)
                np.copyto(self.img_work_layer, self.img_base_layer)
                img_np_array = self.img_work_layer

                # Prepare and execute data extraction tasks
                path_enabled = await self._prepare_data_tasks(m_json, entity_dict)

                # Draw dynamic elements
                img_np_array = await self._draw_dynamic_elements(
                    img_np_array, m_json, entity_dict, colors, path_enabled
                )

                # Draw robot
                img_np_array = await self._draw_robot_if_enabled(
                    img_np_array,
                    robot_pos,
                    robot_position,
                    robot_position_angle,
                    colors,
                )

                # Synchronize zooming state from ImageDraw to handler before auto-crop
                self.zooming = self.imd.img_h.zooming

                # Resize the image - auto_trim_and_zoom_image creates a new array
                # Note: img_np_array currently references self.img_work_layer
                img_np_array = self.auto_trim_and_zoom_image(
                    img_np_array,
                    colors["background"],
                    int(self.shared.margins),
                    int(self.shared.image_rotate),
                    self.zooming,
                )
                # After cropping, we no longer need the full-size working layer
                # Clear it to free memory (will be recreated next frame if needed)
                self.img_work_layer = None
                # Update shared trims with calculated crop values
                self.update_trims()
            # If the image is None return None and log the error.
            if img_np_array is None:
                LOGGER.warning("%s: Image array is None.", self.file_name)
                return None

            # Handle resizing if needed, then return based on format preference
            if self.check_zoom_and_aspect_ratio():
                # Convert to PIL for resizing
                pil_img = await AsyncPIL.async_fromarray(img_np_array, mode="RGBA")
                del img_np_array
                resize_params = self.prepare_resize_params(pil_img)
                resized_image = await self.async_resize_images(resize_params)

                # Return PIL Image
                return resized_image

            # Return PIL Image (convert from NumPy)
            pil_img = await AsyncPIL.async_fromarray(img_np_array, mode="RGBA")
            del img_np_array
            return pil_img
        except (RuntimeError, RuntimeWarning) as e:
            LOGGER.debug(
                "%s: Error %s during image creation.",
                self.file_name,
                str(e),
                exc_info=True,
            )
            return None

    async def async_get_rooms_attributes(self) -> RoomsProperties:
        """Get the rooms attributes from the JSON data.
        :return: The rooms attribute's."""
        if self.room_propriety:
            return self.room_propriety
        if self.json_data:
            self.room_propriety = await self.async_extract_room_properties(
                self.json_data.json_data
            )
        return self.room_propriety

    def get_calibration_data(self, rotation_angle: int = 0) -> CalibrationPoints:
        """Get the calibration data from the JSON data.
        this will create the attribute calibration points."""
        calibration_data = []

        # Define the map points (fixed)
        map_points = self.get_map_points()
        # Calculate the calibration points in the vacuum coordinate system
        vacuum_points = self.get_vacuum_points(rotation_angle)

        # Create the calibration data for each point
        for vacuum_point, map_point in zip(vacuum_points, map_points):
            calibration_point = {"vacuum": vacuum_point, "map": map_point}
            calibration_data.append(calibration_point)
        del vacuum_points, map_points, calibration_point, rotation_angle  # free memory.
        return calibration_data

    async def _prepare_zone_data(self, m_json):
        """Prepare zone data for parallel processing."""
        await asyncio.sleep(0)  # Yield control
        try:
            return self.data.find_zone_entities(m_json)
        except (ValueError, KeyError):
            return None

    @staticmethod
    async def _prepare_goto_data(entity_dict):
        """Prepare go-to flag data for parallel processing."""
        await asyncio.sleep(0)  # Yield control
        # Extract go-to target data from entity_dict
        return entity_dict.get("go_to_target", None)

    async def _prepare_path_data(self, m_json):
        """Prepare path data for parallel processing."""
        await asyncio.sleep(0)  # Yield control
        try:
            return self.data.find_paths_entities(m_json)
        except (ValueError, KeyError):
            return None
