import io
import itertools
import os
import pickle
import uuid
from importlib import reload
from pathlib import Path
from unittest import mock
from urllib.parse import urlparse
from urllib.request import url2pathname

import pytest

import mlflow
from mlflow.environment_variables import (
    MLFLOW_ENABLE_WORKSPACES,
    MLFLOW_TRACKING_INSECURE_TLS,
    MLFLOW_TRACKING_PASSWORD,
    MLFLOW_TRACKING_TOKEN,
    MLFLOW_TRACKING_URI,
    MLFLOW_TRACKING_USERNAME,
)
from mlflow.exceptions import MlflowException
from mlflow.server import ARTIFACT_ROOT_ENV_VAR
from mlflow.store.db.db_types import DATABASE_ENGINES
from mlflow.store.tracking.databricks_rest_store import DatabricksTracingRestStore
from mlflow.store.tracking.file_store import FileStore
from mlflow.store.tracking.rest_store import RestStore
from mlflow.store.tracking.sqlalchemy_store import SqlAlchemyStore
from mlflow.tracking._tracking_service.registry import TrackingStoreRegistry
from mlflow.tracking._tracking_service.utils import (
    _get_store,
    _get_tracking_scheme,
    _resolve_custom_scheme,
    _resolve_tracking_uri,
    _use_tracking_uri,
    get_tracking_uri,
    set_tracking_uri,
)
from mlflow.tracking.registry import UnsupportedModelRegistryStoreURIException
from mlflow.utils.file_utils import path_to_local_file_uri
from mlflow.utils.os import is_windows

from tests.helpers.db_mocks import mock_get_managed_session_maker
from tests.tracing.helper import get_tracer_tracking_uri

# Disable mocking tracking URI here, as we want to test setting the tracking URI via
# environment variable. See
# http://doc.pytest.org/en/latest/skipping.html#skip-all-test-functions-of-a-class-or-module
# and https://github.com/mlflow/mlflow/blob/master/CONTRIBUTING.md#writing-python-tests
# for more information.
pytestmark = pytest.mark.notrackingurimock


@pytest.mark.skip(reason="FileStore is no longer supported")
def test_tracking_scheme_with_existing_mlruns(tmp_path, monkeypatch):
    monkeypatch.chdir(tmp_path)
    mlruns_dir = tmp_path / "mlruns"
    mlruns_dir.mkdir()
    exp_dir = mlruns_dir / "0"
    exp_dir.mkdir()
    (exp_dir / "meta.yaml").touch()
    store = _get_store()
    assert isinstance(store, FileStore)


def test_tracking_scheme_without_existing_mlruns(tmp_path, monkeypatch):
    monkeypatch.chdir(tmp_path)
    store = _get_store()
    assert isinstance(store, SqlAlchemyStore)


@pytest.mark.skip(reason="FileStore is no longer supported")
def test_get_store_with_existing_mlruns_data(tmp_path, monkeypatch):
    monkeypatch.chdir(tmp_path)
    mlruns_dir = tmp_path / "mlruns"
    mlruns_dir.mkdir()
    exp_dir = mlruns_dir / "0"
    exp_dir.mkdir()
    (exp_dir / "meta.yaml").touch()

    store = _get_store()
    assert isinstance(store, FileStore)
    assert os.path.abspath(store.root_directory) == os.path.abspath("mlruns")


def test_get_store_with_empty_mlruns(tmp_path, monkeypatch):
    monkeypatch.chdir(tmp_path)
    mlruns_dir = tmp_path / "mlruns"
    mlruns_dir.mkdir()

    store = _get_store()
    assert isinstance(store, SqlAlchemyStore)


def test_get_store_with_mlruns_dir_but_no_meta_yaml(tmp_path, monkeypatch):
    monkeypatch.chdir(tmp_path)
    mlruns_dir = tmp_path / "mlruns"
    mlruns_dir.mkdir()
    (mlruns_dir / "0").mkdir()

    store = _get_store()
    assert isinstance(store, SqlAlchemyStore)


def test_default_sqlite_tracking_uri_respects_cwd(tmp_path, monkeypatch):
    monkeypatch.chdir(tmp_path)
    with _use_tracking_uri(None):
        store = _get_store()

    assert isinstance(store, SqlAlchemyStore)
    sqlite_uri = store.db_uri
    assert sqlite_uri.startswith("sqlite:")
    parsed = urlparse(sqlite_uri)
    path = parsed.path
    if not parsed.netloc and path.startswith("//"):
        path = path[1:]
    if parsed.netloc:
        path = f"//{parsed.netloc}{path}"
    db_path = Path(url2pathname(path))
    assert db_path.parent == tmp_path


