from os.path import isfile
import os
import mlflow
import copy
from pathlib import Path
from typing import Any, Dict
import pandas as pd
import cloudpickle
from kedro.io.core import AbstractVersionedDataSet, DataSetError, Version
import mlflow.pyfunc
import zipfile



class ZipDataSet(AbstractVersionedDataSet):
    def __init__(self, filepath, version=None):
        super().__init__(Path(filepath), version)

    def _load(self) -> pd.DataFrame:
        load_path = self._get_load_path()
        return zipfile.ZipFile(load_path,"r")

    def _save(self, zipfile) -> None:
        save_path = self._get_save_path()
        
        with zipfile as f:
            f.extractall(save_path)

    def _describe(self):
        return dict(version=self._version)



class ScikitWrapper(mlflow.pyfunc.PythonModel):
    def load_context(self, context):
        import cloudpickle

        # print(context.artifacts["model_path"])
        with open(context.artifacts["model_path"], "rb+") as f:
            self.model = cloudpickle.load(f)

    def predict(self, context, model_input=None, X=None):
        if X is None:
            X = model_input
        return self.model.predict(X)


class SklearnLocalDataSet(AbstractVersionedDataSet):
    """``SklearnLocalDataSet`` loads and saves data to a local Sklearn models and pipelines. The
    underlying functionality is supported by mlflow
    
    https://www.mlflow.org/docs/latest/python_api/mlflow.sklearn.html#module-mlflow.sklearn.
    ::

        >>> from kedro.io import SklearnLocalDataSet
        >>> import sklearn.ensemble import RandomForestClassifier
        >>>
        >>> model = RandomForestClassifier().fit(X, y)
        >>> data_set = SklearnLocalDataSet('myModel')
        >>> data_set.save(model)
        >>> loaded_model = data_set.load()
        >>> assert model.equals(loaded_model)
    """

    DEFAULT_LOAD_ARGS = {}  # type: Dict[str, Any]
    DEFAULT_SAVE_ARGS = {"serialization_format": "cloudpickle"}  # type: Dict[str, Any]

    # pylint: disable=too-many-arguments
    def __init__(
        self,
        filepath: str,
        conda_env: str = "src/conda.yml",
        load_args: Dict[str, Any] = None,
        save_args: Dict[str, Any] = None,
        version: Version = None,
    ) -> None:
        """Creates a new instance of ``SklearnLocalDataSet`` pointing to a
        concrete filepath.

        Args:
            filepath: Path to a parquet file or a metadata file of a multipart
                parquet collection or the directory of a multipart parquet.
                
            conda_env: Path to anaconda environement required to use the model
            
            load_args: Additional loading options

            save_args: Additional saving options

            version: If specified, should be an instance of
                ``kedro.io.core.Version``. If its ``load`` attribute is
                None, the latest version will be loaded. If its ``save``
                attribute is None, save version will be autogenerated.

        """
        super().__init__(Path(filepath), version)

        # Handle default load and save arguments
        self._load_args = copy.deepcopy(self.DEFAULT_LOAD_ARGS)
        if load_args is not None:
            self._load_args.update(load_args)
        self._save_args = copy.deepcopy(self.DEFAULT_SAVE_ARGS)
        if save_args is not None:
            self._save_args.update(save_args)

        if conda_env is None:
            self.conda_env = mlflow.sklearn.get_default_conda_env(
                include_cloudpickle=True
            )
        else:
            self.conda_env = conda_env

    def _describe(self) -> Dict[str, Any]:
        return dict(
            filepath=self._filepath,
            load_args=self._load_args,
            save_args=self._save_args,
            version=self._version,
        )

    def _load(self) -> pd.DataFrame:
        load_path = self._get_load_path()
        print(load_path)
        model_path = os.path.join(load_path, "model.pkl")

        # return mlflow.sklearn.load_model(model_uri=str(load_path),
        #                                 **self._load_args)
        return mlflow.pyfunc.load_model(str(load_path))

    def _save(self, data: pd.DataFrame) -> None:
        save_path = self._get_save_path()
        print(save_path)
        # mlflow.sklearn.save_model(sk_model=data,
        #                          path=save_path,
        #                          conda_env = self.conda_env,
        #                          **self._save_args)

        if not os.path.exists(".tmp"):
            os.makedirs(".tmp")

        model_path = os.path.join(".tmp", "model.pkl")

        with open(model_path, "wb+") as f:
            cloudpickle.dump(data, f)

        mlflow.pyfunc.save_model(
            save_path,
            python_model=ScikitWrapper(),
            artifacts={"model_path": model_path},
            code_path=["src/stitch_classify"],
            conda_env="src/environment.yml",
        )

    def _exists(self) -> bool:
        try:
            path = self._get_load_path()
        except DataSetError:
            return False
        return Path(path).is_file()