"""
tests/test_modelling.py
-----------------------
Tests for MultilingualE5LargeInstructModelling.
All heavy dependencies (SentenceTransformer, PEFT, trainer) are mocked
so no GPU or model download is needed.
"""

from __future__ import annotations

import pytest
from unittest.mock import MagicMock, patch

from fiesta.modelling import (
    LoraSettings,
    TrainingSettings,
    TrainResult,
    MultilingualE5LargeInstructModelling,
)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _make_modeller(lora: LoraSettings | None = None, training: TrainingSettings | None = None):
    """Return a Modelling instance with all heavy __init__ side-effects mocked."""
    mock_model = MagicMock()
    mock_model.__getitem__ = MagicMock(return_value=MagicMock())

    with patch("fiesta.modelling.SentenceTransformer", return_value=mock_model), \
         patch("fiesta.modelling.get_peft_model", return_value=MagicMock()), \
         patch("fiesta.modelling.torch.cuda.is_available", return_value=False):
        m = MultilingualE5LargeInstructModelling(lora_settings=lora, training_settings=training)
    return m


# ---------------------------------------------------------------------------
# Tests: __init__ / settings
# ---------------------------------------------------------------------------

class TestModellingInit:
    def test_default_settings_applied(self):
        m = _make_modeller()
        assert m.lora_settings == LoraSettings()
        assert m.training_settings == TrainingSettings()

    def test_custom_lora_settings_stored(self):
        custom = LoraSettings(r=32, lora_alpha=64)
        m = _make_modeller(lora=custom)
        assert m.lora_settings.r == 32
        assert m.lora_settings.lora_alpha == 64

    def test_custom_training_settings_stored(self):
        custom = TrainingSettings(max_steps=50, fp16=False)
        m = _make_modeller(training=custom)
        assert m.training_settings.max_steps == 50
        assert m.training_settings.fp16 is False

    def test_cpu_selected_when_no_cuda(self):
        with patch("fiesta.modelling.torch.cuda.is_available", return_value=False), \
             patch("fiesta.modelling.SentenceTransformer") as mock_st, \
             patch("fiesta.modelling.get_peft_model", return_value=MagicMock()):
            mock_st.return_value = MagicMock()
            mock_st.return_value.__getitem__ = MagicMock(return_value=MagicMock())
            MultilingualE5LargeInstructModelling()
            mock_st.assert_called_once_with(MultilingualE5LargeInstructModelling.MODEL_ID, device="cpu")

    def test_cuda_selected_when_available(self):
        with patch("fiesta.modelling.torch.cuda.is_available", return_value=True), \
             patch("fiesta.modelling.SentenceTransformer") as mock_st, \
             patch("fiesta.modelling.get_peft_model", return_value=MagicMock()):
            mock_st.return_value = MagicMock()
            mock_st.return_value.__getitem__ = MagicMock(return_value=MagicMock())
            MultilingualE5LargeInstructModelling()
            mock_st.assert_called_once_with(MultilingualE5LargeInstructModelling.MODEL_ID, device="cuda")


# ---------------------------------------------------------------------------
# Tests: _build_training_args
# ---------------------------------------------------------------------------

class TestBuildTrainingArgs:
    def test_output_dir_set(self, tmp_path):
        m = _make_modeller()
        with patch("fiesta.modelling.SentenceTransformerTrainingArguments") as mock_args:
            mock_args.return_value = MagicMock()
            m._build_training_args(str(tmp_path))
            call_kwargs = mock_args.call_args.kwargs
            assert call_kwargs["output_dir"] == str(tmp_path)

    def test_training_settings_propagated(self):
        ts = TrainingSettings(max_steps=99, learning_rate=5e-5, fp16=False, eval_steps=7)
        m = _make_modeller(training=ts)
        with patch("fiesta.modelling.SentenceTransformerTrainingArguments") as mock_args:
            mock_args.return_value = MagicMock()
            m._build_training_args("./out")
            kw = mock_args.call_args.kwargs
            assert kw["max_steps"] == 99
            assert kw["learning_rate"] == 5e-5
            assert kw["fp16"] is False
            assert kw["eval_steps"] == 7
            assert kw["save_steps"] == 7


# ---------------------------------------------------------------------------
# Tests: _build_evaluator
# ---------------------------------------------------------------------------

class TestBuildEvaluator:
    def test_evaluator_constructed_with_correct_args(self):
        m = _make_modeller()
        queries = {"q0": "what is X?"}
        corpus = {"c0": "X is Y."}
        relevant = {"q0": ["c0"]}
        with patch("fiesta.modelling.InformationRetrievalEvaluator") as mock_ev:
            mock_ev.return_value = MagicMock()
            m._build_evaluator(queries, corpus, relevant)
            kw = mock_ev.call_args.kwargs
            assert kw["queries"] is queries
            assert kw["corpus"] is corpus
            # relevant_docs must be converted to sets
            assert kw["relevant_docs"] == {"q0": {"c0"}}
            assert kw["batch_size"] == m.training_settings.evaluator_batch_size


# ---------------------------------------------------------------------------
# Tests: train()
# ---------------------------------------------------------------------------

