"""Tests for dgtw_mlflow_utils.evaluator.
==================================
Covers: TfMLflowEpochLogger callback and evaluate_model_tf function.
MLflow interactions are mocked; TensorFlow is used for model operations.
"""  # noqa: D205

from unittest.mock import MagicMock, patch

import numpy as np
import pytest
import tensorflow as tf
from dgtw_mlflow_utils.evaluator import TfMLflowEpochLogger, evaluate_model_tf

# ======================================================================
# Helper: simple model that produces (batch, batch) logits
# ======================================================================


class DummyTwoTower(tf.keras.Model):
    """Minimal model producing (batch_size, batch_size) logits."""

    def __init__(self):
        """Initialize the DummyTwoTower model."""
        super().__init__()
        self.dense = tf.keras.layers.Dense(16)

    def call(self, inputs, training=False):
        if isinstance(inputs, dict):
            q = tf.cast(inputs["query_id"], tf.float32)
            i = tf.cast(inputs["item_id"], tf.float32)
        else:
            q, i = inputs
            q = tf.cast(q, tf.float32)
            i = tf.cast(i, tf.float32)
        q_emb = self.dense(tf.expand_dims(q, -1))
        i_emb = self.dense(tf.expand_dims(i, -1))
        return tf.matmul(q_emb, i_emb, transpose_b=True)


def _make_dataset(num_samples=32, batch_size=8):
    """Create a synthetic tf.data.Dataset."""
    ds = tf.data.Dataset.from_tensor_slices(
        {
            "query_id": np.random.randint(0, 50, size=num_samples),
            "item_id": np.random.randint(0, 50, size=num_samples),
        }
    )
    return ds.batch(batch_size, drop_remainder=True)


# ======================================================================
# TfMLflowEpochLogger tests
# ======================================================================


