"""Tests for dgtw_mlflow_utils.uploader.

================================
Covers: UploadConfig dataclass and ModelUploader class.
All MLflow interactions are mocked.
"""

import os
from unittest.mock import MagicMock, patch

import pytest
from dgtw_mlflow_utils.naming import Architecture, ModelVariant, RunAction, Stage
from dgtw_mlflow_utils.uploader import ModelUploader, UploadConfig

# ======================================================================
# UploadConfig tests
# ======================================================================


class TestUploadConfig:
    """Tests for the UploadConfig dataclass."""

    def test_default_values(self):
        config = UploadConfig()
        assert config.architecture == Architecture.TWO_TOWER
        assert config.variant == ModelVariant.BASELINE
        assert config.component == "Recommendation"
        assert config.stage == Stage.RESEARCH
        assert config.use_case == "Model"
        assert config.action == RunAction.UPLOAD
        assert config.tracking_uri is None
        assert config.tags == {}
        assert config.params == {}
        assert config.register_model is True
        assert config.artifact_path == "model"
        assert config.run_description == ""

    def test_experiment_name_generated(self):
        config = UploadConfig(component="Ranking", stage=Stage.PRODUCTION)
        assert config.experiment_name == "TAT_Ranking_Production"

    def test_registered_model_name_generated(self):
        config = UploadConfig(
            architecture=Architecture.BERT4REC,
            variant=ModelVariant.CHAMPION,
        )
        assert config.registered_model_name == "TAT_Bert4Rec_Champion"

    def test_run_name_generated(self):
        config = UploadConfig(action=RunAction.TRAIN, use_case="MiniLM")
        assert config.run_name.startswith("Train_MiniLM_")

    def test_custom_string_architecture(self):
        config = UploadConfig(architecture="CustomArch", variant="v2")
        assert config.registered_model_name == "TAT_CustomArch_v2"

    def test_effective_tracking_uri_explicit(self):
        config = UploadConfig(tracking_uri="http://my-server:5000")
        assert config.effective_tracking_uri == "http://my-server:5000"

    def test_effective_tracking_uri_env(self):
        config = UploadConfig()
        with patch.dict(os.environ, {"MLFLOW_TRACKING_URI": "http://env-server:5000"}):
            assert config.effective_tracking_uri == "http://env-server:5000"

    def test_effective_tracking_uri_default(self):
        config = UploadConfig()
        with patch.dict(os.environ, {}, clear=True):
            # Remove MLFLOW_TRACKING_URI if it exists
            os.environ.pop("MLFLOW_TRACKING_URI", None)
            assert config.effective_tracking_uri == "http://127.0.0.1:5001"

    def test_tags_and_params(self):
        config = UploadConfig(
            tags={"team": "ai"},
            params={"lr": 0.001},
        )
        assert config.tags == {"team": "ai"}
        assert config.params == {"lr": 0.001}


# ======================================================================
# ModelUploader tests
# ======================================================================


