import uuid

from fedn.common.log_config import logger
from fedn.network.storage.s3.miniorepository import MINIORepository


class Repository:
    """ Interface for storing model objects and compute packages in S3 compatible storage. """

    def __init__(self, config):

        self.model_bucket = config['storage_bucket']
        self.context_bucket = config['context_bucket']

        # TODO: Make a plug-in solution
        self.client = MINIORepository(config)

        self.client.create_bucket(self.context_bucket)
        self.client.create_bucket(self.model_bucket)

    def get_model(self, model_id):
        """ Retrieve a model with id model_id.

        :param model_id: Unique identifier for model to retrive.
        :return: The model object
        """
        logger.info("Client {} trying to get model with id: {}".format(
            self.client.name, model_id))
        return self.client.get_artifact(model_id, self.model_bucket)

    def get_model_stream(self, model_id):
        """ Retrieve a stream handle to model with id model_id.

        :param model_id:
        :return: Handle to model object
        """
        logger.info("Client {} trying to get model with id: {}".format(
            self.client.name, model_id))
        return self.client.get_artifact_stream(model_id, self.model_bucket)

    def set_model(self, model, is_file=True):
        """ Upload model object.

        :param model: The model object
        :type model: BytesIO or str file name.
        :param is_file: True if model is a file name, else False
        :return: id for the uploaded object (str)
        """
        model_id = uuid.uuid4()

        try:
            self.client.set_artifact(str(model_id), model,
                                     bucket=self.model_bucket, is_file=is_file)
        except Exception:
            logger.error("Failed to upload model with ID {} to repository.".format(model_id))
            raise
        return str(model_id)

    def set_compute_package(self, name, compute_package, is_file=True):
        """ Upload compute package.

        :param name: The name of the compute package.
        :type name: str
        :param compute_package: The compute package
        :type compute_pacakge: BytesIO or str file name.
        :param is_file: True if model is a file name, else False
        """

        try:
            self.client.set_artifact(str(name), compute_package,
                                     bucket=self.context_bucket, is_file=is_file)
        except Exception:
            logger.error("Failed to write compute_package to repository.")
            raise

    def get_compute_package(self, compute_package):
        """ Retrieve compute package from object store.

        :param compute_package: The name of the compute package.
        :type compute_pacakge: str
        :return: Compute package.
        """
        try:
            data = self.client.get_artifact(compute_package, bucket=self.context_bucket)
        except Exception:
            logger.error("Failed to get compute_package from repository.")
            raise
        return data

    def delete_compute_package(self, compute_package):
        """ Delete a compute package from storage.

        :param compute_package: The name of the compute_package
        :type compute_package: str
        """

        try:
            self.client.delete_artifact(compute_package, bucket=[self.context_bucket])
        except Exception:
            logger.error("Failed to delete compute_package from repository.")
            raise
