import json
import pathlib
import pickle

import numpy as np
import pandas as pd
import pytest
from matplotlib.figure import Figure

from mlflow.exceptions import MlflowException
from mlflow.models.evaluation.artifacts import (
    CsvEvaluationArtifact,
    ImageEvaluationArtifact,
    JsonEvaluationArtifact,
    NumpyEvaluationArtifact,
    ParquetEvaluationArtifact,
    PickleEvaluationArtifact,
    TextEvaluationArtifact,
    _infer_artifact_type_and_ext,
)
from mlflow.models.evaluation.default_evaluator import _CustomArtifact


@pytest.fixture
def cm_fn_tuple():
    return _CustomArtifact(lambda: None, "", 0, "")


def __generate_dummy_json_file(path):
    with open(path, "w") as f:
        json.dump([1, 2, 3], f)


class __DummyClass:
    def __init__(self):
        self.test = 1


@pytest.mark.parametrize(
    ("is_file", "artifact", "artifact_type", "ext"),
    [
        (True, lambda path: Figure().savefig(path), ImageEvaluationArtifact, "png"),
        (True, lambda path: Figure().savefig(path), ImageEvaluationArtifact, "jpg"),
        (True, lambda path: Figure().savefig(path), ImageEvaluationArtifact, "jpeg"),
        (True, __generate_dummy_json_file, JsonEvaluationArtifact, "json"),
        (True, lambda path: pathlib.Path(path).write_text("test"), TextEvaluationArtifact, "txt"),
        (
            True,
            lambda path: np.save(path, np.array([1, 2, 3]), allow_pickle=False),
            NumpyEvaluationArtifact,
            "npy",
        ),
        (
            True,
            lambda path: pd.DataFrame({"test": [1, 2, 3]}).to_csv(path, index=False),
            CsvEvaluationArtifact,
            "csv",
        ),
        (
            True,
            lambda path: pd.DataFrame({"test": [1, 2, 3]}).to_parquet(path),
            ParquetEvaluationArtifact,
            "parquet",
        ),
        (False, pd.DataFrame({"test": [1, 2, 3]}), CsvEvaluationArtifact, "csv"),
        (False, np.array([1, 2, 3]), NumpyEvaluationArtifact, "npy"),
        (False, Figure(), ImageEvaluationArtifact, "png"),
        (False, {"a": 1, "b": "e", "c": 1.2, "d": [1, 2]}, JsonEvaluationArtifact, "json"),
        (False, [1, 2, 3, "test"], JsonEvaluationArtifact, "json"),
        (False, '{"a": 1, "b": [1.2, 3]}', JsonEvaluationArtifact, "json"),
        (False, '[1, 2, 3, "test"]', JsonEvaluationArtifact, "json"),
        (False, __DummyClass(), PickleEvaluationArtifact, "pickle"),
    ],
)
def test_infer_artifact_type_and_ext(is_file, artifact, artifact_type, ext, tmp_path, cm_fn_tuple):
    if is_file:
        artifact_representation = tmp_path / f"test.{ext}"
        artifact(artifact_representation)
    else:
        artifact_representation = artifact
    inferred_from_path, inferred_type, inferred_ext = _infer_artifact_type_and_ext(
        f"{ext}_{artifact_type.__name__}_artifact", artifact_representation, cm_fn_tuple
    )
    assert not is_file ^ inferred_from_path
    assert inferred_type is artifact_type
    assert inferred_ext == f".{ext}"


def test_infer_artifact_type_and_ext_raise_exception_for_non_file_non_json_str(cm_fn_tuple):
    with pytest.raises(
        MlflowException,
        match="with string representation 'some random str' that is "
        "neither a valid path to a file nor a JSON string",
    ):
        _infer_artifact_type_and_ext("test_artifact", "some random str", cm_fn_tuple)


def test_infer_artifact_type_and_ext_raise_exception_for_non_existent_path(tmp_path, cm_fn_tuple):
    path = tmp_path / "does_not_exist_path"
    with pytest.raises(MlflowException, match=f"with path '{path}' does not exist"):
        _infer_artifact_type_and_ext("test_artifact", path, cm_fn_tuple)


def test_infer_artifact_type_and_ext_raise_exception_for_non_file_artifact(tmp_path, cm_fn_tuple):
    with pytest.raises(MlflowException, match=f"with path '{tmp_path}' is not a file"):
        _infer_artifact_type_and_ext("non_file_artifact", tmp_path, cm_fn_tuple)


def test_infer_artifact_type_and_ext_raise_exception_for_unsupported_ext(tmp_path, cm_fn_tuple):
    path = tmp_path / "invalid_ext_example.some_ext"
    with open(path, "w") as f:
        f.write("some stuff that shouldn't be read")
    with pytest.raises(
        MlflowException,
        match=f"with path '{path}' does not match any of the supported file extensions",
    ):
        _infer_artifact_type_and_ext("invalid_ext_artifact", path, cm_fn_tuple)


def test_pickle_evaluation_artifact_load_raises_when_pickle_deserialization_disabled(
    tmp_path, monkeypatch
):
    monkeypatch.setenv("MLFLOW_ALLOW_PICKLE_DESERIALIZATION", "false")

    artifact_path = tmp_path / "artifact.pickle"
    with open(artifact_path, "wb") as f:
        pickle.dump({"key": "value"}, f)

    artifact = PickleEvaluationArtifact(uri="test_uri")
    with pytest.raises(MlflowException, match="MLFLOW_ALLOW_PICKLE_DESERIALIZATION"):
        artifact._load_content_from_file(artifact_path)