# 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 from a single json file, which
is generated by the export method of Label Studio for one task
"""

import logging
import os
from typing import List, Optional, Tuple

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.exceptions import ForbiddenClassError
from mlcvzoo_base.configuration.annotation_handler_config import (
    AnnotationHandlerSingleFileInputDataConfig,
)
from mlcvzoo_base.data_preparation.annotation_parser.label_studio_annotation_parser import (
    get_label_studio_annotation_dict,
    parse_from_annotation,
)
from mlcvzoo_base.utils.file_utils import get_image_file_dict

logger = logging.getLogger(__name__)


class LabelStudioAnnotationParserSingle(AnnotationParser):
    """
    Parser for the Label Studio json format of a single json file that is
    generated by the export method of Label Studio for one Project

    Support data types:
    - BoundingBox
    """

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

        self.label_studio_input_data: List[AnnotationHandlerSingleFileInputDataConfig] = (
            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_root_dir, image_format=input_data.image_format
            )
            label_studio_dict = get_label_studio_annotation_dict(
                label_studio_annotation_path=input_data.input_path
            )
            if label_studio_dict is None:
                continue

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

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

            for label_studio_task in label_studio_dict:
                label_studio_image_name = os.path.basename(label_studio_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]] = None
                bounding_boxes: List[BoundingBox] = []
                try:
                    for annotation in label_studio_task["annotations"]:
                        # REMARK: image_shape will be set multiple times.
                        #         It is assumed that all annotations belong to
                        #         the same image and therefore only ONE image_shape
                        #         has to be set.

                        (
                            image_shape,
                            _bounding_boxes,
                        ) = parse_from_annotation(
                            results=annotation["result"],
                            mapper=self.mapper,
                        )
                        bounding_boxes.extend(_bounding_boxes)

                    if image_shape is not None:
                        annotations.append(
                            BaseAnnotation(
                                image_path=image_path,
                                annotation_path=input_data.input_path,
                                image_shape=image_shape,
                                classifications=[],
                                bounding_boxes=bounding_boxes,
                                segmentations=[],
                                image_dir=os.path.dirname(image_path),
                                annotation_dir=os.path.dirname(input_data.input_path),
                                # TODO: How to set this?
                                replacement_string="",
                            )
                        )
                except (ValueError, ForbiddenClassError, KeyError) as error:
                    logger.warning("%s, annotation will be skipped" % str(error))
                    continue

        return annotations
