"""Service layer for handling transform discovery, parameter processing, and execution."""

import base64
import inspect
import logging
from collections.abc import Callable, Sequence
from typing import Any, Literal

import albumentations as A
import cv2
import numpy as np
from numpy.typing import NDArray

from src.config import settings
from src.exceptions import TransformNotFoundError
from src.models.image import Annotations, ImageData
from src.models.transform import TransformExecutionConfig, TransformInfo
from src.services.image_service import ImageService
from src.transforms.constants import CV2ConstantHandler
from src.transforms.fields import FieldExtractor
from src.transforms.info import get_transform_info

logger = logging.getLogger(__name__)

IGNORED_CLASSES = {
    "Lambda",
    "BasicTransform",
    "DualTransform",
    "ImageOnlyTransform",
    "Transform3D",
    "TextImage",
    "PiecewiseAffine",
    "OverlayElements",
}

image_service = ImageService()


class TransformService:
    """Service for handling transform operations."""

    def __init__(self) -> None:
        self.field_extractor = FieldExtractor()
        self.cv2_handler = CV2ConstantHandler()
        self.image_service = ImageService()

    # --- Helper Methods for Param Processing --- #

    def _convert_cv2_constants(self, params: dict[str, Any] | None) -> dict[str, Any]:
        """Converts string CV2 constants in params to their integer values."""
        params = params or {}
        cv2_param_mappings: dict[str, Callable[[str], int]] = {
            "interpolation": self.cv2_handler.get_interpolation_constant,
            "mask_interpolation": self.cv2_handler.get_interpolation_constant,
            "border_mode": self.cv2_handler.get_border_constant,
        }
        for param_name, converter in cv2_param_mappings.items():
            if param_name in params and isinstance(params[param_name], str):
                params[param_name] = converter(params[param_name])
        return params

    # --- Main Parameter Update Method --- #

    def _update_transform_params(
        self,
        params: dict[str, Any] | None,
        transform_class: type[A.BasicTransform],
    ) -> dict[str, Any]:
        """Updates transform parameters by handling CV2 constants and adding metadata_key."""
        params = params or {}
        params = self._convert_cv2_constants(params)

        # Add metadata_key for mixing transforms defined in settings
        # Exclude Mosaic for now as its input mechanism might differ
        if transform_class in settings.MIXING_TRANSFORMS:
            params["metadata_key"] = "metadata"

        return params

    def get_transform_class(self, transform_name: str) -> type[Any]:
        """Get transform class by name."""
        transform_class = getattr(A, transform_name, None)
        if transform_class is None:
            raise ValueError(f"Transform {transform_name} not found")
        return transform_class

    def get_transform_info(self, transform_name: str) -> TransformInfo:
        """Get detailed information about a specific transform."""
        transform_class = getattr(A, transform_name, None)

        if not transform_class:
            raise TransformNotFoundError(f"Transform '{transform_name}' not found")

        return get_transform_info(transform_class)

    def get_all_transforms_info(self) -> dict[Literal["image_only", "dual", "transforms_3d"], list[TransformInfo]]:
        """Get information about all available transforms, grouped by type."""
        image_only_transforms: list[TransformInfo] = []
        dual_transforms: list[TransformInfo] = []
        transforms_3d: list[TransformInfo] = []

        # Collect all transforms in a single pass
        for name, cls in inspect.getmembers(A, predicate=inspect.isclass):
            if not inspect.isclass(cls) or name in IGNORED_CLASSES:
                continue

            if issubclass(cls, A.Transform3D) and cls is not A.Transform3D:
                transforms_3d.append(
                    TransformInfo(
                        name=name,
                        transform_type="transforms_3d",
                        targets=[t.value.lower() for t in cls._targets],
                    ),
                )
            elif issubclass(cls, A.BasicTransform) and cls is not A.BasicTransform:
                if issubclass(cls, A.ImageOnlyTransform):
                    image_only_transforms.append(
                        TransformInfo(
                            name=name,
                            transform_type="image_only",
                            targets=["image", "volume"],
                        ),
                    )
                elif issubclass(cls, A.DualTransform):
                    dual_transforms.append(
                        TransformInfo(
                            name=name,
                            transform_type="dual",
                            targets=[t.value.lower() for t in cls._targets],
                        ),
                    )

        return {
            "image_only": image_only_transforms,
            "dual": dual_transforms,
            "transforms_3d": transforms_3d,
        }

    # --- Helper Methods for execute_transform --- #

    def _decode_user_refs(self, user_provided_refs: Any) -> list[NDArray[np.uint8]]:  # noqa: ANN401
        """Validates and decodes user-provided reference images."""
        if not isinstance(user_provided_refs, list):
            raise TypeError("Provided reference_images must be a list")

        if not user_provided_refs:
            return []  # Empty list provided

        first_item = user_provided_refs[0]
        if isinstance(first_item, str):
            # Decode list of base64
            try:
                return [self.image_service.decode_base64_to_image(ref) for ref in user_provided_refs]
            except Exception as e:
                logger.exception("Failed to decode provided reference images")
                raise ValueError("Invalid base64 format for provided reference_images") from e
        elif isinstance(first_item, np.ndarray):
            # Already numpy arrays, assume correct type
            return user_provided_refs
        else:
            msg = "Invalid format for provided reference_images list (must be list of base64 strings or numpy arrays)"
            logger.error(msg)
            raise TypeError(msg)

    def _decode_default_refs(self, transform_name: str) -> list[NDArray[np.uint8]]:
        """Loads and decodes default reference images."""
        logger.info("%s needs reference images and none provided, loading default(s).", transform_name)
        try:
            reference_data_list = self.image_service.get_reference_image_data(transform_name)
            if reference_data_list:
                return [self.image_service.decode_base64_to_image(ref_data.image) for ref_data in reference_data_list]
            logger.warning("get_reference_image_data returned empty list for %s", transform_name)
            return []  # Return empty list if no defaults found
        except Exception as e:
            logger.exception("Failed to load/decode default reference images for %s", transform_name)
            raise ValueError(f"Failed to prepare default reference images for {transform_name}") from e

    def _prepare_metadata(
        self,
        original_request_params: dict[str, Any] | None,
        transform_class: type[A.BasicTransform],
    ) -> dict[str, Any] | None:
        """Prepares the metadata dictionary with decoded reference images/data if needed."""
        if transform_class not in settings.MIXING_TRANSFORMS:
            return None

        original_request_params = original_request_params or {}
        user_provided_refs = original_request_params.get("reference_images")

        # --- Special Handling for Mosaic --- #
        if transform_class is A.Mosaic:
            # Currently only supports loading configured defaults for Mosaic
            logger.debug("Preparing metadata for Mosaic using configured inputs.")
            try:
                reference_data_list: list[ImageData] = self.image_service.get_reference_image_data("Mosaic")
                # Use helper to format each item
                mosaic_inputs: list[dict[str, Any]] = [
                    self._format_mosaic_item(img_data) for img_data in reference_data_list
                ]
                return {"metadata": mosaic_inputs}
            except Exception as e:
                logger.exception("Failed to prepare metadata for Mosaic")
                raise ValueError("Failed to prepare metadata for Mosaic") from e

        # --- Handling for other mixing transforms --- #
        else:
            loaded_references: list[NDArray[np.uint8]] | None = None
            if user_provided_refs is not None:
                loaded_references = self._decode_user_refs(user_provided_refs)
            else:
                loaded_references = self._decode_default_refs(transform_class.__name__)

            if loaded_references is not None:
                return {"metadata": loaded_references}

        return None

    def _format_mosaic_item(self, img_data: ImageData) -> dict[str, Any]:
        """Decodes and formats a single item for Mosaic metadata."""
        # Decode image
        decoded_img = self.image_service.decode_base64_to_image(img_data.image)

        # Decode mask explicitly as grayscale
        decoded_mask = None
        if img_data.mask:
            try:
                mask_base64_string = img_data.mask
                if "base64," in mask_base64_string:
                    mask_base64_string = mask_base64_string.split("base64,")[1]
                mask_bytes = base64.b64decode(mask_base64_string)
                mask_np_arr = np.frombuffer(mask_bytes, np.uint8)
                decoded_mask = cv2.imdecode(mask_np_arr, cv2.IMREAD_GRAYSCALE)
                if decoded_mask is None:
                    raise ValueError("cv2.imdecode returned None for mask")
            except Exception:
                logger.exception(
                    "Failed to decode mask for Mosaic input %s",
                    img_data.name if img_data.name else "Unknown",
                )
                # Skip mask if decoding fails
                decoded_mask = None

        # Format annotations
        item_dict: dict[str, Any] = {"image": decoded_img, "mask": decoded_mask}
        if img_data.bboxes:
            item_dict["bboxes"] = [(bbox.x_min, bbox.y_min, bbox.x_max, bbox.y_max) for bbox in img_data.bboxes]
            item_dict["class_labels"] = [bbox.label for bbox in img_data.bboxes]
        else:
            item_dict["bboxes"] = []
            item_dict["class_labels"] = []

        if img_data.keypoints:
            item_dict["keypoints"] = [(kp.x, kp.y) for kp in img_data.keypoints]
            item_dict["keypoint_labels"] = [kp.label for kp in img_data.keypoints]
        else:
            item_dict["keypoints"] = []
            item_dict["keypoint_labels"] = []

        return item_dict

    def _load_annotations_if_needed(
        self,
        current_data: dict[str, Any],
        transform_class: type[A.BasicTransform],
        targets: Sequence[str],
    ) -> Annotations | None:
        """Loads annotations if any target requires them and they are missing."""
        annotations_needed = (
            ("mask" in targets and "mask" not in current_data)
            or ("bboxes" in targets and "bboxes" not in current_data)
            or ("keypoints" in targets and "keypoints" not in current_data)
        )

        if not annotations_needed:
            return None

        primary_image_set = settings.get_image_set(transform_class.__name__)
        annotation_path = settings.get_annotation_path(primary_image_set)

        if not annotation_path:
            logger.info("No annotation path configured for %s", transform_class.__name__)
            return None

        try:
            logger.debug("Attempting to load annotations from: %s", annotation_path)
            return self.image_service.load_annotations(annotation_path)
        except Exception as e:
            logger.warning("Could not load annotations from %s: %s", annotation_path, e)
            return None

    def _prepare_transform_input_data(
        self,
        initial_data: dict[str, Any],
        original_request_params: dict[str, Any] | None,
        transform_class: type[A.BasicTransform],
        config: TransformExecutionConfig,
    ) -> dict[str, Any]:
        """Ensures required data (image, mask, annotations) is loaded and formatted."""
        prepared_data = initial_data.copy()
        original_request_params = original_request_params or {}

        # Determine expected label field names from config
        bbox_label_field = (
            config.bbox_params.get("label_fields", ["class_labels"])[0] if config.bbox_params else "class_labels"
        )
        keypoint_label_field = (
            config.keypoint_params.get("label_fields", ["keypoint_labels"])[0]
            if config.keypoint_params
            else "keypoint_labels"
        )

        # 1. Ensure base image exists and is decoded
        if "image" not in prepared_data or prepared_data.get("image") is None:
            logger.warning("Image not provided, loading default for %s", transform_class.__name__)
            primary_image_data = self.image_service.get_original_image_data(transform_class.__name__)
            prepared_data["image"] = self.image_service.decode_base64_to_image(primary_image_data.image)
        elif isinstance(prepared_data.get("image"), str):
            prepared_data["image"] = self.image_service.decode_base64_to_image(prepared_data["image"])
        # Ensure it's a numpy array
        if not isinstance(prepared_data.get("image"), np.ndarray):
            raise TypeError("Image data is not a valid numpy array after loading/decoding.")

        # 2. Load and decode reference images if transform requires them
        metadata = self._prepare_metadata(original_request_params, transform_class)

        # 3. Add loaded references to prepared_data using the consistent "metadata" key
        if metadata:
            prepared_data.update(metadata)

        # 4. Load annotations *only if* needed
        annotations = self._load_annotations_if_needed(prepared_data, transform_class, config.targets)

        # 5. Populate prepared_data dict with necessary items from annotations
        if annotations:
            if "mask" in config.targets and "mask" not in prepared_data and annotations.mask is not None:
                prepared_data["mask"] = annotations.mask
            if "bboxes" in config.targets and "bboxes" not in prepared_data and annotations.bboxes:
                prepared_data["bboxes"] = [
                    (bbox.x_min, bbox.y_min, bbox.x_max, bbox.y_max) for bbox in annotations.bboxes
                ]
                prepared_data[bbox_label_field] = [bbox.label for bbox in annotations.bboxes]
            if "keypoints" in config.targets and "keypoints" not in prepared_data and annotations.keypoints:
                prepared_data["keypoints"] = [(kp.x, kp.y) for kp in annotations.keypoints]
                prepared_data[keypoint_label_field] = [kp.label for kp in annotations.keypoints]

        return prepared_data

    def _prepare_compose_args(
        self,
        prepared_data: dict[str, Any],
        config: TransformExecutionConfig,
    ) -> tuple[dict[str, Any], dict[str, Any]]:
        """Prepares the data and processor arguments for A.Compose based on targets."""
        final_data_for_compose: dict[str, Any] = {}
        # Ensure image is present
        if "image" in prepared_data and isinstance(prepared_data["image"], np.ndarray):
            final_data_for_compose["image"] = prepared_data["image"]
        else:
            raise ValueError("Prepared data must contain a valid numpy array 'image'.")

        processors: dict[str, Any] = {}

        # Call helpers to conditionally add components
        self._add_mask_to_compose(config.targets, prepared_data, final_data_for_compose)
        self._add_metadata_to_compose(prepared_data, final_data_for_compose)
        self._add_bboxes_to_compose(config, prepared_data, final_data_for_compose, processors)
        self._add_keypoints_to_compose(config, prepared_data, final_data_for_compose, processors)

        return final_data_for_compose, processors

    # --- Helpers for preparing Compose args ---

    def _get_bbox_compose_args(
        self,
        prepared_data: dict[str, Any],
        bbox_params: dict[str, Any],  # Assumed to be non-None when called
    ) -> tuple[dict[str, Any], A.BboxParams | None]:
        """Prepares bbox data and processor if validation passes."""
        data_to_add: dict[str, Any] = {}
        processor: A.BboxParams | None = None
        bbox_label_field = bbox_params.get("label_fields", ["class_labels"])[0]

        # Use .get() for safer access
        bboxes = prepared_data.get("bboxes")
        labels = prepared_data.get(bbox_label_field)

        if bboxes is not None and labels is not None:
            data_to_add["bboxes"] = bboxes
            data_to_add[bbox_label_field] = labels
            processor = A.BboxParams(
                format=bbox_params.get("format", "pascal_voc"),
                label_fields=[bbox_label_field],
                min_area=bbox_params.get("minArea", 0),
                min_visibility=bbox_params.get("minVisibility", 0),
            )
        else:
            logger.warning(
                "Skipping bboxes target: 'bboxes' or label field '%s' not found in prepared data.",
                bbox_label_field,
            )
        return data_to_add, processor

    def _get_keypoint_compose_args(
        self,
        prepared_data: dict[str, Any],
        keypoint_params: dict[str, Any],  # Assumed to be non-None when called
    ) -> tuple[dict[str, Any], A.KeypointParams | None]:
        """Prepares keypoint data and processor if validation passes."""
        data_to_add: dict[str, Any] = {}
        processor: A.KeypointParams | None = None
        keypoint_label_field = keypoint_params.get("label_fields", ["keypoint_labels"])[0]
        label_fields_param: list[str] = []

        # Use .get() for safer access
        keypoints = prepared_data.get("keypoints")
        labels = prepared_data.get(keypoint_label_field)

        if keypoints is not None:
            data_to_add["keypoints"] = keypoints
            if labels is not None:
                data_to_add[keypoint_label_field] = labels
                label_fields_param = [keypoint_label_field]

            processor = A.KeypointParams(
                format=keypoint_params.get("format", "xy"),
                label_fields=label_fields_param,
                remove_invisible=keypoint_params.get("removeInvisible", False),
            )
        else:
            logger.warning("Skipping keypoints target: 'keypoints' not found in prepared data.")

        return data_to_add, processor

    # --- Helpers for preparing _prepare_compose_args --- #

    def _add_mask_to_compose(
        self,
        targets: Sequence[str],
        prepared_data: dict[str, Any],
        final_data_for_compose: dict[str, Any],
    ) -> None:
        """Adds mask to compose data if targeted and available."""
        if "mask" in targets and "mask" in prepared_data:
            final_data_for_compose["mask"] = prepared_data["mask"]

    def _add_metadata_to_compose(self, prepared_data: dict[str, Any], final_data_for_compose: dict[str, Any]) -> None:
        """Adds metadata to compose data if available."""
        if "metadata" in prepared_data:
            final_data_for_compose["metadata"] = prepared_data["metadata"]

    def _add_bboxes_to_compose(
        self,
        config: TransformExecutionConfig,
        prepared_data: dict[str, Any],
        final_data_for_compose: dict[str, Any],
        processors: dict[str, Any],
    ) -> None:
        """Adds bboxes and bbox_params to compose args if valid."""
        if "bboxes" not in config.targets:
            return

        if "bboxes" not in prepared_data:
            logger.warning("Skipping bboxes target: No 'bboxes' found in prepared data.")
            return

        if not config.bbox_params:
            logger.warning("Skipping bboxes target: bbox_params not provided.")
            return

        bbox_data_to_add, bbox_processor = self._get_bbox_compose_args(prepared_data, config.bbox_params)
        if bbox_processor:
            final_data_for_compose.update(bbox_data_to_add)
            processors["bbox_params"] = bbox_processor

    def _add_keypoints_to_compose(
        self,
        config: TransformExecutionConfig,
        prepared_data: dict[str, Any],
        final_data_for_compose: dict[str, Any],
        processors: dict[str, Any],
    ) -> None:
        """Adds keypoints and keypoint_params to compose args if valid."""
        if "keypoints" not in config.targets:
            return

        if "keypoints" not in prepared_data:
            logger.warning("Skipping keypoints target: No 'keypoints' found in prepared data.")
            return

        if not config.keypoint_params:
            logger.warning("Skipping keypoints target: keypoint_params not provided.")
            return

        kp_data_to_add, kp_processor = self._get_keypoint_compose_args(prepared_data, config.keypoint_params)
        if kp_processor:
            final_data_for_compose.update(kp_data_to_add)
            processors["keypoint_params"] = kp_processor

    # --- Main Execution Method --- #

    def execute_transform(
        self,
        transform_class: type[A.BasicTransform],
        data: dict[str, Any],
        params: dict[str, Any],
        config: TransformExecutionConfig,
    ) -> dict[str, Any]:
        """Execute transform with given parameters."""
        original_params = params.copy() if params else {}

        # 1. Update transform-specific parameters (adds metadata_key, converts CV2)
        constructor_params = self._update_transform_params(
            params=params,
            transform_class=transform_class,
        )

        # 2. Instantiate the transform
        transform = transform_class(**constructor_params, p=1)

        # 3. Prepare the input data dictionary (load defaults, annotations, refs)
        prepared_data = self._prepare_transform_input_data(
            initial_data=data,
            original_request_params=original_params,
            transform_class=transform_class,
            config=config,
        )

        # 4. Prepare arguments for A.Compose
        final_data_for_compose, processors = self._prepare_compose_args(
            prepared_data=prepared_data,
            config=config,
        )

        # 5. Execute using Compose
        return A.Compose([transform], **processors, strict=True, telemetry=False)(**final_data_for_compose)
