import os
import pathlib
import uuid
from typing import NamedTuple
from unittest import mock

import pytest

import mlflow
from mlflow.exceptions import MlflowException
from mlflow.utils.file_utils import mkdir, path_to_local_file_uri
from mlflow.utils.os import is_windows


class Artifact(NamedTuple):
    uri: str
    content: str


@pytest.fixture
def run_with_artifact(tmp_path):
    artifact_path = "test"
    artifact_content = "content"
    local_path = tmp_path.joinpath("file.txt")
    local_path.write_text(artifact_content)
    with mlflow.start_run() as run:
        mlflow.log_artifact(local_path, artifact_path)

    return (run, artifact_path, artifact_content)


@pytest.fixture
def run_with_artifacts(tmp_path):
    artifact_path = "test"
    artifact_content = "content"

    local_dir = tmp_path / "dir"
    local_dir.mkdir()
    local_dir.joinpath("file.txt").write_text(artifact_content)
    local_dir.joinpath("subdir").mkdir()
    local_dir.joinpath("subdir").joinpath("text.txt").write_text(artifact_content)

    with mlflow.start_run() as run:
        mlflow.log_artifact(local_dir, artifact_path)

    return (run, artifact_path)


def test_download_artifacts_with_uri(run_with_artifact):
    run, artifact_path, artifact_content = run_with_artifact
    run_uri = f"runs:/{run.info.run_id}/{artifact_path}"
    actual_uri = str(pathlib.PurePosixPath(run.info.artifact_uri) / artifact_path)
    for uri in (run_uri, actual_uri):
        download_output_path = mlflow.artifacts.download_artifacts(artifact_uri=uri)
        downloaded_artifact_path = next(pathlib.Path(download_output_path).iterdir())
        assert downloaded_artifact_path.read_text() == artifact_content


def test_download_artifacts_with_run_id_and_path(run_with_artifact):
    run, artifact_path, artifact_content = run_with_artifact
    download_output_path = mlflow.artifacts.download_artifacts(
        run_id=run.info.run_id, artifact_path=artifact_path
    )
    downloaded_artifact_path = next(pathlib.Path(download_output_path).iterdir())
    assert downloaded_artifact_path.read_text() == artifact_content


def test_download_artifacts_with_run_id_no_path(run_with_artifact):
    run, artifact_path, _ = run_with_artifact
    artifact_relative_path_top_level_dir = pathlib.PurePosixPath(artifact_path).parts[0]
    downloaded_output_path = mlflow.artifacts.download_artifacts(run_id=run.info.run_id)
    downloaded_artifact_directory_name = next(pathlib.Path(downloaded_output_path).iterdir()).name
    assert downloaded_artifact_directory_name == artifact_relative_path_top_level_dir


@pytest.mark.parametrize("dst_subdir_path", [None, "doesnt_exist_yet/subdiir"])
def test_download_artifacts_with_dst_path(run_with_artifact, tmp_path, dst_subdir_path):
    run, artifact_path, _ = run_with_artifact
    dst_path = tmp_path / dst_subdir_path if dst_subdir_path else tmp_path

    download_output_path = mlflow.artifacts.download_artifacts(
        run_id=run.info.run_id, artifact_path=artifact_path, dst_path=dst_path
    )
    assert pathlib.Path(download_output_path).samefile(dst_path / artifact_path)


def test_download_artifacts_throws_for_invalid_arguments():
    with pytest.raises(MlflowException, match="Exactly one of"):
        mlflow.artifacts.download_artifacts(
            run_id="run_id", artifact_path="path", artifact_uri="uri"
        )

    with pytest.raises(MlflowException, match="Exactly one of"):
        mlflow.artifacts.download_artifacts()

    with pytest.raises(MlflowException, match="`artifact_path` cannot be specified"):
        mlflow.artifacts.download_artifacts(artifact_path="path", artifact_uri="uri")


@pytest.fixture
def run_with_text_artifact():
    artifact_path = "test/file.txt"
    artifact_content = "This is a sentence"
    with mlflow.start_run() as run:
        mlflow.log_text(artifact_content, artifact_path)

    artifact_uri = str(pathlib.PurePosixPath(run.info.artifact_uri) / artifact_path)
    return Artifact(artifact_uri, artifact_content)


