import io
import os
import pickle
from abc import ABC, abstractmethod

from conductor.utils.sagemaker_training_environment import create_trainer_environment


class Model(ABC):
    def __init__(self) -> None:
        self.train_env = create_trainer_environment()

    @abstractmethod
    def train(self):
        pass

    @abstractmethod
    def predict(self, data: io.StringIO) -> io.StringIO:
        pass

    def predict_mimetype(self):
        return None

    def dump(self) -> None:
        with open(os.path.join(self.train_env.model_dir, "model.pkl"), "wb+") as out:
            pickle.dump(self, out)

    @staticmethod
    def load(path: str):
        with open(path, "rb") as inp:
            return pickle.load(inp)
