import os
from shlex import quote
from unittest import mock

import pytest

from mlflow.exceptions import ExecutionException
from mlflow.projects._project_spec import EntryPoint
from mlflow.utils.file_utils import TempDir, path_to_local_file_uri

from tests.projects.utils import TEST_PROJECT_DIR, load_project


def test_entry_point_compute_params():
    """
    Tests that EntryPoint correctly computes a final set of parameters to use when running a project
    """
    project = load_project()
    entry_point = project.get_entry_point("greeter")
    # Pass extra "excitement" param, use default value for `greeting` param
    with TempDir() as storage_dir:
        params, extra_params = entry_point.compute_parameters(
            {"name": "friend", "excitement": 10}, storage_dir
        )
        assert params == {"name": "friend", "greeting": "hi"}
        assert extra_params == {quote("excitement"): quote("10")}
        # Don't pass extra "excitement" param, pass value for `greeting`
        params, extra_params = entry_point.compute_parameters(
            {"name": "friend", "greeting": "hello"}, storage_dir
        )
        assert params == {"name": "friend", "greeting": "hello"}
        assert extra_params == {}
        # Raise exception on missing required parameter
        with pytest.raises(
            ExecutionException, match="No value given for missing parameters: 'name'"
        ):
            entry_point.compute_parameters({}, storage_dir)


def test_entry_point_compute_command():
    """
    Tests that EntryPoint correctly computes the command to execute in order to run the entry point.
    """
    project = load_project()
    entry_point = project.get_entry_point("greeter")
    with TempDir() as tmp:
        storage_dir = tmp.path()
        command = entry_point.compute_command({"name": "friend", "excitement": 10}, storage_dir)
        assert command == "python greeter.py hi friend --excitement 10"
        with pytest.raises(
            ExecutionException, match="No value given for missing parameters: 'name'"
        ):
            entry_point.compute_command({}, storage_dir)
        # Test shell escaping
        name_value = "friend; echo 'hi'"
        command = entry_point.compute_command({"name": name_value}, storage_dir)
        assert command == "python greeter.py {} {}".format(quote("hi"), quote(name_value))


@pytest.mark.parametrize(
    "malicious_key",
    ["x;touch /tmp/pwned", "a&b", "a|b", "a`b`", "a$(b)", "a>b"],
)
def test_compute_command_shell_escapes_extra_param_keys(malicious_key):
    entry_point = load_project().get_entry_point("greeter")
    with TempDir() as tmp:
        command = entry_point.compute_command({"name": "friend", malicious_key: "1"}, tmp.path())
        # Raw command must contain the shell-quoted key so bash parses it as a
        # single (nonsensical) flag token rather than executing the payload.
        assert f"--{quote(malicious_key)}" in command


def test_path_parameter():
    """
    Tests that MLflow file-download APIs get called when necessary for arguments of type `path`.
    """
    project = load_project()
    entry_point = project.get_entry_point("line_count")
    with mock.patch(
        "mlflow.tracking.artifact_utils._download_artifact_from_uri", return_value=0
    ) as download_uri_mock:
        # Verify that we don't attempt to call download_uri when passing a local file to a
        # parameter of type "path"
        with TempDir() as tmp:
            dst_dir = tmp.path()
            local_path = os.path.join(TEST_PROJECT_DIR, "MLproject")
            params, _ = entry_point.compute_parameters(
                user_parameters={"path": local_path}, storage_dir=dst_dir
            )
            assert params["path"] == os.path.abspath(local_path)
            assert download_uri_mock.call_count == 0

            params, _ = entry_point.compute_parameters(
                user_parameters={"path": path_to_local_file_uri(local_path)}, storage_dir=dst_dir
            )
            assert params["path"] == os.path.abspath(local_path)
            assert download_uri_mock.call_count == 0

        # Verify that we raise an exception when passing a non-existent local file to a
        # parameter of type "path"
        with TempDir() as tmp:
            dst_dir = tmp.path()
            with pytest.raises(ExecutionException, match="no such file or directory"):
                entry_point.compute_parameters(
                    user_parameters={"path": os.path.join(dst_dir, "some/nonexistent/file")},
                    storage_dir=dst_dir,
                )
        # Verify that we do call `download_uri` when passing a URI to a parameter of type "path"
        for i, prefix in enumerate(["dbfs:/", "s3://", "gs://"]):
            with TempDir() as tmp:
                dst_dir = tmp.path()
                file_to_download = "images.tgz"
                download_path = f"{dst_dir}/{file_to_download}"
                download_uri_mock.return_value = download_path
                params, _ = entry_point.compute_parameters(
                    user_parameters={"path": os.path.join(prefix, file_to_download)},
                    storage_dir=dst_dir,
                )
                assert params["path"] == download_path
                assert download_uri_mock.call_count == i + 1


def test_uri_parameter():
    project = load_project()
    entry_point = project.get_entry_point("download_uri")
    with (
        mock.patch(
            "mlflow.tracking.artifact_utils._download_artifact_from_uri"
        ) as download_uri_mock,
        TempDir() as tmp,
    ):
        dst_dir = tmp.path()
        # Test that we don't attempt to locally download parameters of type URI
        entry_point.compute_command(
            user_parameters={"uri": f"file://{dst_dir}"}, storage_dir=dst_dir
        )
        assert download_uri_mock.call_count == 0
        # Test that we raise an exception if a local path is passed to a parameter of type URI
        with pytest.raises(ExecutionException, match="Expected URI for parameter uri"):
            entry_point.compute_command(user_parameters={"uri": dst_dir}, storage_dir=dst_dir)


