# Copyright Open Logistics Foundation
#
# Licensed under the Open Logistics Foundation License 1.3.
# For details on the licensing terms, see the LICENSE file.
# SPDX-License-Identifier: OLFL-1.3

"""Module for parsing Label Studio formatted annotations of single json files that are
generated by Label Studio per image
"""

import json
import logging
import os
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union, cast

import cv2
import numpy as np

from mlcvzoo_base.api.data.annotation import BaseAnnotation
from mlcvzoo_base.api.data.annotation_class_mapper import AnnotationClassMapper
from mlcvzoo_base.api.data.annotation_parser import AnnotationParser
from mlcvzoo_base.api.data.bounding_box import BoundingBox
from mlcvzoo_base.api.data.box import Box, rotate_point
from mlcvzoo_base.api.data.class_identifier import ClassIdentifier
from mlcvzoo_base.api.data.types import Point2f
from mlcvzoo_base.api.exceptions import ClassMappingNotFoundError, ForbiddenClassError
from mlcvzoo_base.configuration.annotation_handler_config import (
    AnnotationHandlerLabelStudioInputDataConfig,
)
from mlcvzoo_base.configuration.structs import ObjectDetectionBBoxFormats
from mlcvzoo_base.utils.file_utils import get_image_file_dict

logger = logging.getLogger(__name__)


def get_label_studio_annotation_dict(
    label_studio_annotation_path: str,
) -> Optional[Union[Dict[str, Any], List[Dict[str, Any]]]]:
    with open(label_studio_annotation_path, "r") as label_studio_file:
        try:
            return cast(List[Dict[str, Any]], json.load(label_studio_file))
        except json.decoder.JSONDecodeError as json_error:
            logger.exception(
                msg="Could not parse annotation file '%s'. Annotation will be skipped."
                % label_studio_annotation_path,
                exc_info=json_error,
            )

    return None


def parse_bounding_box(
    mapper: AnnotationClassMapper, label_studio_value: Dict[str, Any], image_shape: Tuple[int, int]
) -> Optional[BoundingBox]:
    label_studio_class_name: str = label_studio_value["rectanglelabels"][0]

    try:
        class_name = mapper.map_annotation_class_name_to_model_class_name(
            class_name=label_studio_class_name
        )
        class_id = mapper.map_annotation_class_name_to_model_class_id(
            class_name=label_studio_class_name
        )
    except ClassMappingNotFoundError:
        logger.debug(
            "Could not find a valid class-mapping for class-name '%s'. "
            "BoundingBox will be skipped",
            label_studio_class_name,
        )
        return None

    # LabelStudio has defined its rotation around the top-left point of the bounding box.
    # The RotatedBox of the MLCVZoo uses the opencv definition and rotates around the center
    # point of the bounding box. Therefore, we have to transform the LabelStudio definition
    # to the opencv definition.

    top_left_x = float(image_shape[1] * label_studio_value["x"] / 100)
    top_left_y = float(image_shape[0] * label_studio_value["y"] / 100)

    dim_x = float(image_shape[1] * label_studio_value["width"] / 100)
    dim_y = float(image_shape[0] * label_studio_value["height"] / 100)

    rotation = label_studio_value["rotation"]

    if rotation > 0.0:
        top_left: Point2f = (top_left_x, top_left_y)
        top_right: Point2f = (top_left_x + dim_x, top_left_y)
        bottom_right: Point2f = (top_left_x + dim_x, top_left_y + dim_y)
        bottom_left: Point2f = (top_left_x, top_left_y + dim_y)

        top_right = rotate_point(top_right, top_left, label_studio_value["rotation"])
        bottom_right = rotate_point(bottom_right, top_left, label_studio_value["rotation"])
        bottom_left = rotate_point(bottom_left, top_left, label_studio_value["rotation"])

        box = np.asarray([top_left, top_right, bottom_right, bottom_left], dtype=np.int_)  # type: ignore[var-annotated]
        (center_x, center_y), (w, h), angle = cv2.minAreaRect(box)

        return BoundingBox(
            class_identifier=ClassIdentifier(
                class_id=class_id,
                class_name=class_name,
            ),
            box=Box.init_format_based(
                box_list=(float(center_x), float(center_y), float(w), float(h)),
                box_format=ObjectDetectionBBoxFormats.CXCYWH,
                angle=angle,
            ),
            score=1.0,
            difficult=False,
            occluded=False,
            content="",
            background=False,
        )

    return BoundingBox(
        class_identifier=ClassIdentifier(
            class_id=class_id,
            class_name=class_name,
        ),
        box=Box.init_format_based(
            box_list=(top_left_x, top_left_y, dim_x, dim_y),
            box_format=ObjectDetectionBBoxFormats.XYWH,
            angle=0.0,
        ),
        score=1.0,
        difficult=False,
        occluded=False,
        content="",
        background=False,
    )


