"""Tests for dgtw_mlflow_utils.registry.

=================================
Covers: ModelRegistry class — all read, write, and load operations.
All MLflow interactions are mocked.
"""

from types import SimpleNamespace
from unittest.mock import MagicMock, patch

import pytest
from dgtw_mlflow_utils.registry import (
    ALIAS_CANARY,
    ALIAS_CHALLENGER,
    ALIAS_CHAMPION,
    ModelRegistry,
)


def _make_model_version(version: str, run_id: str = "run-abc", last_updated_timestamp: int = 1700000000000):
    """Helper to create a fake ModelVersion-like object."""
    mv = SimpleNamespace(
        version=version,
        run_id=run_id,
        last_updated_timestamp=last_updated_timestamp,
    )
    return mv


# ======================================================================
# Alias Constants
# ======================================================================


class TestAliasConstants:
    def test_champion(self):
        assert ALIAS_CHAMPION == "champion"

    def test_challenger(self):
        assert ALIAS_CHALLENGER == "challenger"

    def test_canary(self):
        assert ALIAS_CANARY == "canary"


# ======================================================================
# ModelRegistry tests
# ======================================================================


class TestModelRegistry:
    """Tests for ModelRegistry with mocked MlflowClient."""

    @patch("dgtw_mlflow_utils.registry.MlflowClient")
    @patch("dgtw_mlflow_utils.registry.mlflow")
    def _create_registry(self, mock_mlflow, mock_client_cls, tracking_uri="http://test:5000"):
        """Helper to create a ModelRegistry with mocked deps."""
        mock_client = MagicMock()
        mock_client_cls.return_value = mock_client
        registry = ModelRegistry(tracking_uri=tracking_uri)
        return registry, mock_client, mock_mlflow

    # --- Init ---

    def test_init_sets_tracking_uri(self):
        _registry, _, mock_mlflow = self._create_registry()
        mock_mlflow.set_tracking_uri.assert_called_once_with("http://test:5000")

    def test_class_level_alias_constants(self):
        assert ModelRegistry.ALIAS_CHAMPION == "champion"
        assert ModelRegistry.ALIAS_CHALLENGER == "challenger"
        assert ModelRegistry.ALIAS_CANARY == "canary"

    # --- list_versions ---

    def test_list_versions_sorted(self):
        registry, mock_client, _ = self._create_registry()

        mock_client.search_model_versions.return_value = [
            _make_model_version("1"),
            _make_model_version("3"),
            _make_model_version("2"),
        ]

        versions = registry.list_versions("TestModel")
        assert [v.version for v in versions] == ["3", "2", "1"]
        mock_client.search_model_versions.assert_called_once_with("name='TestModel'")

    def test_list_versions_empty(self):
        registry, mock_client, _ = self._create_registry()
        mock_client.search_model_versions.return_value = []

        versions = registry.list_versions("TestModel")
        assert versions == []

    # --- get_latest_version ---

    def test_get_latest_version(self):
        registry, mock_client, _ = self._create_registry()
        mock_client.search_model_versions.return_value = [
            _make_model_version("1"),
            _make_model_version("5"),
            _make_model_version("3"),
        ]

        latest = registry.get_latest_version("TestModel")
        assert latest.version == "5"

    def test_get_latest_version_none(self):
        registry, mock_client, _ = self._create_registry()
        mock_client.search_model_versions.return_value = []

        latest = registry.get_latest_version("TestModel")
        assert latest is None

    # --- get_version_by_alias ---

    def test_get_version_by_alias_found(self):
        registry, mock_client, _ = self._create_registry()
        mock_client.get_model_version_by_alias.return_value = _make_model_version("3")

        mv = registry.get_version_by_alias("TestModel", "champion")
        assert mv.version == "3"
        mock_client.get_model_version_by_alias.assert_called_once_with("TestModel", "champion")

    def test_get_version_by_alias_not_found(self):
        registry, mock_client, _mock_mlflow = self._create_registry()

        import mlflow.exceptions

        mock_client.get_model_version_by_alias.side_effect = mlflow.exceptions.MlflowException("not found")

        mv = registry.get_version_by_alias("TestModel", "champion")
        assert mv is None

    # --- promote ---

    def test_promote_with_explicit_version(self):
        registry, mock_client, _ = self._create_registry()
        mock_client.get_model_version.return_value = _make_model_version("3")

        mv = registry.promote("TestModel", alias="champion", version="3")
        assert mv.version == "3"
        mock_client.set_registered_model_alias.assert_called_once_with("TestModel", "champion", "3")

    def test_promote_uses_latest_version(self):
        registry, mock_client, _ = self._create_registry()
        mock_client.search_model_versions.return_value = [
            _make_model_version("1"),
            _make_model_version("5"),
        ]
        mock_client.get_model_version.return_value = _make_model_version("5")

        registry.promote("TestModel")
        mock_client.set_registered_model_alias.assert_called_once_with("TestModel", "champion", "5")

    def test_promote_no_versions_raises(self):
        registry, mock_client, _ = self._create_registry()
        mock_client.search_model_versions.return_value = []

        with pytest.raises(RuntimeError, match="No versions found"):
            registry.promote("TestModel")

    def test_promote_default_alias_is_champion(self):
        registry, mock_client, _ = self._create_registry()
        mock_client.search_model_versions.return_value = [
            _make_model_version("1"),
        ]
        mock_client.get_model_version.return_value = _make_model_version("1")

        registry.promote("TestModel")
        mock_client.set_registered_model_alias.assert_called_once_with("TestModel", "champion", "1")

    # --- set_tag ---

    def test_set_tag(self):
        registry, mock_client, _ = self._create_registry()
        registry.set_tag("TestModel", "3", "approved", "yes")
        mock_client.set_model_version_tag.assert_called_once_with("TestModel", "3", "approved", "yes")

    # --- delete_version ---

    def test_delete_version(self):
        registry, mock_client, _ = self._create_registry()
        registry.delete_version("TestModel", "2")
        mock_client.delete_model_version.assert_called_once_with("TestModel", "2")

    # --- cleanup_old_versions ---

    def test_cleanup_dry_run(self):
        registry, mock_client, _ = self._create_registry()
        mock_client.search_model_versions.return_value = [
            _make_model_version("1"),
            _make_model_version("2"),
            _make_model_version("3"),
            _make_model_version("4"),
            _make_model_version("5"),
        ]

        deleted = registry.cleanup_old_versions("TestModel", keep=2, dry_run=True)
        # Sorted desc: 5, 4, 3, 2, 1 → keep 5, 4 → delete 3, 2, 1
        assert len(deleted) == 3
        mock_client.delete_model_version.assert_not_called()

    def test_cleanup_actual_delete(self):
        registry, mock_client, _ = self._create_registry()
        mock_client.search_model_versions.return_value = [
            _make_model_version("1"),
            _make_model_version("2"),
            _make_model_version("3"),
        ]

        deleted = registry.cleanup_old_versions("TestModel", keep=1, dry_run=False)
        # Sorted desc: 3, 2, 1 → keep 3 → delete 2, 1
        assert len(deleted) == 2
        assert mock_client.delete_model_version.call_count == 2

    def test_cleanup_nothing_to_delete(self):
        registry, mock_client, _ = self._create_registry()
        mock_client.search_model_versions.return_value = [
            _make_model_version("1"),
        ]

        deleted = registry.cleanup_old_versions("TestModel", keep=5)
        assert deleted == []

    # --- load_champion ---

    @patch("mlflow.pyfunc.load_model")
    def test_load_champion(self, mock_load_model):
        registry, _, _ = self._create_registry()
        mock_load_model.return_value = MagicMock()

        registry.load_champion("TAT_TwoTower_Baseline")
        mock_load_model.assert_called_once_with("models:/TAT_TwoTower_Baseline@champion")

    # --- load_version ---

    @patch("mlflow.pyfunc.load_model")
    def test_load_version(self, mock_load_model):
        registry, _, _ = self._create_registry()
        mock_load_model.return_value = MagicMock()

        registry.load_version("TAT_TwoTower_Baseline", "3")
        mock_load_model.assert_called_once_with("models:/TAT_TwoTower_Baseline/3")

    # --- print_summary ---

    def test_print_summary(self, capsys):
        registry, mock_client, _ = self._create_registry()
        mock_client.search_model_versions.return_value = [
            _make_model_version("1", run_id="run-abc", last_updated_timestamp=1700000000000),
            _make_model_version("2", run_id="run-def", last_updated_timestamp=1700100000000),
        ]

        registry.print_summary("TestModel")
        captured = capsys.readouterr()
        assert "Model Registry Summary: TestModel" in captured.out
        assert "run-abc" in captured.out or "run-def" in captured.out