@pytest.fixture
def run_with_json_artifact():
    artifact_path = "test/config.json"
    artifact_content = {"mlflow-version": "0.28", "n_cores": "10"}
    with mlflow.start_run() as run:
        mlflow.log_dict(artifact_content, artifact_path)

    artifact_uri = str(pathlib.PurePosixPath(run.info.artifact_uri) / artifact_path)
    return Artifact(artifact_uri, artifact_content)


@pytest.fixture
def run_with_image_artifact():
    from PIL import Image

    artifact_path = "test/image.png"
    image = Image.new("RGB", (100, 100))
    with mlflow.start_run() as run:
        mlflow.log_image(image, artifact_path)

    artifact_uri = str(pathlib.PurePosixPath(run.info.artifact_uri) / artifact_path)
    return Artifact(artifact_uri, image)


def test_load_text(run_with_text_artifact):
    artifact = run_with_text_artifact
    assert mlflow.artifacts.load_text(artifact.uri) == artifact.content


def test_load_dict(run_with_json_artifact):
    artifact = run_with_json_artifact
    assert mlflow.artifacts.load_dict(artifact.uri) == artifact.content


def test_load_json_invalid_json(run_with_text_artifact):
    artifact = run_with_text_artifact
    with pytest.raises(mlflow.exceptions.MlflowException, match="Unable to form a JSON object"):
        mlflow.artifacts.load_dict(artifact.uri)


def test_load_image(run_with_image_artifact):
    from PIL import Image

    artifact = run_with_image_artifact
    assert isinstance(mlflow.artifacts.load_image(artifact.uri), Image.Image)


def test_load_image_invalid_image(run_with_text_artifact):
    artifact = run_with_text_artifact
    with pytest.raises(
        mlflow.exceptions.MlflowException, match="Unable to form a PIL Image object"
    ):
        mlflow.artifacts.load_image(artifact.uri)


class ArtifactReturnType(NamedTuple):
    tmp_path: pathlib.Path
    artifact_path: pathlib.Path
    artifact_name: str


@pytest.fixture
def text_artifact(tmp_path):
    artifact_name = "test.txt"
    artifacts_root_tmp = mkdir(tmp_path.joinpath(str(uuid.uuid4())))
    test_artifact_path = artifacts_root_tmp.joinpath(artifact_name)
    test_artifact_path.write_text("test")
    return ArtifactReturnType(artifacts_root_tmp, test_artifact_path, artifact_name)


def _assert_artifact_uri(expected_artifact_path: pathlib.Path, test_artifact):
    mlflow.log_artifact(test_artifact.artifact_path)
    assert expected_artifact_path.exists()


def test_default_relative_artifact_uri_resolves(text_artifact, tmp_path, monkeypatch):
    pytest.skip("FileStore is no longer supported.")
    tracking_uri = path_to_local_file_uri(text_artifact.tmp_path.joinpath("mlruns"))
    mlflow.set_tracking_uri(tracking_uri)
    monkeypatch.chdir(tmp_path)
    experiment_id = mlflow.create_experiment("test_exp_a", "test_artifacts_root")
    with mlflow.start_run(experiment_id=experiment_id) as run:
        _assert_artifact_uri(
            tmp_path.joinpath(
                "test_artifacts_root",
                run.info.run_id,
                "artifacts",
                text_artifact.artifact_name,
            ),
            text_artifact,
        )


def test_custom_relative_artifact_uri_resolves(text_artifact):
    pytest.skip("FileStore is no longer supported.")
    tracking_uri = path_to_local_file_uri(text_artifact.tmp_path.joinpath("tracking"))
    artifacts_root_path = text_artifact.tmp_path.joinpath("test_artifacts")
    artifacts_root_uri = path_to_local_file_uri(artifacts_root_path)
    mlflow.set_tracking_uri(tracking_uri)
    experiment_id = mlflow.create_experiment("test_exp_b", artifacts_root_uri)
    with mlflow.start_run(experiment_id=experiment_id) as run:
        _assert_artifact_uri(
            artifacts_root_path.joinpath(run.info.run_id, "artifacts", text_artifact.artifact_name),
            text_artifact,
        )