def parse_from_annotation(
    results: Dict[Any, Any], mapper: AnnotationClassMapper
) -> Tuple[Optional[Tuple[int, int]], List[BoundingBox]]:
    image_shape: Optional[Tuple[int, int]] = None
    bounding_boxes: List[BoundingBox] = []

    for result in results:
        if image_shape is None:
            image_shape = (int(result["original_height"]), int(result["original_width"]))

        if result["type"] == "rectanglelabels":
            bounding_box: Optional[BoundingBox]
            bounding_box = parse_bounding_box(
                mapper=mapper,
                label_studio_value=result["value"],
                image_shape=image_shape,
            )
            if bounding_box is not None:
                bounding_boxes.append(bounding_box)

    return image_shape, bounding_boxes


class LabelStudioAnnotationParser(AnnotationParser):
    """
    Parser for the Label Studio json format of a json file that is
    generated by Label Studio per image

    Support data types:
    - BoundingBox
    """

    def __init__(
        self,
        mapper: AnnotationClassMapper,
        label_studio_input_data: List[AnnotationHandlerLabelStudioInputDataConfig],
    ):
        AnnotationParser.__init__(self, mapper=mapper)

        self.label_studio_input_data = label_studio_input_data

    def parse(self) -> List[BaseAnnotation]:
        annotations: List[BaseAnnotation] = []

        for dataset_count, input_data in enumerate(self.label_studio_input_data):
            image_file_dict = get_image_file_dict(
                input_dir=input_data.input_image_dir, image_format=input_data.image_format
            )

            label_studio_annotation_paths = [
                str(p) for p in Path(input_data.input_annotation_dir).glob("**/*")
            ]

            for label_studio_annotation_path in label_studio_annotation_paths:
                label_studio_dict = get_label_studio_annotation_dict(
                    label_studio_annotation_path=label_studio_annotation_path
                )

                if label_studio_dict is None:
                    logger.warning(
                        "Label Studio annotation content could not be loaded from: '%s', "
                        "Annotation will be skipped." % label_studio_annotation_path
                    )
                    continue

                if not isinstance(label_studio_dict, dict):
                    logger.warning(
                        "Loaded Label Studio annotation content is not of type dict, file: '%s', "
                        "Annotation will be skipped." % label_studio_annotation_path
                    )
                    continue

                try:
                    label_studio_image_name = os.path.basename(
                        label_studio_dict["task"]["data"]["image"]
                    )

                    if label_studio_image_name in image_file_dict:
                        image_path = image_file_dict[label_studio_image_name]
                    else:
                        logger.warning(
                            "Label Studio annotation contains non-existing image: '%s', "
                            "Annotation will be skipped." % label_studio_image_name
                        )
                        continue

                    image_shape: Optional[Tuple[int, int]]
                    bounding_boxes: List[BoundingBox]
                    (
                        image_shape,
                        bounding_boxes,
                    ) = parse_from_annotation(
                        results=label_studio_dict["result"],
                        mapper=self.mapper,
                    )

                    if image_shape is not None:
                        annotations.append(
                            BaseAnnotation(
                                image_path=image_path,
                                annotation_path=label_studio_annotation_path,
                                image_shape=image_shape,
                                classifications=[],
                                bounding_boxes=bounding_boxes,
                                segmentations=[],
                                image_dir=os.path.dirname(image_path),
                                annotation_dir=os.path.dirname(label_studio_annotation_path),
                                # TODO: How to set this?
                                replacement_string="",
                            )
                        )
                except (ValueError, ForbiddenClassError, KeyError) as error:
                    logger.exception(
                        msg="Could not parse annotation file '%s'. Annotation will be skipped."
                        % label_studio_annotation_path,
                        exc_info=error,
                    )

        return annotations
