import json

from .api_dto import ApiDto
from .pipeline import Pipeline
from .script import Script
from .mlmodel import MLModel

import io
import tarfile
import gzip
import dill
import pickle


def files_from_gzipped_data(gzipped_data):
    with gzip.GzipFile(fileobj=io.BytesIO(gzipped_data)) as gz:
        with tarfile.open(fileobj=gz, mode='r') as tar:
            files = {}
            for tar_info in tar.getmembers():
                if tar_info.isfile():
                    file_content = tar.extractfile(tar_info).read()
                    files[tar_info.name] = file_content
            return files


class PipelineImage(ApiDto):

    def __init__(self,
                 pipeline_image_id: str = None,
                 pipeline: Pipeline = None):

        # image definition
        self.pipeline_image_id = pipeline_image_id

        # loaded entities
        self.pipeline = pipeline
        self.scripts = []
        self.models = {}

    def get_script(self, name: str) -> Script:
        """
        get a script stored on the image from its name.
        :param name: name of the script.
        :return: loaded script.
        """
        script = None
        for i_script in self.scripts:
            if i_script.name == name:
                script = i_script
        if script is None:
            raise RuntimeError(f'not able to find {name} in pipeline image')
        return script

    @classmethod
    def load(cls,
             pipeline_image_id: str = None,
             g_stream=None) -> 'PipelineImage':
        """
        deserialize from a stream (e.g. file handle).
        :param pipeline_image_id: id of the image.
        :param g_stream: readable stream.
        :return: PipelineImage
        """
        if g_stream is None:
            raise ValueError("No stream provided")
        g_bytes = g_stream.read()
        return cls.loads(pipeline_image_id=pipeline_image_id, g_bytes=g_bytes)

    @classmethod
    def loads(cls,
              pipeline_image_id: str = None,
              g_bytes=None) -> 'PipelineImage':
        """
        deserialize from a bytes array an image.
        :param pipeline_image_id: id of the image.
        :param g_bytes: from loaded content as bytes.
        :return: PipelineImage
        """
        # unpack the content
        image = cls(pipeline_image_id=pipeline_image_id)
        if g_bytes is not None:
            g_files = files_from_gzipped_data(g_bytes)
        else:
            raise ValueError(f'please provide a valid bytes array or dict of files')

        # load the content
        try:

            # load the pipeline
            if 'pipeline.json' not in g_files:
                raise ValueError(f'invalid image')
            image.pipeline = Pipeline.from_dict(json.loads(g_files['pipeline.json']))

            # process the scripts
            script_files = {k: v for k, v in g_files.items() if k.startswith('scripts/')}

            json_files = {k for k in script_files if k.endswith('.json')}
            pkl_files = {k for k in script_files if k.endswith('.pkl')}

            for json_file in json_files:
                base_name = json_file[:-5]
                pkl_file = base_name + '.pkl'

                if pkl_file in script_files:
                    script = Script.from_dict(json.loads(g_files[json_file]))
                    script.function = dill.loads(g_files[pkl_file])
                    image.scripts.append(script)
                else:
                    raise KeyError(f"missing pkl for script json {json_file}")

            for pkl_file in pkl_files:
                base_name = pkl_file[:-4]
                json_file = base_name + '.json'

                if json_file not in json_files:
                    raise KeyError(f"missing json file for script pkl {pkl_file}")

            # process the models
            model_files = {k: v for k, v in g_files.items() if k.startswith('models/')}
            for pkl_file in model_files:
                image.models[pkl_file[7:-4]] = pickle.loads(g_files[pkl_file])

            return image

        except Exception as e:
            raise Exception(f'invalid image - {e}')