def test_artifact_logging_resolution_works_with_non_root_working_directory(tmp_path, monkeypatch):
    text_file = tmp_path.joinpath("test.txt")
    text_file.write_text("test")
    cwd = tmp_path.joinpath("cwd")
    cwd.mkdir()
    monkeypatch.chdir(cwd)
    experiment_id = mlflow.create_experiment("test_exp_c", "some_path")
    not_cwd = tmp_path.joinpath("not_cwd")
    not_cwd.mkdir()
    monkeypatch.chdir(not_cwd)

    with mlflow.start_run(experiment_id=experiment_id) as run:
        _assert_artifact_uri(
            cwd.joinpath("some_path", run.info.run_id, "artifacts", text_file.name),
            ArtifactReturnType(tmp_path, text_file, text_file.name),
        )


@pytest.mark.skipif(not is_windows(), reason="This test only passes on Windows")
def test_log_artifact_windows_path_with_hostname(text_artifact):
    experiment_test_1_artifact_location = r"\\my_server\my_path\my_sub_path\1"
    experiment_test_1_id = mlflow.create_experiment(
        "test_exp_d", experiment_test_1_artifact_location
    )
    with mlflow.start_run(experiment_id=experiment_test_1_id) as run:
        artifact_dir = rf"{experiment_test_1_artifact_location}\{run.info.run_id}\artifacts"
        temp_artifact_path = rf"{artifact_dir}\.artifact.uploading.mock"
        destination_artifact_path = rf"{artifact_dir}\{text_artifact.artifact_name}"
        real_exists = os.path.exists

        def _exists(path):
            if os.path.normcase(path) == os.path.normcase(artifact_dir):
                return mkdir_mock.called
            return real_exists(path)

        fdopen_mock = mock.MagicMock()
        fdopen_mock.__enter__.return_value = fdopen_mock
        with (
            mock.patch("shutil.copyfileobj") as copyfileobj_mock,
            mock.patch(
                "mlflow.store.artifact.local_artifact_repo.os.path.exists", side_effect=_exists
            ),
            mock.patch("mlflow.store.artifact.local_artifact_repo.mkdir") as mkdir_mock,
            mock.patch(
                "mlflow.store.artifact.local_artifact_repo.tempfile.mkstemp",
                return_value=(123, temp_artifact_path),
            ) as mkstemp_mock,
            mock.patch(
                "mlflow.store.artifact.local_artifact_repo.os.fdopen", return_value=fdopen_mock
            ) as fdopen_fn_mock,
            mock.patch("mlflow.store.artifact.local_artifact_repo.os.replace") as replace_mock,
            mock.patch("mlflow.store.artifact.local_artifact_repo.shutil.copystat"),
        ):
            mlflow.log_artifact(text_artifact.artifact_path)
            mkdir_mock.assert_called_once_with(artifact_dir)
            mkstemp_mock.assert_called_once_with(dir=artifact_dir, prefix=".artifact.uploading.")
            fdopen_fn_mock.assert_called_once_with(123, "wb")
            copyfileobj_mock.assert_called_once()
            assert copyfileobj_mock.call_args.args[1] is fdopen_mock
            replace_mock.assert_called_once_with(temp_artifact_path, destination_artifact_path)


def test_list_artifacts_with_artifact_uri(run_with_artifacts):
    run, artifact_path = run_with_artifacts
    run_uri = f"runs:/{run.info.run_id}/{artifact_path}"
    actual_uri = str(pathlib.PurePosixPath(run.info.artifact_uri) / artifact_path)
    for uri in (run_uri, actual_uri):
        artifacts = mlflow.artifacts.list_artifacts(artifact_uri=uri)
        assert len(artifacts) == 1
        assert artifacts[0].path == f"{artifact_path}/dir"

        artifacts = mlflow.artifacts.list_artifacts(artifact_uri=f"{uri}/dir")
        assert len(artifacts) == 2
        assert artifacts[0].path == "dir/file.txt"
        assert artifacts[1].path == "dir/subdir"

        artifacts = mlflow.artifacts.list_artifacts(artifact_uri=f"{uri}/dir/subdir")
        assert len(artifacts) == 1
        assert artifacts[0].path == "subdir/text.txt"

        artifacts = mlflow.artifacts.list_artifacts(artifact_uri=f"{uri}/non-exist-path")
        assert len(artifacts) == 0