class TestModelUploader:
    """Tests for ModelUploader with mocked MLflow."""

    @patch("dgtw_mlflow_utils.uploader.mlflow")
    def test_init_sets_tracking_uri(self, mock_mlflow):
        config = UploadConfig(tracking_uri="http://test:5000")
        ModelUploader(config)

        mock_mlflow.set_tracking_uri.assert_called_once_with("http://test:5000")
        mock_mlflow.set_experiment.assert_called_once_with(config.experiment_name)

    @patch("dgtw_mlflow_utils.uploader.mlflow")
    def test_summary(self, mock_mlflow):
        config = UploadConfig(
            architecture=Architecture.TWO_TOWER,
            variant=ModelVariant.BASELINE,
            tracking_uri="http://test:5000",
        )
        uploader = ModelUploader(config)
        summary = uploader.summary()

        assert "http://test:5000" in summary
        assert "TAT_Recommendation_Research" in summary
        assert "TAT_TwoTower_Baseline" in summary
        assert "MLflow Upload Configuration" in summary

    @patch("dgtw_mlflow_utils.uploader.mlflow")
    def test_upload_pyfunc(self, mock_mlflow):
        config = UploadConfig(tracking_uri="http://test:5000")
        uploader = ModelUploader(config)

        # Mock the run context
        mock_run = MagicMock()
        mock_run.info.run_id = "run-123"
        mock_mlflow.start_run.return_value.__enter__ = MagicMock(return_value=mock_run)
        mock_mlflow.start_run.return_value.__exit__ = MagicMock(return_value=False)

        # Mock pyfunc log_model
        mock_model_info = MagicMock()
        mock_model_info.model_uri = "runs:/run-123/model"
        mock_mlflow.pyfunc.log_model.return_value = mock_model_info

        python_model = MagicMock()
        result = uploader.upload_pyfunc(python_model=python_model)

        assert result["run_id"] == "run-123"
        assert result["model_uri"] == "runs:/run-123/model"
        assert result["registered_model_name"] == config.registered_model_name
        mock_mlflow.pyfunc.log_model.assert_called_once()

    @patch("dgtw_mlflow_utils.uploader.mlflow")
    def test_upload_pyfunc_no_register(self, mock_mlflow):
        config = UploadConfig(tracking_uri="http://test:5000", register_model=False)
        uploader = ModelUploader(config)

        mock_run = MagicMock()
        mock_run.info.run_id = "run-456"
        mock_mlflow.start_run.return_value.__enter__ = MagicMock(return_value=mock_run)
        mock_mlflow.start_run.return_value.__exit__ = MagicMock(return_value=False)

        mock_model_info = MagicMock()
        mock_model_info.model_uri = "runs:/run-456/model"
        mock_mlflow.pyfunc.log_model.return_value = mock_model_info

        python_model = MagicMock()
        result = uploader.upload_pyfunc(python_model=python_model)

        assert result["registered_model_name"] is None

    @patch("dgtw_mlflow_utils.uploader.mlflow")
    def test_upload_pyfunc_with_all_options(self, mock_mlflow):
        config = UploadConfig(
            tracking_uri="http://test:5000",
            tags={"team": "ai"},
            params={"lr": 0.001},
        )
        uploader = ModelUploader(config)

        mock_run = MagicMock()
        mock_run.info.run_id = "run-789"
        mock_mlflow.start_run.return_value.__enter__ = MagicMock(return_value=mock_run)
        mock_mlflow.start_run.return_value.__exit__ = MagicMock(return_value=False)

        mock_model_info = MagicMock()
        mock_model_info.model_uri = "runs:/run-789/model"
        mock_mlflow.pyfunc.log_model.return_value = mock_model_info

        python_model = MagicMock()
        result = uploader.upload_pyfunc(
            python_model=python_model,
            artifacts={"weights": "/path/to/model.pt"},
            signature=MagicMock(),
            input_example={"text": "hello"},
            code_paths=["src/model.py"],
            pip_requirements=["torch>=1.12"],
        )

        call_kwargs = mock_mlflow.pyfunc.log_model.call_args
        assert "artifacts" in call_kwargs.kwargs or call_kwargs[1].get("artifacts")
        assert result["run_id"] == "run-789"

    @patch("dgtw_mlflow_utils.uploader.mlflow")
    def test_log_common_tags(self, mock_mlflow):
        config = UploadConfig(
            tracking_uri="http://test:5000",
            tags={"custom_tag": "value"},
            params={"lr": 0.001},
        )
        uploader = ModelUploader(config)

        mock_run = MagicMock()
        mock_run.info.run_id = "run-tags"
        mock_mlflow.start_run.return_value.__enter__ = MagicMock(return_value=mock_run)
        mock_mlflow.start_run.return_value.__exit__ = MagicMock(return_value=False)

        mock_model_info = MagicMock()
        mock_model_info.model_uri = "runs:/run-tags/model"
        mock_mlflow.pyfunc.log_model.return_value = mock_model_info

        uploader.upload_pyfunc(python_model=MagicMock())

        # Verify set_tags was called with default + custom tags
        mock_mlflow.set_tags.assert_called_once()
        tags_arg = mock_mlflow.set_tags.call_args[0][0]
        assert tags_arg["project"] == "TAT-Personalization"
        assert tags_arg["custom_tag"] == "value"
        assert "architecture" in tags_arg
        assert "variant" in tags_arg

        # Verify log_params was called
        mock_mlflow.log_params.assert_called_once_with({"lr": 0.001})

    @patch("dgtw_mlflow_utils.uploader.mlflow")
    def test_upload_artifacts(self, mock_mlflow):
        config = UploadConfig(tracking_uri="http://test:5000")
        uploader = ModelUploader(config)

        mock_run = MagicMock()
        mock_run.info.run_id = "run-art"
        mock_mlflow.start_run.return_value.__enter__ = MagicMock(return_value=mock_run)
        mock_mlflow.start_run.return_value.__exit__ = MagicMock(return_value=False)

        # Use an existing file (this test file itself) instead of creating a temp file
        # to avoid disk space issues
        this_file = os.path.abspath(__file__)

        result = uploader.upload_artifacts(
            files=[this_file],
            artifact_subdir="test_dir",
            extra_metrics={"score": 0.95},
        )
        assert result["run_id"] == "run-art"
        assert len(result["artifact_uris"]) == 1
        mock_mlflow.log_artifact.assert_called_once()
        mock_mlflow.log_metrics.assert_called_once_with({"score": 0.95})

    @patch("dgtw_mlflow_utils.uploader.mlflow")
    def test_upload_artifacts_file_not_found(self, mock_mlflow):
        config = UploadConfig(tracking_uri="http://test:5000")
        uploader = ModelUploader(config)

        mock_run = MagicMock()
        mock_run.info.run_id = "run-err"
        mock_mlflow.start_run.return_value.__enter__ = MagicMock(return_value=mock_run)
        mock_mlflow.start_run.return_value.__exit__ = MagicMock(return_value=False)

        with pytest.raises(FileNotFoundError, match="Artifact not found"):
            uploader.upload_artifacts(files=["/nonexistent/file.txt"])