class TestTfMLflowEpochLogger:
    """Tests for the TfMLflowEpochLogger Keras callback."""

    def test_init_defaults(self):
        callback = TfMLflowEpochLogger()
        assert callback.enable_mlflow is False
        assert callback.experiment_name is None
        assert callback.run_name == "Training_Run"
        assert callback.log_params == {}

    def test_init_with_params(self):
        callback = TfMLflowEpochLogger(
            enable_mlflow=True,
            experiment_name="TestExp",
            run_name="TestRun",
            log_params={"lr": 0.001, "epochs": 10},
        )
        assert callback.enable_mlflow is True
        assert callback.experiment_name == "TestExp"
        assert callback.run_name == "TestRun"
        assert callback.log_params == {"lr": 0.001, "epochs": 10}

    # --- on_train_begin ---

    @patch("dgtw_mlflow_utils.evaluator.mlflow")
    def test_on_train_begin_mlflow_enabled(self, mock_mlflow):
        mock_mlflow.active_run.return_value = None
        mock_mlflow.get_experiment_by_name.return_value = None

        callback = TfMLflowEpochLogger(
            enable_mlflow=True,
            experiment_name="TestExp",
            run_name="TestRun",
            log_params={"lr": 0.001, "epochs": 10},
        )
        callback.on_train_begin()

        mock_mlflow.create_experiment.assert_called_once_with("TestExp")
        mock_mlflow.set_experiment.assert_called_once_with("TestExp")
        mock_mlflow.start_run.assert_called_once_with(run_name="TestRun")
        mock_mlflow.log_params.assert_called_once_with({"lr": 0.001, "epochs": 10})

    @patch("dgtw_mlflow_utils.evaluator.mlflow")
    def test_on_train_begin_existing_experiment(self, mock_mlflow):
        mock_mlflow.active_run.return_value = None
        mock_mlflow.get_experiment_by_name.return_value = MagicMock()  # exists

        callback = TfMLflowEpochLogger(enable_mlflow=True, experiment_name="Existing")
        callback.on_train_begin()

        mock_mlflow.create_experiment.assert_not_called()
        mock_mlflow.set_experiment.assert_called_once_with("Existing")

    @patch("dgtw_mlflow_utils.evaluator.mlflow")
    def test_on_train_begin_ends_active_run(self, mock_mlflow):
        mock_mlflow.active_run.return_value = MagicMock()  # active run exists

        callback = TfMLflowEpochLogger(enable_mlflow=True)
        callback.on_train_begin()

        mock_mlflow.end_run.assert_called()

    @patch("dgtw_mlflow_utils.evaluator.mlflow")
    def test_on_train_begin_disabled(self, mock_mlflow):
        callback = TfMLflowEpochLogger(enable_mlflow=False)
        callback.on_train_begin()

        mock_mlflow.start_run.assert_not_called()

    @patch("dgtw_mlflow_utils.evaluator.mlflow")
    def test_on_train_begin_no_experiment_name(self, mock_mlflow):
        mock_mlflow.active_run.return_value = None

        callback = TfMLflowEpochLogger(enable_mlflow=True, experiment_name=None)
        callback.on_train_begin()

        mock_mlflow.set_experiment.assert_not_called()
        mock_mlflow.start_run.assert_called_once()

    # --- on_train_end ---

    @patch("dgtw_mlflow_utils.evaluator.mlflow")
    def test_on_train_end_enabled(self, mock_mlflow):
        callback = TfMLflowEpochLogger(enable_mlflow=True)
        callback.on_train_end()
        mock_mlflow.end_run.assert_called_once()

    @patch("dgtw_mlflow_utils.evaluator.mlflow")
    def test_on_train_end_disabled(self, mock_mlflow):
        callback = TfMLflowEpochLogger(enable_mlflow=False)
        callback.on_train_end()
        mock_mlflow.end_run.assert_not_called()

    # --- on_epoch_end ---

    @patch("dgtw_mlflow_utils.evaluator.mlflow")
    def test_on_epoch_end_logs_metrics(self, mock_mlflow):
        callback = TfMLflowEpochLogger(enable_mlflow=True)
        logs = {"loss": 0.5, "accuracy": 0.8}

        callback.on_epoch_end(epoch=2, logs=logs)
        mock_mlflow.log_metrics.assert_called_once_with(logs, step=2)

    @patch("dgtw_mlflow_utils.evaluator.mlflow")
    def test_on_epoch_end_none_logs(self, mock_mlflow):
        callback = TfMLflowEpochLogger(enable_mlflow=True)
        callback.on_epoch_end(epoch=0, logs=None)
        mock_mlflow.log_metrics.assert_not_called()

    @patch("dgtw_mlflow_utils.evaluator.mlflow")
    def test_on_epoch_end_disabled(self, mock_mlflow):
        callback = TfMLflowEpochLogger(enable_mlflow=False)
        callback.on_epoch_end(epoch=0, logs={"loss": 1.0})
        mock_mlflow.log_metrics.assert_not_called()

    # --- on_test_end ---

    # @patch("dgtw_mlflow_utils.evaluator.mlflow")
    # def test_on_test_end_logs_with_prefix(self, mock_mlflow):
    #     callback = TfMLflowEpochLogger(enable_mlflow=True)
    #     logs = {"loss": 0.3, "accuracy": 0.9}

    #     callback.on_test_end(logs=logs)

    #     expected = {"validation_loss": 0.3, "validation_accuracy": 0.9}
    #     mock_mlflow.log_metrics.assert_called_once_with(expected)

    # @patch("dgtw_mlflow_utils.evaluator.mlflow")
    # def test_on_test_end_none_logs(self, mock_mlflow):
    #     callback = TfMLflowEpochLogger(enable_mlflow=True)
    #     callback.on_test_end(logs=None)
    #     mock_mlflow.log_metrics.assert_not_called()

    # @patch("dgtw_mlflow_utils.evaluator.mlflow")
    # def test_on_test_end_disabled(self, mock_mlflow):
    #     callback = TfMLflowEpochLogger(enable_mlflow=False)
    #     callback.on_test_end(logs={"loss": 0.5})
    #     mock_mlflow.log_metrics.assert_not_called()


# ======================================================================
# evaluate_model_tf tests
# ======================================================================