class TestModellingTrain:
    def _make_dataset(self):
        from datasets import Dataset
        return Dataset.from_dict({"anchor": ["q1", "q2"], "positive": ["p1", "p2"]})

    def test_train_returns_train_result(self, tmp_path):
        m = _make_modeller()
        train_ds = self._make_dataset()
        test_ds = self._make_dataset()
        evaluator_data = {
            "test_queries": {0: "q"},
            "train_test_chunks": {0: "p"},
            "testQueryId_to_chunkId": {0: [0]},
        }

        mock_evaluator = MagicMock(side_effect=[{"cosine_accuracy@10": 0.4}, {"cosine_accuracy@10": 0.6}])
        mock_trainer = MagicMock()

        with patch.object(m, "_build_training_args", return_value=MagicMock()), \
             patch.object(m, "_build_evaluator", return_value=mock_evaluator), \
             patch("fiesta.modelling.CachedMultipleNegativesRankingLoss", return_value=MagicMock()), \
             patch("fiesta.modelling.SentenceTransformerTrainer", return_value=mock_trainer):

            result = m.train(train_ds, test_ds, evaluator_data, str(tmp_path))

        assert isinstance(result, TrainResult)
        assert result.baseline == {"cosine_accuracy@10": 0.4}
        assert result.final_metrics == {"cosine_accuracy@10": 0.6}

    def test_train_calls_trainer_train_and_save(self, tmp_path):
        m = _make_modeller()
        train_ds = self._make_dataset()
        test_ds = self._make_dataset()
        evaluator_data = {
            "test_queries": {0: "q"},
            "train_test_chunks": {0: "p"},
            "testQueryId_to_chunkId": {0: [0]},
        }

        mock_trainer = MagicMock()
        mock_evaluator = MagicMock(side_effect=[{}, {}])

        with patch.object(m, "_build_training_args", return_value=MagicMock()), \
             patch.object(m, "_build_evaluator", return_value=mock_evaluator), \
             patch("fiesta.modelling.CachedMultipleNegativesRankingLoss", return_value=MagicMock()), \
             patch("fiesta.modelling.SentenceTransformerTrainer", return_value=mock_trainer):

            m.train(train_ds, test_ds, evaluator_data, str(tmp_path))

        mock_trainer.train.assert_called_once()
        mock_trainer.save_model.assert_called_once()

    def test_evaluator_called_twice(self, tmp_path):
        """Baseline and final evaluation must each call the evaluator once."""
        m = _make_modeller()
        train_ds = self._make_dataset()
        test_ds = self._make_dataset()
        evaluator_data = {
            "test_queries": {0: "q"},
            "train_test_chunks": {0: "p"},
            "testQueryId_to_chunkId": {0: [0]},
        }

        mock_evaluator = MagicMock(side_effect=[{"k": 0.1}, {"k": 0.2}])

        with patch.object(m, "_build_training_args", return_value=MagicMock()), \
             patch.object(m, "_build_evaluator", return_value=mock_evaluator), \
             patch("fiesta.modelling.CachedMultipleNegativesRankingLoss", return_value=MagicMock()), \
             patch("fiesta.modelling.SentenceTransformerTrainer", return_value=MagicMock()):

            m.train(train_ds, test_ds, evaluator_data, str(tmp_path))

        assert mock_evaluator.call_count == 2


# ---------------------------------------------------------------------------
# Tests: _build_evaluator — relevant_docs must be converted to sets
# ---------------------------------------------------------------------------

class TestBuildEvaluatorSets:
    def test_relevant_docs_converted_to_sets(self):
        m = _make_modeller()
        queries = {"q0": "q?"}
        corpus = {"c0": "passage.", "c1": "other."}
        relevant = {"q0": ["c0", "c1"]}  # list input

        with patch("fiesta.modelling.InformationRetrievalEvaluator") as mock_ev:
            mock_ev.return_value = MagicMock()
            m._build_evaluator(queries, corpus, relevant)
            kw = mock_ev.call_args.kwargs
            # Must be a set, not a list
            assert isinstance(kw["relevant_docs"]["q0"], set)
            assert kw["relevant_docs"]["q0"] == {"c0", "c1"}


# ---------------------------------------------------------------------------
# Tests: TrainResult — disjoint metric keys
# ---------------------------------------------------------------------------

class TestTrainResultDisjointKeys:
    def test_missing_baseline_key_defaults_to_zero(self):
        r = TrainResult(
            baseline={},
            final_metrics={"cosine_accuracy@10": 0.5},
        )
        delta = r.improvement()
        assert delta["cosine_accuracy@10"] == 0.5

    def test_improvement_empty_when_no_final_metrics(self):
        r = TrainResult(
            baseline={"cosine_accuracy@10": 0.9},
            final_metrics={},
        )
        assert r.improvement() == {}


# ---------------------------------------------------------------------------
# Tests: fp16 CPU warning
# ---------------------------------------------------------------------------

class TestFp16Warning:
    def test_fp16_cpu_logs_warning(self, caplog):
        import logging
        ts = TrainingSettings(fp16=True)
        with patch("fiesta.modelling.torch.cuda.is_available", return_value=False), \
             patch("fiesta.modelling.SentenceTransformer", return_value=MagicMock()), \
             patch("fiesta.modelling.get_peft_model", return_value=MagicMock()):
            with caplog.at_level(logging.WARNING, logger="fiesta.modelling"):
                MultilingualE5LargeInstructModelling(training_settings=ts)
        assert any("fp16" in r.message for r in caplog.records)

    def test_no_warning_when_fp16_false(self, caplog):
        import logging
        ts = TrainingSettings(fp16=False)
        with patch("fiesta.modelling.torch.cuda.is_available", return_value=False), \
             patch("fiesta.modelling.SentenceTransformer", return_value=MagicMock()), \
             patch("fiesta.modelling.get_peft_model", return_value=MagicMock()):
            with caplog.at_level(logging.WARNING, logger="fiesta.modelling"):
                MultilingualE5LargeInstructModelling(training_settings=ts)
        assert not any("fp16" in r.message for r in caplog.records)