def test_list_artifacts_with_run_id(run_with_artifacts):
    run, artifact_path = run_with_artifacts
    artifacts = mlflow.artifacts.list_artifacts(run_id=run.info.run_id)
    assert len(artifacts) == 1
    assert artifacts[0].path == artifact_path

    artifacts = mlflow.artifacts.list_artifacts(run_id=run.info.run_id, artifact_path=artifact_path)
    assert len(artifacts) == 1
    assert artifacts[0].path == f"{artifact_path}/dir"


def test_list_artifacts_throws_for_invalid_arguments():
    with pytest.raises(MlflowException, match="Exactly one of"):
        mlflow.artifacts.list_artifacts(
            artifact_uri="uri",
            run_id="run_id",
            artifact_path="path",
        )

    with pytest.raises(MlflowException, match="Exactly one of"):
        mlflow.artifacts.list_artifacts()

    with pytest.raises(MlflowException, match="`artifact_path` cannot be specified"):
        mlflow.artifacts.list_artifacts(artifact_uri="uri", artifact_path="path")


def test_download_artifacts_with_run_id_and_artifact_path(tmp_path):
    class DummyModel(mlflow.pyfunc.PythonModel):
        def predict(self, context, model_input: list[str]) -> list[str]:
            return model_input

    with mlflow.start_run() as run:
        mlflow.pyfunc.log_model(name="model", python_model=DummyModel())
    mlflow.artifacts.download_artifacts(
        run_id=run.info.run_id, artifact_path="model", dst_path=tmp_path
    )


def test_list_artifacts_with_client_and_tracking_uri(tmp_path: pathlib.Path):
    tracking_uri = f"sqlite:///{tmp_path}/mlflow-{uuid.uuid4().hex}.db"
    assert mlflow.get_tracking_uri() != tracking_uri
    client = mlflow.MlflowClient(tracking_uri)
    experiment_id = client.create_experiment("my_experiment")
    run = client.create_run(experiment_id)
    tmp_dir = tmp_path / "subdir"
    tmp_dir.mkdir()
    tmp_file = tmp_dir / "file.txt"
    tmp_file.touch()
    client.log_artifacts(run.info.run_id, tmp_dir, "subdir")

    artifacts = mlflow.artifacts.list_artifacts(
        run_id=run.info.run_id,
        tracking_uri=tracking_uri,
    )
    assert [p.path for p in artifacts] == ["subdir"]

    artifacts = mlflow.artifacts.list_artifacts(
        run_id=run.info.run_id,
        artifact_path="subdir",
        tracking_uri=tracking_uri,
    )
    assert [p.path for p in artifacts] == ["subdir/file.txt"]


def test_download_artifacts_with_client_and_tracking_uri(tmp_path: pathlib.Path):
    tracking_uri = f"sqlite:///{tmp_path}/mlflow-{uuid.uuid4().hex}.db"
    assert mlflow.get_tracking_uri() != tracking_uri
    client = mlflow.MlflowClient(tracking_uri)
    experiment_id = client.create_experiment("my_experiment")
    run = client.create_run(experiment_id)
    tmp_dir = tmp_path / "subdir"
    tmp_dir.mkdir()
    tmp_file = tmp_dir / "file.txt"
    tmp_file.touch()
    client.log_artifacts(run.info.run_id, tmp_dir, "subdir")

    dst_path = tmp_path / "dst"
    mlflow.artifacts.download_artifacts(
        run_id=run.info.run_id,
        artifact_path="subdir",
        tracking_uri=tracking_uri,
        dst_path=dst_path,
    )
    assert tmp_file.name in [p.name for p in dst_path.rglob("*")]


def test_single_run_artifact_download_when_both_run_and_model_artifacts_exist(tmp_path):
    class DummyModel(mlflow.pyfunc.PythonModel):
        def predict(self, model_input: list[str]) -> list[str]:
            return model_input

    with mlflow.start_run() as run:
        mlflow.pyfunc.log_model(python_model=DummyModel(), name="model")
        mlflow.log_text("test", "model/file.txt")

    out = mlflow.artifacts.download_artifacts(
        run_id=run.info.run_id, artifact_path="model/file.txt", dst_path=tmp_path
    )
    assert pathlib.Path(out).read_text() == "test"