def test_params():
    defaults = {
        "alpha": "float",
        "l1_ratio": {"type": "float", "default": 0.1},
        "l2_ratio": {"type": "float", "default": 0.0003},
        "random_str": {"type": "string", "default": "hello"},
    }
    entry_point = EntryPoint("entry_point_name", defaults, "command_name script.py")

    user1 = {}
    with pytest.raises(ExecutionException, match="No value given for missing parameters"):
        entry_point._validate_parameters(user1)

    user_2 = {"beta": 0.004}
    with pytest.raises(ExecutionException, match="No value given for missing parameters"):
        entry_point._validate_parameters(user_2)

    user_3 = {"alpha": 0.004, "gamma": 0.89}
    expected_final_3 = {
        "alpha": "0.004",
        "l1_ratio": "0.1",
        "l2_ratio": "0.0003",
        "random_str": "hello",
    }
    expected_extra_3 = {quote("gamma"): quote("0.89")}
    final_3, extra_3 = entry_point.compute_parameters(user_3, None)
    assert expected_extra_3 == extra_3
    assert expected_final_3 == final_3

    user_4 = {"alpha": 0.004, "l1_ratio": 0.0008, "random_str_2": "hello"}
    expected_final_4 = {
        "alpha": "0.004",
        "l1_ratio": "0.0008",
        "l2_ratio": "0.0003",
        "random_str": "hello",
    }
    expected_extra_4 = {quote("random_str_2"): quote("hello")}
    final_4, extra_4 = entry_point.compute_parameters(user_4, None)
    assert expected_extra_4 == extra_4
    assert expected_final_4 == final_4

    user_5 = {"alpha": -0.99, "random_str": "hi"}
    expected_final_5 = {
        "alpha": "-0.99",
        "l1_ratio": "0.1",
        "l2_ratio": "0.0003",
        "random_str": "hi",
    }
    expected_extra_5 = {}
    final_5, extra_5 = entry_point.compute_parameters(user_5, None)
    assert expected_final_5 == final_5
    assert expected_extra_5 == extra_5

    user_6 = {"alpha": 0.77, "ALPHA": 0.89}
    expected_final_6 = {
        "alpha": "0.77",
        "l1_ratio": "0.1",
        "l2_ratio": "0.0003",
        "random_str": "hello",
    }
    expected_extra_6 = {quote("ALPHA"): quote("0.89")}
    final_6, extra_6 = entry_point.compute_parameters(user_6, None)
    assert expected_extra_6 == extra_6
    assert expected_final_6 == final_6


def test_path_params():
    data_file = "s3://path.test/resources/data_file.csv"
    defaults = {
        "constants": {"type": "uri", "default": "s3://path.test/b1"},
        "data": {"type": "path", "default": data_file},
    }
    entry_point = EntryPoint("entry_point_name", defaults, "command_name script.py")

    with mock.patch(
        "mlflow.tracking.artifact_utils._download_artifact_from_uri", return_value=None
    ) as download_uri_mock:
        final_1, extra_1 = entry_point.compute_parameters({}, None)
        assert final_1 == {"constants": "s3://path.test/b1", "data": data_file}
        assert extra_1 == {}
        assert download_uri_mock.call_count == 0

    with mock.patch(
        "mlflow.tracking.artifact_utils._download_artifact_from_uri"
    ) as download_uri_mock:
        user_2 = {"alpha": 0.001, "constants": "s3://path.test/b_two"}
        final_2, extra_2 = entry_point.compute_parameters(user_2, None)
        assert final_2 == {"constants": "s3://path.test/b_two", "data": data_file}
        assert extra_2 == {"alpha": "0.001"}
        assert download_uri_mock.call_count == 0

    with (
        mock.patch(
            "mlflow.tracking.artifact_utils._download_artifact_from_uri"
        ) as download_uri_mock,
        TempDir() as tmp,
    ):
        dest_path = tmp.path()
        download_path = f"{dest_path}/data_file.csv"
        download_uri_mock.return_value = download_path
        user_3 = {"alpha": 0.001}
        final_3, extra_3 = entry_point.compute_parameters(user_3, dest_path)
        assert final_3 == {"constants": "s3://path.test/b1", "data": download_path}
        assert extra_3 == {"alpha": "0.001"}
        assert download_uri_mock.call_count == 1

    with (
        mock.patch(
            "mlflow.tracking.artifact_utils._download_artifact_from_uri"
        ) as download_uri_mock,
        TempDir() as tmp,
    ):
        dest_path = tmp.path()
        download_path = f"{dest_path}/images.tgz"
        download_uri_mock.return_value = download_path
        user_4 = {"data": "s3://another.example.test/data_stash/images.tgz"}
        final_4, extra_4 = entry_point.compute_parameters(user_4, dest_path)
        assert final_4 == {"constants": "s3://path.test/b1", "data": download_path}
        assert extra_4 == {}
        assert download_uri_mock.call_count == 1