# 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 to provide a registry for MLCVZoo data classes"""

import logging
from typing import Any, Dict, Optional, Type, Union

from mlcvzoo_base.api.data.bounding_box import BoundingBox
from mlcvzoo_base.api.data.classification import Classification
from mlcvzoo_base.api.data.ocr_perception import OCRPerception
from mlcvzoo_base.api.data.segmentation import Segmentation
from mlcvzoo_base.api.registry import MLCVZooRegistry

logger = logging.getLogger(__name__)


class DataRegistry(MLCVZooRegistry[Type[Classification]]):
    """
    Class to provide a registry for MLCVZoo data classes that is
    handling the encoding and decoding of such classes.
    """

    __type_key__ = "type"
    __data_key__ = "data"

    def __init__(self) -> None:
        MLCVZooRegistry.__init__(self)

        self.__data_types = [Classification, BoundingBox, Segmentation, OCRPerception]

        for d in self.__data_types:
            self.register_module(d.__name__, d)

    def encode(
        self, prediction: Union[Classification, BoundingBox, Segmentation, OCRPerception]
    ) -> Dict[str, Any]:
        """
        Encode the given data object to a dictionary, it has to be in one of
        the following formats: Classification, BoundingBox, Segmentation, OCRPerception.

        Args:
            prediction: The data to encode

        Returns:
            The encoded data
        """
        prediction_type = prediction.__class__.__name__

        if prediction_type in self._registry:
            return {
                DataRegistry.__type_key__: prediction_type,
                DataRegistry.__data_key__: prediction.to_dict(),
            }

        raise ValueError(
            "Type='%s' is not registered in the MLCVZoo DataRegistry", prediction_type
        )

    def decode(
        self, data: Dict[str, Any]
    ) -> Optional[Union[Classification, BoundingBox, Segmentation, OCRPerception]]:
        """
        Decode the data object from a dictionary

        Args:
            data: The data to decode

        Returns:
            The decoded object or None if the dictionary data does not
            contain decodable data
        """

        try:
            if data[DataRegistry.__type_key__] in self._registry:
                return self._registry[data[DataRegistry.__type_key__]].from_dict(
                    data[DataRegistry.__data_key__]
                )
        except KeyError as key_error:
            logger.error("The key '%s' is missing in the data dictionary", key_error)

        return None
