# 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 writing MMPretrain formatted annotations."""

import logging
import os
from typing import List, Optional

from mlcvzoo_base.api.data.annotation import BaseAnnotation
from mlcvzoo_base.api.data.annotation_writer import AnnotationWriter
from mlcvzoo_base.utils import encoding_safe_imwrite

logger = logging.getLogger(__name__)


class ClassificationAnnotationWriter(AnnotationWriter):
    """
    Class for writing classification based annotations.
    """

    def __init__(
        self,
        output_file_path: str,
        data_root_dir: str,
        seperator: str,
    ):
        self.output_file_path = output_file_path
        self.data_root_dir = data_root_dir
        self.seperator = seperator

    def write(self, annotations: List[BaseAnnotation]) -> Optional[str]:

        os.makedirs(os.path.dirname(self.output_file_path), exist_ok=True)

        logger.info("Write annotations to %s", self.output_file_path)

        with open(self.output_file_path, "w") as annotation_file:
            for annotation in annotations:
                for classification in annotation.classifications:
                    annotation_file.write(
                        f"{os.path.relpath(annotation.image_path, self.data_root_dir)}"
                        f"{self.seperator}"
                        f"{classification.class_id}\n"
                    )

        return self.output_file_path


class TextCropAnnotationWriter(AnnotationWriter):
    """
    Class for writing annotations for cropped text images.
    """

    def __init__(
        self,
        output_file_path: str,
        data_root_dir: str,
        seperator: str,
        rel_image_output_dir: str,
        padding_ratio: float = 0.1,
    ):
        self.output_file_path: str = output_file_path
        self.data_root_dir: str = data_root_dir
        self.seperator: str = seperator
        self.rel_image_output_dir = rel_image_output_dir
        self.padding_ratio: float = padding_ratio

    @staticmethod
    def __create_cropped_image_filepath(
        image_path: str,
        content: str,
        class_name: str,
        index: int,
        image_output_dir: str,
    ) -> str:
        """Create image path for cropped image.

        Args:
            image_path: The annotations to crop
            content: The bounding box or segmentation to crop
            class_name: The class name for the cropped image
            index: The index of the cropped image
            image_output_dir: The directory to save cropped images to,
                if not given save cropped image to original image path

        Returns:
            The created image path
        """
        image_base_name = os.path.basename(image_path)

        image_name_split = image_base_name.split(".")
        assert len(image_name_split) == 2
        return os.path.join(
            image_output_dir,
            # Replace spaces with underscores to avoid failures
            image_name_split[0].replace(" ", "_"),
            f"{class_name}_{index}_{content}.jpg",
        )

    def write(self, annotations: List[BaseAnnotation]) -> Optional[str]:

        os.makedirs(os.path.dirname(self.output_file_path), exist_ok=True)

        logger.info("Write annotations to %s", self.output_file_path)

        with open(self.output_file_path, "w") as annotation_file:
            for annotation in annotations:
                cropped_image_infos = annotation.crop_images(padding_ratio=self.padding_ratio)
                for index, cropped_image_info in enumerate(cropped_image_infos):
                    if cropped_image_info.geometric_object.content == "":
                        continue

                    if self.rel_image_output_dir is None:
                        image_output_dir = os.path.join(
                            self.data_root_dir,
                            os.path.basename(self.output_file_path).split(".")[0],
                        )
                    else:
                        image_output_dir = os.path.join(
                            self.data_root_dir,
                            self.rel_image_output_dir,
                        )

                    cropped_image_path = TextCropAnnotationWriter.__create_cropped_image_filepath(
                        image_path=annotation.image_path,
                        content=cropped_image_info.geometric_object.content,
                        class_name=cropped_image_info.geometric_object.class_name,
                        index=index,
                        image_output_dir=image_output_dir,
                    )

                    os.makedirs(os.path.dirname(cropped_image_path), exist_ok=True)

                    encoding_safe_imwrite(
                        filename=cropped_image_path,
                        image=cropped_image_info.image,
                    )

                    annotation_file.write(
                        f"{os.path.relpath(cropped_image_path, self.data_root_dir)}"
                        f"{self.seperator}"
                        f"{cropped_image_info.geometric_object.content}\n"
                    )

        return self.output_file_path