@pytest.mark.skip(reason="FileStore is no longer supported")
def test_get_store_file_store_from_arg(tmp_path, monkeypatch):
    monkeypatch.chdir(tmp_path)
    store = _get_store("other/path")
    assert isinstance(store, FileStore)
    assert os.path.abspath(store.root_directory) == os.path.abspath("other/path")


@pytest.mark.skip(reason="FileStore is no longer supported")
@pytest.mark.parametrize("uri", ["other/path", "file:other/path"])
def test_get_store_file_store_from_env(tmp_path, monkeypatch, uri):
    monkeypatch.chdir(tmp_path)
    monkeypatch.setenv(MLFLOW_TRACKING_URI.name, uri)
    store = _get_store()
    assert isinstance(store, FileStore)
    assert os.path.abspath(store.root_directory) == os.path.abspath("other/path")


def test_get_store_basic_rest_store(monkeypatch):
    monkeypatch.setenv(MLFLOW_TRACKING_URI.name, "https://my-tracking-server:5050")
    store = _get_store()
    assert isinstance(store, RestStore)
    assert store.get_host_creds().host == "https://my-tracking-server:5050"
    assert store.get_host_creds().token is None
    assert _get_tracking_scheme() == "https"


def test_get_store_rest_store_with_password(monkeypatch):
    for k, v in {
        MLFLOW_TRACKING_URI.name: "https://my-tracking-server:5050",
        MLFLOW_TRACKING_USERNAME.name: "Bob",
        MLFLOW_TRACKING_PASSWORD.name: "Ross",
    }.items():
        monkeypatch.setenv(k, v)

    store = _get_store()
    assert isinstance(store, RestStore)
    assert store.get_host_creds().host == "https://my-tracking-server:5050"
    assert store.get_host_creds().username == "Bob"
    assert store.get_host_creds().password == "Ross"


def test_get_store_rest_store_with_token(monkeypatch):
    for k, v in {
        MLFLOW_TRACKING_URI.name: "https://my-tracking-server:5050",
        MLFLOW_TRACKING_TOKEN.name: "my-token",
    }.items():
        monkeypatch.setenv(k, v)

    store = _get_store()
    assert isinstance(store, RestStore)
    assert store.get_host_creds().token == "my-token"


def test_get_store_rest_store_with_insecure(monkeypatch):
    for k, v in {
        MLFLOW_TRACKING_URI.name: "https://my-tracking-server:5050",
        MLFLOW_TRACKING_INSECURE_TLS.name: "true",
    }.items():
        monkeypatch.setenv(k, v)
    store = _get_store()
    assert isinstance(store, RestStore)
    assert store.get_host_creds().ignore_tls_verification


def test_get_store_rest_store_with_no_insecure(monkeypatch):
    with monkeypatch.context() as m:
        for k, v in {
            MLFLOW_TRACKING_URI.name: "https://my-tracking-server:5050",
            MLFLOW_TRACKING_INSECURE_TLS.name: "false",
        }.items():
            m.setenv(k, v)
        store = _get_store()
        assert isinstance(store, RestStore)
        assert not store.get_host_creds().ignore_tls_verification

    # By default, should not ignore verification.
    with monkeypatch.context() as m:
        monkeypatch.setenv(MLFLOW_TRACKING_URI.name, "https://my-tracking-server:5050")
        store = _get_store()
        assert isinstance(store, RestStore)
        assert not store.get_host_creds().ignore_tls_verification


@pytest.mark.parametrize("db_type", DATABASE_ENGINES)
def test_get_store_sqlalchemy_store(tmp_path, monkeypatch, db_type):
    monkeypatch.chdir(tmp_path)
    monkeypatch.setenv(MLFLOW_ENABLE_WORKSPACES.name, "false")
    uri = f"{db_type}://hostname/database-{uuid.uuid4().hex}"
    monkeypatch.setenv(MLFLOW_TRACKING_URI.name, uri)
    monkeypatch.delenv("MLFLOW_SQLALCHEMYSTORE_POOLCLASS", raising=False)
    with (
        mock.patch("sqlalchemy.create_engine") as mock_create_engine,
        mock.patch("sqlalchemy.event.listens_for"),
        mock.patch("mlflow.store.db.utils._verify_schema"),
        mock.patch("mlflow.store.db.utils._initialize_tables"),
        mock.patch(
            "mlflow.store.db.utils._get_managed_session_maker",
            new=mock_get_managed_session_maker,
        ),
        mock.patch(
            # In sqlalchemy 1.4.0, `SqlAlchemyStore.search_experiments`, which is called when
            # fetching the store, results in an error when called with a mocked sqlalchemy engine.
            # Accordingly, we mock `SqlAlchemyStore.search_experiments`
            "mlflow.store.tracking.sqlalchemy_store.SqlAlchemyStore.search_experiments",
            return_value=[],
        ),
        mock.patch(
            "mlflow.store.tracking.sqlalchemy_store.SqlAlchemyStore._initialize_store_state",
            return_value=None,
        ),
    ):
        store = _get_store()
        assert isinstance(store, SqlAlchemyStore)
        assert store.db_uri == uri
        # Create another store to ensure the engine is cached
        another_store = _get_store()
        assert store.engine is another_store.engine
        if is_windows():
            assert store.artifact_root_uri == Path.cwd().joinpath("mlruns").as_uri()
        else:
            assert store.artifact_root_uri == Path.cwd().joinpath("mlruns").as_posix()
        assert _get_tracking_scheme() == db_type

    mock_create_engine.assert_called_once_with(uri, pool_pre_ping=True)