class TestEvaluateModel:
    """Tests for the evaluate_model_tf function."""

    def test_returns_metric_dict(self):
        model = DummyTwoTower()
        dataset = _make_dataset(num_samples=16, batch_size=8)

        metrics = evaluate_model_tf(
            model=model,
            dataset=dataset,
            k=3,
            enable_mlflow=False,
        )

        assert isinstance(metrics, dict)
        assert "recall_3" in metrics
        assert "mrr_3" in metrics
        assert "ndcg_3" in metrics
        assert "loss" in metrics

    def test_metrics_are_floats(self):
        model = DummyTwoTower()
        dataset = _make_dataset(num_samples=16, batch_size=8)

        metrics = evaluate_model_tf(model=model, dataset=dataset, k=5, enable_mlflow=False)

        for key, val in metrics.items():
            assert isinstance(val, float), f"{key} is {type(val)}, expected float"

    def test_metrics_bounded(self):
        model = DummyTwoTower()
        dataset = _make_dataset(num_samples=16, batch_size=8)

        metrics = evaluate_model_tf(model=model, dataset=dataset, k=5, enable_mlflow=False)

        assert 0.0 <= metrics["recall_5"] <= 1.0
        assert 0.0 < metrics["mrr_5"] <= 1.0
        assert metrics["ndcg_5"] > 0.0
        assert metrics["loss"] >= 0.0

    def test_k_reflected_in_key(self):
        model = DummyTwoTower()
        dataset = _make_dataset(num_samples=32, batch_size=16)

        metrics = evaluate_model_tf(model=model, dataset=dataset, k=5, enable_mlflow=False)
        assert "recall_5" in metrics

    @patch("dgtw_mlflow_utils.evaluator.mlflow")
    def test_mlflow_enabled(self, mock_mlflow):
        mock_mlflow.get_experiment_by_name.return_value = MagicMock()

        model = DummyTwoTower()
        dataset = _make_dataset(num_samples=16, batch_size=8)

        evaluate_model_tf(
            model=model,
            dataset=dataset,
            k=5,
            enable_mlflow=True,
            experiment_name="TestExp",
            run_name="TestEval",
            log_params={"lr": 0.001},
        )

        mock_mlflow.set_experiment.assert_called_with("TestExp")
        mock_mlflow.start_run.assert_called_once_with(run_name="TestEval")
        mock_mlflow.log_params.assert_called_once_with({"lr": 0.001})
        mock_mlflow.log_metrics.assert_called_once()
        mock_mlflow.end_run.assert_called_once()

    @patch("dgtw_mlflow_utils.evaluator.mlflow")
    def test_mlflow_creates_experiment_if_missing(self, mock_mlflow):
        mock_mlflow.get_experiment_by_name.return_value = None

        model = DummyTwoTower()
        dataset = _make_dataset(num_samples=16, batch_size=8)

        evaluate_model_tf(
            model=model,
            dataset=dataset,
            k=5,
            enable_mlflow=True,
            experiment_name="NewExp",
        )

        mock_mlflow.create_experiment.assert_called_once_with("NewExp")

    @patch("dgtw_mlflow_utils.evaluator.mlflow")
    def test_mlflow_end_run_called_on_exception(self, mock_mlflow):
        """MLflow run should be ended even if an error occurs during evaluation."""
        mock_mlflow.get_experiment_by_name.return_value = MagicMock()

        # Create a model that raises on call
        model = MagicMock()
        model.side_effect = RuntimeError("boom")

        dataset = _make_dataset(num_samples=8, batch_size=8)

        with pytest.raises(RuntimeError, match="boom"):
            evaluate_model_tf(
                model=model,
                dataset=dataset,
                k=5,
                enable_mlflow=True,
                experiment_name="TestExp",
            )

        mock_mlflow.end_run.assert_called_once()

    def test_disabled_mlflow_no_error(self):
        """Should run cleanly without MLflow tracking."""
        model = DummyTwoTower()
        dataset = _make_dataset(num_samples=16, batch_size=8)

        metrics = evaluate_model_tf(model=model, dataset=dataset, k=5, enable_mlflow=False)
        assert len(metrics) == 4


# ======================================================================
# NumpyMLflowLogger tests
# ======================================================================


class TestNumpyMLflowLogger:
    """Tests for the NumpyMLflowLogger class."""

    def test_init_defaults(self):
        from dgtw_mlflow_utils.evaluator import NumpyMLflowLogger

        logger = NumpyMLflowLogger()
        assert logger.enable_mlflow is False
        assert logger.experiment_name is None
        assert logger.run_name == "Training_Run_Numpy"
        assert logger.log_params == {}

    @patch("dgtw_mlflow_utils.evaluator.mlflow")
    def test_on_train_begin_mlflow_enabled(self, mock_mlflow):
        from dgtw_mlflow_utils.evaluator import NumpyMLflowLogger

        mock_mlflow.active_run.return_value = None
        logger = NumpyMLflowLogger(enable_mlflow=True, experiment_name="Exp", run_name="Run")
        logger.on_train_begin()
        mock_mlflow.set_experiment.assert_called_once_with("Exp")
        mock_mlflow.start_run.assert_called_once_with(run_name="Run")

    @patch("dgtw_mlflow_utils.evaluator.mlflow")
    def test_on_epoch_end_logs_metrics(self, mock_mlflow):
        from dgtw_mlflow_utils.evaluator import NumpyMLflowLogger

        logger = NumpyMLflowLogger(enable_mlflow=True)
        metrics = {"loss": 0.5}
        logger.on_epoch_end(epoch=1, metrics=metrics)
        mock_mlflow.log_metrics.assert_called_once_with(metrics, step=1)


# ======================================================================
# evaluate_model_numpy tests
# ======================================================================


class TestEvaluateModelNumpy:
    """Tests for the evaluate_model_numpy function."""

    def test_returns_metric_dict(self):
        import numpy as np
        from dgtw_mlflow_utils.evaluator import evaluate_model_numpy

        def mock_predict(batch):
            return np.random.randn(8, 8)

        dataset = [{"batch": i} for i in range(2)]

        metrics = evaluate_model_numpy(
            predict_fn=mock_predict,
            dataset=dataset,
            k=3,
            enable_mlflow=False,
        )
        assert isinstance(metrics, dict)
        assert "recall_3" in metrics
