from typing import List
from pathlib import Path
from abc import abstractmethod
import shutil
from loguru import logger

from .transformer import ModelTransformer
from .model import Model
from pydlmeta.identifier.types import ModelFormat
from pydlmeta.identifier.model  import identify
from ..common import get_tmp_path


class SerializerRegistry(type):

    REGISTRY: List = []

    def __new__(cls, name, bases, attrs):

        new_cls = type.__new__(cls, name, bases, attrs)
        cls.REGISTRY.append(new_cls)
        return new_cls


class Serializer(ModelTransformer, metaclass=SerializerRegistry):

    FORMAT = None

    @classmethod
    def is_me(cls, model: Model) -> bool:
        return identify(model.src) == cls.FORMAT

    @abstractmethod
    def transform(self, model: Model) -> Model:
        raise NotImplementedError("`transform` has to be implemented")

    def serialize(self, model: Model, dest: Path) -> Model:
        self.add_param('dest', dest)  # type: ignore[attr-defined]
        m = self.transform(model)
        return m


def serializer_selector(model: Model) -> Serializer:
    for serializer in Serializer.REGISTRY:
        if serializer.is_me(model):
            return serializer
    raise NotImplementedError(f"Unable to serialize {model}")


def serialize(model: Model, dest: Path) -> Model:
    for serializer in Serializer.REGISTRY:
        if serializer.is_me(model):
            return serializer().serialize(model, dest)

    raise NotImplementedError(f"Unable to serialize {model}")


class FileSerializer(Serializer):

    FORMAT: ModelFormat = ModelFormat.NON_SPECIFIED

    def transform(self, model: Model):
        dest = self.get_param('dest')

        if model.src.is_file():
            dest = dest.with_suffix(model.src.suffix)
            shutil.copy(model.src, dest)
        else:
            raise Exception(f'src [`{model.src}`] should be a path to a file.')

        return Model(dest)

    def serialize(self, model: Model, dest: Path) -> Model:
        # return super().serialize(model, dest)
        self.add_param('dest', dest)  # type: ignore[attr-defined]
        m = self.transform(model)
        return m


class DirSerializer(Serializer):

    FORMAT: ModelFormat = ModelFormat.NON_SPECIFIED

    def transform(self, model: Model):
        dest = self.get_param('dest')

        if Path(dest).exists() and Path(dest).is_dir():
            raise Exception(f'dest `{dest}` should be a file path.')
        elif not Path(dest).exists:
            if dest[-1] in ['\\', '/']:
                raise Exception(f"dest `{dest}` should file path, not a dir")

        if model.src.is_dir():
            # Zip the dir if the model is in dir form
            dest = dest.with_suffix('.zip')
            tmp_path = Path(get_tmp_path()[:-1])
            shutil.make_archive(str(tmp_path), 'zip', model.src)
            shutil.move(str(tmp_path) + '.zip', str(dest))
        else:
            raise Exception(
                f'model source `{model.src}` should be a path to a dir.')

        return Model(dest)

    def serialize(self, model: Model, dest: Path) -> Model:
        # return super().serialize(model, dest)
        self.add_param('dest', dest)  # type: ignore[attr-defined]
        m = self.transform(model)
        return m


class CaffeDir(DirSerializer):

    FORMAT = ModelFormat.CAFFE_DIR


class H5(FileSerializer):

    FORMAT = ModelFormat.H5


class ONNX(FileSerializer):

    FORMAT = ModelFormat.ONNX


class PTH(FileSerializer):

    FORMAT = ModelFormat.PTH


class PB(FileSerializer):

    FORMAT = ModelFormat.PB


class TorchTraced(FileSerializer):

    FORMAT = ModelFormat.TORCH_TRACED


class TFLITE(FileSerializer):

    FORMAT = ModelFormat.TFLITE


class OpenvinoIRDir(DirSerializer):

    FORMAT = ModelFormat.OPENVINO_IRDIR


class ZippedOpenvinoIRDir(FileSerializer):

    FORMAT = ModelFormat.ZIPPED_OPENVINO_IRDIR


class SavedModel(DirSerializer):

    FORMAT = ModelFormat.SAVED_MODEL


class ZippedSavedModel(FileSerializer):

    FORMAT = ModelFormat.ZIPPED_SAVED_MODEL


class TFKerasModel(Serializer):
    '''
    '''

    FORMAT = ModelFormat.TF_KERAS_MODEL

    def transform(self, model: Model):
        dest = self.get_param('dest')

        # TF model.save uses file ext to determine the format to
        # be saved. if not specified, SavedModel willl be used.
        model.src.save(dest.with_suffix('.h5'))
        shutil.move(dest.with_suffix('.h5'), str(dest))
        return Model(dest)


class KerasModel(Serializer):
    '''
    Keras 2.5.0 Serializer
    '''

    FORMAT = ModelFormat.KERAS_MODEL

    def transform(self, model: Model):
        dest = self.get_param('dest')
        model.src.save(dest.with_suffix('.h5'))
        shutil.move(dest.with_suffix('.h5'), str(dest))
        return Model(dest)


class TF_Session(Serializer):
    '''
    Keras 2.5.0 Serializer
    '''

    FORMAT = ModelFormat.TF_SESSION

    def transform(self, model: Model):
        import tensorflow

        dest = self.get_param('dest')

        saver = tensorflow.train.Saver()

        sess = model.src
        sess.run(tensorflow.initialize_all_variables())

        model.src.save(dest.with_suffix('.h5'))
        shutil.move(dest.with_suffix('.h5'), str(dest))
        return Model(dest)


class PytorchModel(Serializer):
    """Use python MRO to check if it contains specific str"""

    FORMAT = ModelFormat.PT_NN_MODULE

    def transform(self, model: Model):
        import torch

        dest = self.get_param('dest')

        if not model.inputs:
            raise Exception(
                "PytorchModel requires input shpae, please add an input tensor in the Model"
            )

        for input_tensor in model.inputs:
            try:
                if not all(isinstance(x, int) for x in input_tensor.shape):
                    raise Exception(f"Parameter `shape` must be List[int]")
            except Exception as e:
                raise Exception("Parameter `shape` must be List[int]")

        dummy_inputs = []
        if not len(model.inputs) > 0:
            raise Exception(
                "Pytorch model has to have at least one input tensor")

        for input_tensor in model.inputs:
            shape = list(input_tensor.shape)
            shape[0] = 1
            shape = tuple(shape)
            dummy_inputs.append(torch.rand(*shape))

        # torch.onnx.export(model.src, tuple(dummy_inputs), dest,
        #                   input_names=[x.name for x in model.inputs],
        #                   output_names=[x.name for x in model.outputs])

        try:
            traced_module = torch.jit.trace(model.src, tuple(dummy_inputs))
        except ValueError as ve:
            logger.error("Unable to save the Pytorch model. "
                         "Specify `inputs` in project.add_model() "
                         "may solve this issue.")
            raise ve

        torch.jit.save(traced_module, dest)

        return Model(dest)