@pytest.mark.parametrize("db_type", DATABASE_ENGINES)
def test_get_store_sqlalchemy_store_with_artifact_uri(tmp_path, monkeypatch, db_type):
    monkeypatch.chdir(tmp_path)
    monkeypatch.setenv(MLFLOW_ENABLE_WORKSPACES.name, "false")
    uri = f"{db_type}://hostname/database-{uuid.uuid4().hex}"
    artifact_uri = "file:artifact/path"
    monkeypatch.setenv(MLFLOW_TRACKING_URI.name, uri)
    monkeypatch.delenv("MLFLOW_SQLALCHEMYSTORE_POOLCLASS", raising=False)
    with (
        mock.patch("sqlalchemy.create_engine") as mock_create_engine,
        mock.patch("sqlalchemy.event.listens_for"),
        mock.patch("mlflow.store.db.utils._verify_schema"),
        mock.patch("mlflow.store.db.utils._initialize_tables"),
        mock.patch(
            "mlflow.store.db.utils._get_managed_session_maker",
            new=mock_get_managed_session_maker,
        ),
        mock.patch(
            "mlflow.store.tracking.sqlalchemy_store.SqlAlchemyStore.search_experiments",
            return_value=[],
        ),
        mock.patch(
            "mlflow.store.tracking.sqlalchemy_store.SqlAlchemyStore._initialize_store_state",
            return_value=None,
        ),
    ):
        store = _get_store(artifact_uri=artifact_uri)
        assert isinstance(store, SqlAlchemyStore)
        assert store.db_uri == uri
        if is_windows():
            assert store.artifact_root_uri == Path.cwd().joinpath("artifact", "path").as_uri()
        else:
            assert store.artifact_root_uri == path_to_local_file_uri(
                Path.cwd().joinpath("artifact", "path")
            )

    mock_create_engine.assert_called_once_with(uri, pool_pre_ping=True)


def test_get_sqlalchemy_store_uses_server_artifact_root(tmp_path, monkeypatch):
    store_uri = f"sqlite:///{tmp_path.joinpath('backend_store.db')}"
    artifact_path = tmp_path / "server-artifacts"
    artifact_uri = path_to_local_file_uri(artifact_path)
    monkeypatch.setenv(ARTIFACT_ROOT_ENV_VAR, artifact_uri)

    with mock.patch("mlflow.store.tracking.sqlalchemy_store.SqlAlchemyStore") as mock_store:
        mlflow.tracking._tracking_service.utils._get_sqlalchemy_store(
            store_uri=store_uri, artifact_uri=None
        )

    mock_store.assert_called_once()
    assert mock_store.call_args.args[1] == artifact_uri
    monkeypatch.delenv(ARTIFACT_ROOT_ENV_VAR, raising=False)


def test_get_store_databricks(monkeypatch):
    for k, v in {
        MLFLOW_TRACKING_URI.name: "databricks",
        "DATABRICKS_HOST": "https://my-tracking-server",
        "DATABRICKS_TOKEN": "abcdef",
    }.items():
        monkeypatch.setenv(k, v)
    store = _get_store()
    assert isinstance(store, DatabricksTracingRestStore)
    assert store.get_host_creds().use_databricks_sdk
    assert _get_tracking_scheme() == "databricks"


def test_get_store_databricks_profile(monkeypatch):
    monkeypatch.setenv(MLFLOW_TRACKING_URI.name, "databricks://mycoolprofile")
    # It's kind of annoying to setup a profile, and we're not really trying to test
    # that anyway, so just check if we raise a relevant exception.
    store = _get_store()
    assert isinstance(store, DatabricksTracingRestStore)
    with pytest.raises(MlflowException, match="mycoolprofile"):
        store.get_host_creds()


