import os
from mlflow_migration.common.source_tags import ExportTags
from mlflow_migration.common import MlflowExportImportException
from mlflow_migration.common.model_utils import model_names_same_registry
from mlflow_migration.model.export_model import export_model
from mlflow_migration.model.import_model import import_model
from mlflow_migration.model.import_model import _extract_model_path, _path_join

from tests.open_source.oss_utils_test import create_simple_run, create_version
from tests.open_source.oss_utils_test import (
    mk_test_object_name_default,
    mk_dst_model_name,
)
from tests.compare_utils import (
    compare_models_with_versions,
    compare_models,
    compare_versions,
)
from tests.open_source.init_tests import mlflow_context


# == Test stages


def test_export_import_model_1_stage(mlflow_context):
    model_src, model_dst = _run_test_export_import_model_stages(
        mlflow_context, stages=["Production"]
    )
    assert len(model_dst.latest_versions) == 1
    compare_models_with_versions(mlflow_context, model_src, model_dst)


def test_export_import_model_2_stages(mlflow_context):
    model_src, model_dst = _run_test_export_import_model_stages(
        mlflow_context, stages=["Production", "Staging"]
    )
    assert len(model_dst.latest_versions) == 2
    compare_models_with_versions(mlflow_context, model_src, model_dst)


def test_export_import_model_all_stages(mlflow_context):
    model_src, model_dst = _run_test_export_import_model_stages(
        mlflow_context, stages=None
    )
    assert len(model_dst.latest_versions) == 4
    compare_models_with_versions(mlflow_context, model_src, model_dst)


# == Test stages and versions


def test_export_import_model_both_stages(mlflow_context):
    try:
        _run_test_export_import_model_stages(
            mlflow_context, stages=["Production"], versions=[1]
        )
    except MlflowExportImportException:
        # "Both stages {self.stages} and versions {self.versions} cannot be set")
        pass


# == Test versions


def _get_version_ids(model):
    return [vr.version for vr in model.latest_versions]


def _get_version_ids_dst(model):
    return [
        vr.tags[f"{ExportTags.PREFIX_FIELD}.version"] for vr in model.latest_versions
    ]


def test_export_import_model_first_two_versions(mlflow_context):
    model_src, model_dst = _run_test_export_import_model_stages(
        mlflow_context, versions=["1", "2"]
    )
    assert len(model_dst.latest_versions) == 2
    compare_models_with_versions(mlflow_context, model_src, model_dst)
    ids_src = _get_version_ids(model_src)
    ids_dst = _get_version_ids(model_dst)
    for j, id in enumerate(ids_dst):
        assert id == ids_src[j]


def test_export_import_model_two_from_middle_versions(mlflow_context):
    model_src, model_dst = _run_test_export_import_model_stages(
        mlflow_context, versions=["2", "3", "4"]
    )
    assert len(model_dst.latest_versions) == 3
    ids_src = _get_version_ids(model_src)
    ids_dst = _get_version_ids_dst(model_dst)
    assert set(ids_dst).issubset(set(ids_src))

    compare_models(
        model_src, model_dst, mlflow_context.client_src != mlflow_context.client_dst
    )
    for vr_dst in model_dst.latest_versions:
        vr_src_id = vr_dst.tags[f"{ExportTags.PREFIX_FIELD}.version"]
        vr_src = [vr for vr in model_src.latest_versions if vr.version == vr_src_id]
        assert (len(vr_src)) == 1
        vr_src = vr_src[0]
        assert vr_src.version == vr_src_id
        compare_versions(mlflow_context, vr_src, vr_dst)


# == Test export deleted runs


def _run_test_deleted_runs(mlflow_context, delete_run, export_deleted_runs):
    model_name_src = mk_test_object_name_default()
    mlflow_context.client_src.create_registered_model(model_name_src)

    create_version(mlflow_context.client_src, model_name_src, "Production")
    _, run2 = create_version(mlflow_context.client_src, model_name_src, "Staging")
    versions = mlflow_context.client_src.search_model_versions(
        filter_string=f"name='{model_name_src}'"
    )
    if delete_run:
        mlflow_context.client_src.delete_run(run2.info.run_id)
    run2 = mlflow_context.client_src.get_run(run2.info.run_id)

    export_model(
        model_name=model_name_src,
        output_dir=mlflow_context.output_dir,
        export_deleted_runs=export_deleted_runs,
        mlflow_client=mlflow_context.client_src,
    )
    model_name_dst = mk_dst_model_name(model_name_src)
    import_model(
        model_name=model_name_dst,
        experiment_name=model_name_dst,
        input_dir=mlflow_context.output_dir,
        mlflow_client=mlflow_context.client_dst,
    )
    versions = mlflow_context.client_dst.search_model_versions(
        filter_string=f"name='{model_name_dst}'"
    )
    run_lifecycle_stages = [
        mlflow_context.client_dst.get_run(vr.run_id).info.lifecycle_stage
        for vr in versions
    ]
    return versions, run_lifecycle_stages


