"""
tests/test_pipeline.py
----------------------
Tests for MultilingualE5LargeInstructPipeline.
Model loading is always mocked — no GPU or download needed.
"""

from __future__ import annotations

from pathlib import Path
from unittest.mock import MagicMock, patch
import pytest

from fiesta.modelling import TrainResult


# ---------------------------------------------------------------------------
# Helper: build pipeline with mocked modeller init
# ---------------------------------------------------------------------------

def _make_pipeline(tmp_path, kb_id="kb", lang="en", **kwargs):
    with patch(
        "fiesta.pipeline.MultilingualE5LargeInstructModelling.__init__",
        return_value=None,
    ):
        from fiesta import MultilingualE5LargeInstructPipeline
        return MultilingualE5LargeInstructPipeline(
            kb_id=kb_id, lang=lang, base_dir=str(tmp_path), **kwargs
        )


# ---------------------------------------------------------------------------
# Tests: init
# ---------------------------------------------------------------------------

class TestPipelineInit:
    def test_save_dir_created(self, tmp_path):
        p = _make_pipeline(tmp_path)
        assert p.save_dir.exists()

    def test_lang_lowercased(self, tmp_path):
        p = _make_pipeline(tmp_path, lang="PT")
        assert p.lang == "pt"

    def test_kb_id_lowercased_in_path(self, tmp_path):
        p = _make_pipeline(tmp_path, kb_id="MY-KB", lang="en")
        assert "my-kb" in str(p.save_dir)

    def test_save_dir_structure(self, tmp_path):
        p = _make_pipeline(tmp_path, kb_id="my-kb", lang="fr")
        expected = tmp_path / "fiesta" / "fr" / "my-kb.fr"
        assert p.save_dir == expected

    def test_default_base_dir_is_home(self):
        with patch(
            "fiesta.pipeline.MultilingualE5LargeInstructModelling.__init__",
            return_value=None,
        ):
            from fiesta import MultilingualE5LargeInstructPipeline
            p = MultilingualE5LargeInstructPipeline(kb_id="kb", lang="en")
            assert str(Path.home()) in str(p.save_dir)
            # clean up the created directory
            import shutil
            shutil.rmtree(p.save_dir, ignore_errors=True)

    def test_preprocessing_settings_forwarded(self, tmp_path):
        p = _make_pipeline(
            tmp_path,
            preprocessing_settings={"task_description": "Custom task", "n_test_samples_per_chunk": 3},
        )
        assert p.preprocessor.task_description == "Custom task"
        assert p.preprocessor.n_test_samples_per_chunk == 3


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

class TestPipelineTrain:
    def _raw_data(self, n_chunks=5, n_aug=15, n_noise=5):
        return [
            {
                "chunk_text": f"Passage {i}.",
                "docId": f"doc-{i}",
                "questions": [{
                    "augmented_questions": [f"aug q{j} chunk {i}" for j in range(n_aug)],
                    "noise_questions": [f"noise q{j} chunk {i}" for j in range(n_noise)],
                }],
            }
            for i in range(n_chunks)
        ]

    def test_train_returns_train_result(self, tmp_path):
        p = _make_pipeline(tmp_path)
        expected = TrainResult(
            baseline={"cosine_accuracy@10": 0.3},
            final_metrics={"cosine_accuracy@10": 0.5},
        )
        p.modeller.train = MagicMock(return_value=expected)

        result = p.train(self._raw_data())

        assert isinstance(result, TrainResult)
        assert result.baseline["cosine_accuracy@10"] == 0.3

    def test_train_calls_modeller_with_correct_args(self, tmp_path):
        p = _make_pipeline(tmp_path)
        p.modeller.train = MagicMock(return_value=TrainResult({}, {}))

        p.train(self._raw_data())

        call_kwargs = p.modeller.train.call_args.kwargs
        assert "train_dataset" in call_kwargs
        assert "test_dataset" in call_kwargs
        assert "evaluator_data" in call_kwargs
        assert call_kwargs["save_dir"] == str(p.save_dir)

    def test_train_deletes_checkpoints_after_training(self, tmp_path):
        p = _make_pipeline(tmp_path)
        p.modeller.train = MagicMock(return_value=TrainResult({}, {}))

        # Create fake checkpoint dirs to be cleaned up
        ckpt1 = p.save_dir / "checkpoint-5"
        ckpt2 = p.save_dir / "checkpoint-10"
        ckpt1.mkdir()
        ckpt2.mkdir()

        p.train(self._raw_data())

        assert not ckpt1.exists()
        assert not ckpt2.exists()

    def test_train_ignores_missing_checkpoints(self, tmp_path):
        p = _make_pipeline(tmp_path)
        p.modeller.train = MagicMock(return_value=TrainResult({}, {}))
        # No checkpoint dirs exist — should not raise
        p.train(self._raw_data())


# ---------------------------------------------------------------------------
# Tests: _delete_checkpoints
# ---------------------------------------------------------------------------

class TestDeleteCheckpoints:
    def test_removes_checkpoint_dirs(self, tmp_path):
        p = _make_pipeline(tmp_path)
        ckpt = p.save_dir / "checkpoint-20"
        ckpt.mkdir()
        p._delete_checkpoints()
        assert not ckpt.exists()

    def test_does_not_remove_non_checkpoint_dirs(self, tmp_path):
        p = _make_pipeline(tmp_path)
        keep = p.save_dir / "final_model"
        keep.mkdir()
        p._delete_checkpoints()
        assert keep.exists()

    def test_handles_empty_save_dir(self, tmp_path):
        p = _make_pipeline(tmp_path)
        # Should not raise even with nothing to delete
        p._delete_checkpoints()

    def test_removes_multiple_checkpoints(self, tmp_path):
        p = _make_pipeline(tmp_path)
        ckpts = [p.save_dir / f"checkpoint-{i}" for i in range(3)]
        for c in ckpts:
            c.mkdir()
        p._delete_checkpoints()
        assert all(not c.exists() for c in ckpts)

    def test_oserror_on_delete_is_logged_not_raised(self, tmp_path, caplog):
        """An OSError during rmtree must be logged, not re-raised."""
        import logging
        from unittest.mock import patch
        p = _make_pipeline(tmp_path)
        ckpt = p.save_dir / "checkpoint-99"
        ckpt.mkdir()

        with patch("fiesta.pipeline.shutil.rmtree", side_effect=OSError("boom")):
            with caplog.at_level(logging.ERROR, logger="fiesta.pipeline"):
                p._delete_checkpoints()  # must not raise

        assert any("boom" in r.message or "checkpoint" in r.message.lower() for r in caplog.records)


# ---------------------------------------------------------------------------
# Tests: BasePipeline abstract contract
# ---------------------------------------------------------------------------

class TestBasePipeline:
    def test_cannot_instantiate_base_directly(self):
        from fiesta.pipeline import BasePipeline
        with pytest.raises(TypeError):
            BasePipeline()  # type: ignore[abstract]

    def test_concrete_subclass_works(self, tmp_path):
        from fiesta.pipeline import BasePipeline
        from typing import Any

        class Concrete(BasePipeline):
            def train(self, data: list[dict[str, Any]], **kwargs: Any) -> Any:
                return "ok"

        assert Concrete().train([]) == "ok"
