# 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

import cv2

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.types import FrameShape
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 (
    LabelStudioAnnotationData,
    get_label_studio_annotation_dict_single_file,
    parse_annotation_data,
)
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_single_file(
                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

                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.

                        label_studio_annotation_data: LabelStudioAnnotationData = (
                            parse_annotation_data(
                                results=annotation["result"],
                                mapper=self.mapper,
                            )
                        )

                        if label_studio_annotation_data.image_shape is None:
                            _image_shape = FrameShape(*cv2.imread(image_path).shape)
                            label_studio_annotation_data.image_shape = (
                                _image_shape.height,
                                _image_shape.width,
                            )

                        annotations.append(
                            BaseAnnotation(
                                image_path=image_path,
                                annotation_path=input_data.input_path,
                                image_shape=label_studio_annotation_data.image_shape,
                                classifications=label_studio_annotation_data.classifications,
                                bounding_boxes=label_studio_annotation_data.bounding_boxes,
                                segmentations=label_studio_annotation_data.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
