import abc
import gc
import typing
from pathlib import Path

from source.manager import ModelManager


class BaseModel(abc.ABC):
    """
    Represent a model.
    """

    def __init__(self, manager: ModelManager, configuration: dict[str, typing.Any], path: Path):
        # the environment directory of the model
        self.path = path
        # the model manager
        self.manager = manager
        # the mimetype of the model responses
        self.output_type: str = configuration.get("output_type", "application/json")

        self._loaded = False

    def __repr__(self):
        return f"<{self.__class__.__name__}: {self.name}>"

    @property
    def name(self):
        """
        Get the name of the model
        :return: the name of the model
        """

        return self.path.name

    def get_information(self):
        """
        Get information about the model
        :return: information about the model
        """

        return {
            "name": self.name,
            "output_type": self.output_type,
        }

    def load(self) -> None:
        """
        Load the model within the model manager
        """

        # if we are already loaded, stop
        if self._loaded:
            return

        # check if we are the current loaded model
        if self.manager.current_loaded_model is not self:
            # unload the previous model
            if self.manager.current_loaded_model is not None:
                self.manager.current_loaded_model.unload()

        # model specific loading
        self._load()

        # declare ourselves as the currently loaded model
        self.manager.current_loaded_model = self

        # mark the model as loaded
        self._loaded = True

    @abc.abstractmethod
    def _load(self):
        """
        Load the model
        Do not call manually, use `load` instead.
        """

    def unload(self) -> None:
        """
        Unload the model within the model manager
        """

        # if we are not already loaded, stop
        if not self._loaded:
            return

        # if we were the currently loaded model of the manager, demote ourselves
        if self.manager.current_loaded_model is self:
            self.manager.current_loaded_model = None

        # model specific unloading part
        self._unload()

        # force the garbage collector to clean the memory
        gc.collect()

        # mark the model as unloaded
        self._loaded = False

    @abc.abstractmethod
    def _unload(self):
        """
        Unload the model
        Do not call manually, use `unload` instead.
        """

    def infer(self, *args, **kwargs) -> typing.Iterator[bytes]:
        """
        Infer our payload through the model within the model manager
        :return: the response of the model
        """

        # make sure we are loaded before an inference
        self.load()

        # model specific inference part
        return self._infer(*args, **kwargs)

    @abc.abstractmethod
    def _infer(self, *args, **kwargs) -> typing.Iterator[bytes]:
        """
        Infer our payload through the model
        :return: the response of the model
        """