import pathlib
from json import JSONDecodeError

import pandas as pd
import numpy as np
import json
import pickle
import matplotlib.pyplot as plt
from collections import namedtuple

from mlflow.exceptions import MlflowException
from mlflow.models.evaluation.base import EvaluationArtifact
from mlflow.utils.annotations import developer_stable
from mlflow.utils.proto_json_utils import NumpyEncoder


@developer_stable
class ImageEvaluationArtifact(EvaluationArtifact):
    def _save(self, output_artifact_path):
        self._content.save(output_artifact_path)

    def _load_content_from_file(self, local_artifact_path):
        from PIL.Image import open as open_image

        self._content = open_image(local_artifact_path)
        self._content.load()  # Load image and close the file descriptor.
        return self._content


@developer_stable
class CsvEvaluationArtifact(EvaluationArtifact):
    def _save(self, output_artifact_path):
        self._content.to_csv(output_artifact_path, index=False)

    def _load_content_from_file(self, local_artifact_path):
        self._content = pd.read_csv(local_artifact_path)
        return self._content


@developer_stable
class ParquetEvaluationArtifact(EvaluationArtifact):
    def _save(self, output_artifact_path):
        self._content.to_parquet(output_artifact_path, compression="brotli")

    def _load_content_from_file(self, local_artifact_path):
        self._content = pd.read_parquet(local_artifact_path)
        return self._content


@developer_stable
class NumpyEvaluationArtifact(EvaluationArtifact):
    def _save(self, output_artifact_path):
        np.save(output_artifact_path, self._content, allow_pickle=False)

    def _load_content_from_file(self, local_artifact_path):
        self._content = np.load(local_artifact_path, allow_pickle=False)
        return self._content


@developer_stable
class JsonEvaluationArtifact(EvaluationArtifact):
    def _save(self, output_artifact_path):
        with open(output_artifact_path, "w") as f:
            json.dump(self._content, f)

    def _load_content_from_file(self, local_artifact_path):
        with open(local_artifact_path) as f:
            self._content = json.load(f)
        return self._content


@developer_stable
class TextEvaluationArtifact(EvaluationArtifact):
    def _save(self, output_artifact_path):
        with open(output_artifact_path, "w") as f:
            f.write(self._content)

    def _load_content_from_file(self, local_artifact_path):
        with open(local_artifact_path) as f:
            self._content = f.read()
        return self._content


@developer_stable
class PickleEvaluationArtifact(EvaluationArtifact):
    def _save(self, output_artifact_path):
        with open(output_artifact_path, "wb") as f:
            pickle.dump(self._content, f)

    def _load_content_from_file(self, local_artifact_path):
        with open(local_artifact_path, "rb") as f:
            self._content = pickle.load(f)
        return self._content


_EXT_TO_ARTIFACT_MAP = {
    ".png": ImageEvaluationArtifact,
    ".jpg": ImageEvaluationArtifact,
    ".jpeg": ImageEvaluationArtifact,
    ".json": JsonEvaluationArtifact,
    ".npy": NumpyEvaluationArtifact,
    ".csv": CsvEvaluationArtifact,
    ".parquet": ParquetEvaluationArtifact,
    ".txt": TextEvaluationArtifact,
}

_TYPE_TO_EXT_MAP = {
    pd.DataFrame: ".csv",
    np.ndarray: ".npy",
    plt.Figure: ".png",
}

_TYPE_TO_ARTIFACT_MAP = {
    pd.DataFrame: CsvEvaluationArtifact,
    np.ndarray: NumpyEvaluationArtifact,
    plt.Figure: ImageEvaluationArtifact,
}

_InferredArtifactProperties = namedtuple(
    "_InferredArtifactProperties", ["from_path", "type", "ext"]
)


def _infer_artifact_type_and_ext(artifact_name, raw_artifact, custom_metric_tuple):
    """
    This function performs type and file extension inference on the provided artifact
    :param artifact_name: The name of the provided artifact
    :param raw_artifact: The artifact object
    :param custom_metric_tuple: Containing a user provided function and its index in the
                                ``custom_metrics`` parameter of ``mlflow.evaluate``
    :return: InferredArtifactProperties namedtuple
    """

    exception_header = (
        f"Custom metric function '{custom_metric_tuple.name}' at index "
        f"{custom_metric_tuple.index} in the `custom_metrics` parameter produced an "
        f"artifact '{artifact_name}'"
    )

    # Given a string, first see if it is a path. Otherwise, check if it is a JsonEvaluationArtifact
    if isinstance(raw_artifact, str):
        potential_path = pathlib.Path(raw_artifact)
        if potential_path.exists():
            raw_artifact = potential_path
        else:
            try:
                json.loads(raw_artifact)
                return _InferredArtifactProperties(
                    from_path=False, type=JsonEvaluationArtifact, ext=".json"
                )
            except JSONDecodeError:
                raise MlflowException(
                    f"{exception_header} with string representation '{raw_artifact}' that is "
                    f"neither a valid path to a file nor a JSON string."
                )

    # Type inference based on the file extension
    if isinstance(raw_artifact, pathlib.Path):
        if not raw_artifact.exists():
            raise MlflowException(f"{exception_header} with path '{raw_artifact}' does not exist.")
        if not raw_artifact.is_file():
            raise MlflowException(f"{exception_header} with path '{raw_artifact}' is not a file.")
        if raw_artifact.suffix not in _EXT_TO_ARTIFACT_MAP:
            raise MlflowException(
                f"{exception_header} with path '{raw_artifact}' does not match any of the supported"
                f" file extensions: {', '.join(_EXT_TO_ARTIFACT_MAP.keys())}."
            )
        return _InferredArtifactProperties(
            from_path=True, type=_EXT_TO_ARTIFACT_MAP[raw_artifact.suffix], ext=raw_artifact.suffix
        )

    # Type inference based on object type
    if type(raw_artifact) in _TYPE_TO_ARTIFACT_MAP:
        return _InferredArtifactProperties(
            from_path=False,
            type=_TYPE_TO_ARTIFACT_MAP[type(raw_artifact)],
            ext=_TYPE_TO_EXT_MAP[type(raw_artifact)],
        )

    # Given as other python object, we first attempt to infer as JsonEvaluationArtifact. If that
    # fails, we store it as PickleEvaluationArtifact
    try:
        json.dumps(raw_artifact, cls=NumpyEncoder)
        return _InferredArtifactProperties(
            from_path=False, type=JsonEvaluationArtifact, ext=".json"
        )
    except TypeError:
        return _InferredArtifactProperties(
            from_path=False, type=PickleEvaluationArtifact, ext=".pickle"
        )
