"""API endpoints for executing Albumentations transforms."""

import logging
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from collections.abc import Sequence

import albumentations as A
from fastapi import APIRouter, HTTPException

from src.models.image import BBoxData, KeypointData
from src.models.request import TransformRequest
from src.models.transform import TransformExecutionConfig, TransformResponse
from src.services.image_service import ImageService
from src.services.transform_service import TransformService

router = APIRouter(tags=["execution"])
image_service = ImageService()
transform_service = TransformService()

logger = logging.getLogger(__name__)


@router.post("/execute-transform")
async def execute_transform(request: TransformRequest) -> TransformResponse:
    """Execute a transform with given parameters."""
    try:
        # Validate and prepare for transform execution
        prepared_params = _validate_and_prepare_transform(request)

        # Execute transform
        result = transform_service.execute_transform(
            transform_class=prepared_params["transform_class"],
            data=prepared_params["data"],
            params=prepared_params["params"],
            config=prepared_params["config"],
        )

        # Format response
        return _format_transform_response(request, result)

    except HTTPException:
        # Re-raise HTTP exceptions (like our 400 validation errors)
        raise
    except Exception as e:
        logger.exception("Error executing transform %s", request.transform)
        # All other exceptions become 500 errors
        raise HTTPException(status_code=500, detail=str(e)) from e


def _validate_and_prepare_transform(request: TransformRequest) -> dict[str, Any]:
    """Validate request and prepare parameters for transform execution."""
    # Get transform class - this can raise a ValueError if transform not found
    transform_class = transform_service.get_transform_class(request.transform)

    # Prepare data dictionary
    data: dict[str, Any] = {}

    # Get image from request or use original image
    if request.data and request.data.image is not None:
        data["image"] = image_service.decode_base64_to_image(request.data.image)
    else:
        # Call get_original_image_data and decode the result
        original_image_data = image_service.get_original_image_data(request.transform)
        data["image"] = image_service.decode_base64_to_image(original_image_data.image)

    # Add any additional fields from the request data
    if request.data:
        # We use model_dump() to get all fields including extra ones
        request_data_dict = request.data.model_dump(exclude={"image"})
        # Use dictionary comprehension instead of for-loop
        data.update({key: value for key, value in request_data_dict.items() if key != "image" and value is not None})

    # Validate targets based on transform type
    if issubclass(transform_class, A.ImageOnlyTransform):
        # ImageOnlyTransform can only transform images
        invalid_targets = set(request.targets) - {"image"}
        msg = f"ImageOnlyTransform '{request.transform}' can only transform 'image', but got: {invalid_targets}"
        if invalid_targets:
            raise HTTPException(
                status_code=400,
                detail=msg,
            )

    # Validate that required parameters are provided for requested targets
    if "bboxes" in request.targets and not request.bbox_params:
        raise HTTPException(
            status_code=400,
            detail="bbox_params must be provided when 'bboxes' is in targets",
        )

    if "keypoints" in request.targets and not request.keypoint_params:
        raise HTTPException(
            status_code=400,
            detail="keypoint_params must be provided when 'keypoints' is in targets",
        )

    # Only pass bbox_params and keypoint_params if they're in the targets
    bbox_params = request.bbox_params if "bboxes" in request.targets else None
    keypoint_params = request.keypoint_params if "keypoints" in request.targets else None

    # Create config object
    config = TransformExecutionConfig(
        targets=request.targets,
        bbox_params=bbox_params,
        keypoint_params=keypoint_params,
    )

    # Prepare parameters for transform execution
    return {
        "transform_class": transform_class,
        "data": data,
        "params": request.params,
        "config": config,
    }


def _format_transform_response(request: TransformRequest, result: dict[str, Any]) -> TransformResponse:
    """Format transform result into API response."""
    # Prepare base response data
    response_data: dict[str, Any] = {
        "image": image_service.encode_image_to_base64(result["image"]),
    }

    # Add mask if requested and available
    if "mask" in request.targets and "mask" in result:
        response_data["mask"] = image_service.encode_image_to_base64(result["mask"])
    else:
        response_data["mask"] = None  # Ensure key exists

    # Add bounding boxes if requested and available
    if "bboxes" in request.targets and "bboxes" in result:
        response_data["bboxes"] = _format_bboxes(result)
    else:
        response_data["bboxes"] = None  # Ensure key exists

    # Add keypoints if requested and available
    if "keypoints" in request.targets and "keypoints" in result:
        response_data["keypoints"] = _format_keypoints(result)
    else:
        response_data["keypoints"] = None  # Ensure key exists

    # Create and return TransformResponse object
    return TransformResponse(**response_data)


def _format_bboxes(result: dict[str, Any]) -> list[BBoxData]:
    """Format bounding boxes from transform result."""
    formatted_bboxes: list[BBoxData] = []
    bboxes: Sequence[Sequence[float]] = result.get("bboxes", [])
    labels: Sequence[str] = result.get("class_labels", [])

    for i, bbox in enumerate(bboxes):
        # Check if bbox is in [x_min, y_min, x_max, y_max] format
        if len(bbox) == 4:
            label = labels[i] if i < len(labels) else None
            formatted_bboxes.append(
                BBoxData(
                    x_min=float(bbox[0]),
                    y_min=float(bbox[1]),
                    x_max=float(bbox[2]),
                    y_max=float(bbox[3]),
                    label=label,
                ),
            )
    return formatted_bboxes


def _format_keypoints(result: dict[str, Any]) -> list[KeypointData]:
    """Format keypoints from transform result."""
    formatted_keypoints: list[KeypointData] = []
    keypoints: Sequence[Sequence[float]] = result.get("keypoints", [])
    labels: Sequence[str | None] = result.get("keypoint_labels", [])  # Labels can be None

    for i, keypoint in enumerate(keypoints):
        # Keypoints are in [x, y] format
        if len(keypoint) == 2:
            label = labels[i] if i < len(labels) else None
            formatted_keypoints.append(
                KeypointData(
                    x=float(keypoint[0]),
                    y=float(keypoint[1]),
                    label=label,
                ),
            )
    return formatted_keypoints