def test_get_store_caches_on_store_uri_and_artifact_uri(tmp_path):
    registry = mlflow.tracking._tracking_service.utils._tracking_store_registry

    store_uri_1 = f"sqlite:///{tmp_path.joinpath('backend_store_1.db')}"
    store_uri_2 = f"sqlite:///{tmp_path.joinpath('backend_store_2.db')}"
    stores_uris = [store_uri_1, store_uri_2]
    artifact_uris = [
        None,
        str(tmp_path.joinpath("artifact_root_1")),
        str(tmp_path.joinpath("artifact_root_2")),
    ]

    stores = []
    for args in itertools.product(stores_uris, artifact_uris):
        store1 = registry.get_store(*args)
        store2 = registry.get_store(*args)
        assert store1 is store2
        stores.append(store1)

    assert all(s1 is not s2 for s1, s2 in itertools.combinations(stores, 2))


def test_standard_store_registry_with_mocked_entrypoint():
    mock_entrypoint = mock.Mock()
    mock_entrypoint.name = "mock-scheme"

    with mock.patch("mlflow.utils.plugins._get_entry_points", return_value=[mock_entrypoint]):
        # Entrypoints are registered at import time, so we need to reload the
        # module to register the entrypoint given by the mocked
        # entrypoints.get_group_all
        reload(mlflow.tracking._tracking_service.utils)

        expected_standard_registry = {
            "",
            "file",
            "http",
            "https",
            "postgresql",
            "mysql",
            "sqlite",
            "mssql",
            "databricks",
            "mock-scheme",
        }
        assert expected_standard_registry.issubset(
            mlflow.tracking._tracking_service.utils._tracking_store_registry._registry.keys()
        )


@pytest.mark.skip(reason="FileStore is no longer supported")
def test_standard_store_registry_with_installed_plugin(tmp_path, monkeypatch):
    monkeypatch.chdir(tmp_path)
    reload(mlflow.tracking._tracking_service.utils)
    assert (
        "file-plugin" in mlflow.tracking._tracking_service.utils._tracking_store_registry._registry
    )

    from mlflow_test_plugin.file_store import PluginFileStore

    monkeypatch.setenv(MLFLOW_TRACKING_URI.name, "file-plugin:test-path")
    plugin_file_store = mlflow.tracking._tracking_service.utils._get_store()
    assert isinstance(plugin_file_store, PluginFileStore)
    assert plugin_file_store.is_plugin
    assert _get_tracking_scheme() == "custom_scheme"


def test_plugin_registration():
    tracking_store = TrackingStoreRegistry()

    test_uri = "mock-scheme://fake-host/fake-path"
    test_scheme = "mock-scheme"

    mock_plugin = mock.Mock()
    tracking_store.register(test_scheme, mock_plugin)
    assert test_scheme in tracking_store._registry
    assert tracking_store.get_store(test_uri) == mock_plugin.return_value
    mock_plugin.assert_called_once_with(store_uri=test_uri, artifact_uri=None)


def test_plugin_registration_via_entrypoints():
    mock_plugin_function = mock.Mock()
    mock_entrypoint = mock.Mock(load=mock.Mock(return_value=mock_plugin_function))
    mock_entrypoint.name = "mock-scheme"

    with mock.patch(
        "mlflow.utils.plugins._get_entry_points", return_value=[mock_entrypoint]
    ) as mock_get_group_all:
        tracking_store = TrackingStoreRegistry()
        tracking_store.register_entrypoints()

    assert tracking_store.get_store("mock-scheme://") == mock_plugin_function.return_value

    mock_plugin_function.assert_called_once_with(store_uri="mock-scheme://", artifact_uri=None)
    mock_get_group_all.assert_called_once_with("mlflow.tracking_store")


@pytest.mark.parametrize(
    "exception", [AttributeError("test exception"), ImportError("test exception")]
)
def test_handle_plugin_registration_failure_via_entrypoints(exception):
    mock_entrypoint = mock.Mock(load=mock.Mock(side_effect=exception))
    mock_entrypoint.name = "mock-scheme"

    with mock.patch(
        "mlflow.utils.plugins._get_entry_points", return_value=[mock_entrypoint]
    ) as mock_get_group_all:
        tracking_store = TrackingStoreRegistry()

        # Check that the raised warning contains the message from the original exception
        with pytest.warns(UserWarning, match="test exception"):
            tracking_store.register_entrypoints()

    mock_entrypoint.load.assert_called_once()
    mock_get_group_all.assert_called_once_with("mlflow.tracking_store")


