"""Service layer for handling image loading, processing, and encoding/decoding."""

import base64
import json
import logging
from functools import lru_cache
from pathlib import Path
from typing import NotRequired, TypedDict

import cv2
import numpy as np
import numpy.typing as npt

from src.config import settings
from src.models.image import (
    Annotations,
    BBoxData,
    ImageData,
    ImageSet,
    KeypointData,
)

logger = logging.getLogger(__name__)


class LabelMeShape(TypedDict):
    """Type for LabelMe shape data."""

    label: str
    points: list[list[float]]
    shape_type: str
    mask: NotRequired[None]


class LabelMeData(TypedDict):
    """Type for LabelMe JSON data."""

    version: str
    shapes: list[LabelMeShape]
    imagePath: str
    imageHeight: int
    imageWidth: int


class ImageService:
    """Service for handling image operations."""

    @staticmethod
    def load_image(image_path: str) -> np.ndarray:
        """Load an image from path and convert to RGB."""
        return cv2.imread(str(image_path), cv2.IMREAD_COLOR_RGB)

    @staticmethod
    def encode_image_to_base64(image: np.ndarray, quality: int = 70) -> str:
        """Encode numpy array to base64 string.

        Args:
            image: Numpy array representing the image
            quality: JPEG quality (1-100, higher is better quality but larger file)

        Returns:
            str: Base64 encoded string with data URL prefix
        """
        try:
            # Convert to BGR for OpenCV
            if len(image.shape) == 3 and image.shape[2] == 3:
                # RGB to BGR for OpenCV
                image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
            elif len(image.shape) == 2 or (len(image.shape) == 3 and image.shape[2] == 1):
                # Grayscale to BGR
                image_bgr = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
            elif len(image.shape) == 3 and image.shape[2] == 4:
                # RGBA to BGR
                image_bgr = cv2.cvtColor(image, cv2.COLOR_RGBA2BGR)
            else:
                image_bgr = image  # Use as is

            # Use OpenCV to encode the image to JPEG with specified quality
            encode_param = [int(cv2.IMWRITE_WEBP_QUALITY), quality]
            success, buffer = cv2.imencode(".webp", image_bgr, encode_param)

            if not success:
                raise ValueError("Failed to encode image")

            # Encode to base64
            base64_string = base64.b64encode(buffer).decode("utf-8")
            return f"data:image/webp;base64,{base64_string}"
        except Exception as e:
            logger.exception("Failed to encode image to base64")
            raise ValueError(f"Failed to encode image: {e!s}") from e

    @staticmethod
    def decode_base64_to_image(base64_string: str) -> np.ndarray:
        """Decode base64 string to numpy array."""
        try:
            # Remove data URL prefix if present
            if "base64," in base64_string:
                base64_string = base64_string.split("base64,")[1]

            # Decode base64 to bytes
            image_bytes = base64.b64decode(base64_string)

            # Convert bytes to numpy array
            np_arr = np.frombuffer(image_bytes, np.uint8)

            # Decode image using OpenCV
            image_bgr = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)

            if image_bgr is None:
                raise ValueError("Failed to decode image")

            # Convert from BGR to RGB (OpenCV uses BGR by default)
            return cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)

        except Exception as e:
            logger.exception("Failed to decode base64 image")
            raise ValueError(f"Invalid image data: {e!s}") from e

    def load_annotations(self, annotation_path: Path | None) -> Annotations:
        """Load annotations from JSON file in LabelMe format.

        Args:
            annotation_path: Path to the annotation file or None

        Returns:
            Annotations: Annotations object containing bboxes, keypoints and mask
        """
        if not annotation_path or not annotation_path.exists():
            # Return empty annotations if no annotation file is specified
            return Annotations(bboxes=[], keypoints=[], mask=None)

        with annotation_path.open("r") as f:
            data: LabelMeData = json.load(f)

        image_width = data["imageWidth"]
        image_height = data["imageHeight"]

        # Parse bounding boxes (rectangles)
        bboxes: list[BBoxData] = []
        for shape in data["shapes"]:
            if shape["shape_type"] == "rectangle":
                points = shape["points"]
                x_min = min(points[0][0], points[1][0])
                y_min = min(points[0][1], points[1][1])
                x_max = max(points[0][0], points[1][0])
                y_max = max(points[0][1], points[1][1])

                bboxes.append(
                    BBoxData(
                        label=shape["label"],
                        x_min=x_min,
                        y_min=y_min,
                        x_max=x_max,
                        y_max=y_max,
                    ),
                )

        # Parse keypoints
        keypoints: list[KeypointData] = []
        for shape in data["shapes"]:
            if shape["shape_type"] == "point":
                point = shape["points"][0]
                keypoints.append(
                    KeypointData(
                        label=shape["label"],
                        x=point[0],
                        y=point[1],
                    ),
                )

        # Parse mask (polygon)
        mask: npt.NDArray[np.uint8] | None = None
        for shape in data["shapes"]:
            if shape["shape_type"] == "polygon":
                # Create binary mask from polygon
                mask = np.zeros((image_height, image_width), dtype=np.uint8)

                # Convert polygon points to numpy arrays
                points = np.array(shape["points"], dtype=np.int32)

                # Fill polygon area with ones
                cv2.fillPoly(mask, [points], 255)

        return Annotations(
            bboxes=bboxes,
            keypoints=keypoints,
            mask=mask,
            image_width=image_width,
            image_height=image_height,
        )

    @lru_cache(maxsize=128)  # noqa: B019
    def _create_image_data_from_set(self, image_set: ImageSet) -> ImageData:
        """Load image and annotations from an ImageSet and return ImageData. (Cached)."""
        logger.info("Loading/Creating ImageData for ImageSet: %s", image_set.name)
        image_path = settings.get_image_path(image_set)
        annotation_path = settings.get_annotation_path(image_set)

        image = self.load_image(str(image_path))
        encoded_image = self.encode_image_to_base64(image)

        annotations = self.load_annotations(annotation_path)

        encoded_mask = None
        if annotations.mask is not None:
            encoded_mask = self.encode_image_to_base64(annotations.mask)

        return ImageData(
            image=encoded_image,
            mask=encoded_mask,
            bboxes=annotations.bboxes,
            keypoints=annotations.keypoints,
        )

    def get_original_image_data(self, transform_name: str) -> ImageData:
        """Create an ImageData object for the primary image of a transform."""
        # Get the primary image set
        primary_image_set = settings.get_image_set(transform_name)

        # Create ImageData using the helper
        return self._create_image_data_from_set(primary_image_set)

    def get_image_set(self, transform_name: str) -> ImageSet:
        """Get the appropriate ImageSet for a transform."""
        return settings.get_image_set(transform_name)

    def get_reference_image_data(self, transform_name: str | None = None) -> list[ImageData]:
        """Load reference/secondary images and annotations for a transform."""
        reference_image_sets = settings.get_reference_image_sets(transform_name)

        return [self._create_image_data_from_set(image_set) for image_set in reference_image_sets]