def test_export_delete_Yes_delete_run_No(mlflow_context):
    "Test: export-deleted-runs==Yes and no run deleted"
    versions, run_lifecycle_stages = _run_test_deleted_runs(mlflow_context, False, True)
    assert len(versions) == 2
    assert run_lifecycle_stages == ["active", "active"]


def test_export_delete_Yes_delete_run_Yes(mlflow_context):
    "Test: export-deleted-runs==True and run deleted"
    versions, run_lifecycle_stages = _run_test_deleted_runs(mlflow_context, True, True)
    assert len(versions) == 2
    assert run_lifecycle_stages == ["deleted", "active"]


def test_export_delete_No_delete_run_Yes(mlflow_context):
    "Test: export-deleted-runs==False and no run deleted"
    versions, run_lifecycle_stages = _run_test_deleted_runs(mlflow_context, True, False)
    assert len(versions) == 1
    assert run_lifecycle_stages == ["active"]


# == Internal


def _run_test_export_import_model_stages(mlflow_context, stages=None, versions=None):
    model_name_src = mk_test_object_name_default()
    desc = "Hello decription"
    tags = {"city": "franconia"}
    model_src = mlflow_context.client_src.create_registered_model(
        model_name_src, tags, desc
    )

    create_version(mlflow_context.client_src, model_name_src, "Production")
    create_version(mlflow_context.client_src, model_name_src, "Staging")
    create_version(mlflow_context.client_src, model_name_src, "Archived")
    create_version(mlflow_context.client_src, model_name_src, "None")

    model_src = mlflow_context.client_src.get_registered_model(model_name_src)
    export_model(
        model_name=model_name_src,
        output_dir=mlflow_context.output_dir,
        stages=stages,
        versions=versions,
        mlflow_client=mlflow_context.client_src,
    )

    model_name_dst = mk_dst_model_name(model_name_src)
    import_model(
        model_name=model_name_dst,
        experiment_name=model_name_dst,
        input_dir=mlflow_context.output_dir,
        import_source_tags=True,
        delete_model=True,
        mlflow_client=mlflow_context.client_dst,
    )

    model_dst = mlflow_context.client_dst.get_registered_model(model_name_dst)
    return model_src, model_dst


def _create_run(client):
    _, run = create_simple_run(client)
    return client.get_run(run.info.run_id)


# == Test parsing for _extract_model_path to extract from version `source` field

_exp_id = "1812"
_run_id = "48cf29167ddb4e098da780f0959fb4cf"
_local_path_base = os.path.join("dbfs:/databricks/mlflow-tracking", _exp_id, _run_id)


def test_extract_no_artifacts():
    source = os.path.join(_local_path_base)
    model_path = _extract_model_path(source, _run_id)
    assert model_path == ""


def test_extract_just_artifacts():
    source = os.path.join(_local_path_base, "artifacts")
    model_path = _extract_model_path(source, _run_id)
    assert model_path == ""


def test_extract_just_artifacts_slash():
    source = os.path.join(_local_path_base, "artifacts/")
    model_path = _extract_model_path(source, _run_id)
    assert model_path == ""


def test_extract_model():
    source = os.path.join(_local_path_base, "artifacts", "model")
    model_path = _extract_model_path(source, _run_id)
    assert model_path == "model"


def test_extract_model_sklearn():
    source = os.path.join(_local_path_base, "artifacts", "model/sklearn")
    model_path = _extract_model_path(source, _run_id)
    assert model_path == "model/sklearn"


def test_extract_no_run_id():
    source = os.path.join(_local_path_base, "artifacts")
    try:
        _extract_model_path(source, "1215")
        assert False
    except MlflowExportImportException:
        pass


# == Test that both model names are either UC model names or WS model names


def test_both_model_names_are_uc():
    assert model_names_same_registry("model1", "model2")


def test_both_model_names_are_ws():
    assert model_names_same_registry("andre.models.model1", "andre.models.model2")


def test_both_model_names_are_not_same_1():
    assert not model_names_same_registry("model1", "andre.models.model2")


def test_both_model_names_are_not_same_2():
    assert not model_names_same_registry("andre.models.model2", "model1")


# == Test for DOS path adjustment

_base_dir = "dbfs:/mlflow/1812"
_expected_path = "dbfs:/mlflow/1812/model"


def test_path_join_frontslash():
    res = _path_join(_base_dir, "model")
    assert res == os.path.join(_expected_path)


def test_path_join_backslash():
    dir = _base_dir.replace("/", "\\")
    res = _path_join(dir, "model")
    assert res == os.path.join(_expected_path)