def test_get_store_for_unregistered_scheme():
    tracking_store = TrackingStoreRegistry()

    with pytest.raises(
        UnsupportedModelRegistryStoreURIException,
        match="Model registry functionality is unavailable",
    ):
        tracking_store.get_store("unknown-scheme://")


def test_resolve_tracking_uri_with_param():
    with mock.patch(
        "mlflow.tracking._tracking_service.utils.get_tracking_uri",
        return_value="databricks://tracking_qoeirj",
    ):
        overriding_uri = "databricks://tracking_poiwerow"
        assert _resolve_tracking_uri(overriding_uri) == overriding_uri


def test_resolve_tracking_uri_with_no_param():
    with mock.patch(
        "mlflow.tracking._tracking_service.utils.get_tracking_uri",
        return_value="databricks://tracking_zlkjdas",
    ):
        assert _resolve_tracking_uri() == "databricks://tracking_zlkjdas"


@pytest.mark.skip(reason="FileStore is no longer supported")
def test_store_object_can_be_serialized_by_pickle(tmp_path):
    """
    This test ensures a store object generated by `_get_store` can be serialized by pickle
    to prevent issues such as https://github.com/mlflow/mlflow/issues/2954
    """
    pickle.dump(_get_store(f"file:///{tmp_path.joinpath('mlflow')}"), io.BytesIO())
    pickle.dump(_get_store("databricks"), io.BytesIO())
    pickle.dump(_get_store("https://example.com"), io.BytesIO())
    # pickle.dump(_get_store(f"sqlite:///{tmpdir.strpath}/mlflow.db"), io.BytesIO())
    # This throws `AttributeError: Can't pickle local object 'create_engine.<locals>.connect'`


@pytest.mark.parametrize("absolute", [True, False], ids=["absolute", "relative"])
def test_set_tracking_uri_with_path(tmp_path, monkeypatch, absolute):
    monkeypatch.chdir(tmp_path)
    path = Path("foo/bar")
    if absolute:
        path = tmp_path / path
    with mock.patch("mlflow.tracking._tracking_service.utils._tracking_uri", None):
        set_tracking_uri(path)
        assert get_tracking_uri() == path.absolute().resolve().as_uri()


def test_set_tracking_uri_update_trace_provider(tmp_path):
    default_uri = mlflow.get_tracking_uri()
    sqlite_uri = f"sqlite:///{tmp_path / 'mlflow.db'}"
    try:
        assert get_tracer_tracking_uri() != sqlite_uri

        set_tracking_uri(sqlite_uri)
        assert get_tracer_tracking_uri() == sqlite_uri

        set_tracking_uri("https://foo")
        assert get_tracer_tracking_uri() == "https://foo"
    finally:
        # clean up
        set_tracking_uri(default_uri)


@pytest.mark.parametrize("store_uri", ["databricks-uc", "databricks-uc://profile"])
def test_get_store_raises_on_uc_uri(store_uri):
    set_tracking_uri(store_uri)
    with pytest.raises(
        MlflowException,
        match="Setting the tracking URI to a Unity Catalog backend is not "
        "supported in the current version of the MLflow client",
    ):
        mlflow.tracking.MlflowClient()
    assert _get_tracking_scheme() == "databricks-uc"


@pytest.mark.parametrize("tracking_uri", ["file:///tmp/mlruns", "sqlite:///tmp/mlruns.db", ""])
def test_set_get_tracking_uri_consistency(tracking_uri):
    mlflow.set_tracking_uri(tracking_uri)
    assert mlflow.get_tracking_uri() == tracking_uri


def test_get_tracking_scheme():
    assert _get_tracking_scheme("uc://profile@databricks") == "uc"
    # no builder registered for custom scheme
    assert _get_tracking_scheme("custom-scheme://") == "None"


@pytest.mark.parametrize(
    ("scheme", "uri", "expected"),
    [
        ("arn", "arn:aws:sagemaker:us-east-1:123456789:mlflow-tracking-server/my-server", "aws"),
        ("arn", "arn:aws:sagemaker:eu-west-1:987654321:mlflow-tracking-server/test", "aws"),
        ("azureml", "azureml://eastus.api.azureml.ms/mlflow/v2.0/subscriptions/123", "azure"),
        ("azureml", "azureml://workspace", "azure"),
        ("some-plugin", "some-plugin://host/path", "custom_scheme"),
    ],
)
def test_resolve_custom_scheme(scheme, uri, expected):
    assert _resolve_custom_scheme(scheme, uri) == expected