"""
tests/test_preprocessing.py
----------------------------
Unit tests for MultilingualE5LargeInstructPreProcessing.
No GPU or model download required.
"""

from __future__ import annotations

import pytest
from datasets import Dataset

from fiesta import MultilingualE5LargeInstructPreProcessing


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------

def _make_data(n_chunks: int = 5, n_aug: int = 15, n_noise: int = 5) -> list[dict]:
    """Build a minimal synthetic dataset."""
    return [
        {
            "chunk_text": f"Passage number {i} about topic {i}.",
            "docId": f"doc-{i:03d}",
            "questions": [
                {
                    "augmented_questions": [f"augmented q{j} for chunk {i}" for j in range(n_aug)],
                    "noise_questions": [f"noise q{j} for chunk {i}" for j in range(n_noise)],
                }
            ],
        }
        for i in range(n_chunks)
    ]


@pytest.fixture
def preprocessor() -> MultilingualE5LargeInstructPreProcessing:
    return MultilingualE5LargeInstructPreProcessing(n_test_samples_per_chunk=5)


@pytest.fixture
def raw_data() -> list[dict]:
    return _make_data(n_chunks=5, n_aug=15, n_noise=5)


# ---------------------------------------------------------------------------
# Tests: data_to_df
# ---------------------------------------------------------------------------

class TestDataToDf:
    def test_row_count(self, preprocessor, raw_data):
        df = preprocessor._data_to_df(raw_data)
        # 5 chunks * (15 aug + 5 noise) = 100 rows
        assert len(df) == 100

    def test_required_columns(self, preprocessor, raw_data):
        df = preprocessor._data_to_df(raw_data)
        for col in ("id", "anchor", "positive", "chunk_id", "doc_id"):
            assert col in df.columns, f"Missing column: {col}"

    def test_unique_chunk_ids(self, preprocessor, raw_data):
        df = preprocessor._data_to_df(raw_data)
        assert df["chunk_id"].nunique() == len(raw_data)

    def test_unique_question_ids(self, preprocessor, raw_data):
        df = preprocessor._data_to_df(raw_data)
        assert df["id"].nunique() == len(df)

    def test_empty_data_raises(self, preprocessor):
        with pytest.raises(ValueError, match="empty"):
            preprocessor._data_to_df([])

    def test_no_questions_raises(self, preprocessor):
        bad_data = [{"chunk_text": "text", "docId": "d1", "questions": [
            {"augmented_questions": [], "noise_questions": []}
        ]}]
        with pytest.raises(ValueError):
            preprocessor._data_to_df(bad_data)


# ---------------------------------------------------------------------------
# Tests: _preprocess_queries
# ---------------------------------------------------------------------------

class TestPreprocessQueries:
    def test_prompt_applied(self, preprocessor, raw_data):
        df = preprocessor._data_to_df(raw_data)
        df = preprocessor._preprocess_queries(df)
        assert df["anchor"].str.startswith("Instruct:").all()

    def test_lowercase(self, preprocessor, raw_data):
        df = preprocessor._data_to_df(raw_data)
        df = preprocessor._preprocess_queries(df)
        # everything after "Query: " should be lowercase
        queries = df["anchor"].str.split("Query: ").str[1]
        assert queries.str.lower().equals(queries)

    def test_punctuation_removed(self, preprocessor):
        import pandas as pd
        df = pd.DataFrame({"anchor": ["Hello, World!"], "positive": ["p"], "chunk_id": [0], "doc_id": ["d"], "id": [0]})
        result = preprocessor._preprocess_queries(df)
        query_part = result["anchor"].iloc[0].split("Query: ")[1]
        assert "," not in query_part
        assert "!" not in query_part

    def test_custom_task_description(self, raw_data):
        pp = MultilingualE5LargeInstructPreProcessing(task_description="Custom task")
        df = pp._data_to_df(raw_data)
        df = pp._preprocess_queries(df)
        assert df["anchor"].str.startswith("Instruct: Custom task").all()


# ---------------------------------------------------------------------------
# Tests: _split_by_question
# ---------------------------------------------------------------------------

class TestSplitByQuestion:
    def test_split_sizes(self, preprocessor, raw_data):
        df = preprocessor._data_to_df(raw_data)
        train, test = preprocessor._split_by_question(df)
        # 5 chunks * 5 test samples = 25 test rows
        assert len(test) == 25
        assert len(train) == len(df) - 25

    def test_no_overlap(self, preprocessor, raw_data):
        df = preprocessor._data_to_df(raw_data)
        train, test = preprocessor._split_by_question(df)
        overlap = set(train["id"]) & set(test["id"])
        assert len(overlap) == 0

    def test_single_chunk_raises(self, preprocessor):
        import pandas as pd
        df = pd.DataFrame({
            "id": list(range(20)),
            "anchor": [f"q{i}" for i in range(20)],
            "positive": ["same chunk"] * 20,
            "chunk_id": [0] * 20,
            "doc_id": ["d"] * 20,
        })
        with pytest.raises(ValueError, match="at least two"):
            preprocessor._split_by_question(df)


# ---------------------------------------------------------------------------
# Tests: get_evaluator_data
# ---------------------------------------------------------------------------

class TestGetEvaluatorData:
    def test_keys(self, preprocessor, raw_data):
        df = preprocessor._data_to_df(raw_data)
        train, test = preprocessor._split_by_question(df)
        ev = preprocessor.get_evaluator_data(train, test)
        assert set(ev.keys()) == {"train_test_chunks", "test_queries", "testQueryId_to_chunkId"}

    def test_chunks_cover_train_and_test(self, preprocessor, raw_data):
        df = preprocessor._data_to_df(raw_data)
        train, test = preprocessor._split_by_question(df)
        ev = preprocessor.get_evaluator_data(train, test)
        assert len(ev["train_test_chunks"]) == 5  # 5 unique chunks

    def test_relevant_docs_values_are_lists(self, preprocessor, raw_data):
        df = preprocessor._data_to_df(raw_data)
        train, test = preprocessor._split_by_question(df)
        ev = preprocessor.get_evaluator_data(train, test)
        for v in ev["testQueryId_to_chunkId"].values():
            assert isinstance(v, list)


# ---------------------------------------------------------------------------
# Tests: full preprocess() pipeline
# ---------------------------------------------------------------------------

class TestPreprocess:
    def test_returns_datasets(self, preprocessor, raw_data):
        train_ds, test_ds, ev = preprocessor.preprocess(raw_data)
        assert isinstance(train_ds, Dataset)
        assert isinstance(test_ds, Dataset)

    def test_dataset_columns(self, preprocessor, raw_data):
        train_ds, test_ds, _ = preprocessor.preprocess(raw_data)
        for col in ("anchor", "positive", "chunk_id", "id"):
            assert col in train_ds.column_names
            assert col in test_ds.column_names

    def test_no_train_test_id_overlap(self, preprocessor, raw_data):
        train_ds, test_ds, _ = preprocessor.preprocess(raw_data)
        train_ids = set(train_ds["id"])
        test_ids = set(test_ds["id"])
        assert train_ids.isdisjoint(test_ids)

    def test_empty_raises(self, preprocessor):
        with pytest.raises(ValueError):
            preprocessor.preprocess([])


# ---------------------------------------------------------------------------
# Tests: dataclasses (LoraSettings, TrainingSettings, TrainResult)
# ---------------------------------------------------------------------------

class TestDataclasses:
    def test_lora_defaults(self):
        from fiesta import LoraSettings
        ls = LoraSettings()
        assert ls.r == 16
        assert ls.lora_alpha == 16
        assert ls.lora_dropout == 0.0
        assert ls.bias == "none"
        assert ls.target_modules == ["query", "key", "value", "dense"]

    def test_training_defaults(self):
        from fiesta import TrainingSettings
        ts = TrainingSettings()
        assert ts.max_steps == 200
        assert ts.fp16 is True
        assert ts.optim == "adafactor"

    def test_train_result_improvement(self):
        from fiesta import TrainResult
        r = TrainResult(
            baseline={"cosine_accuracy@10": 0.5},
            final_metrics={"cosine_accuracy@10": 0.7},
        )
        delta = r.improvement()
        assert abs(delta["cosine_accuracy@10"] - 0.2) < 1e-9

    def test_train_result_negative_improvement(self):
        from fiesta import TrainResult
        r = TrainResult(
            baseline={"cosine_accuracy@10": 0.8},
            final_metrics={"cosine_accuracy@10": 0.6},
        )
        delta = r.improvement()
        assert delta["cosine_accuracy@10"] < 0


# ---------------------------------------------------------------------------
# Tests: pipeline init (no model load)
# ---------------------------------------------------------------------------

class TestPipelineInit:
    def test_save_dir_created(self, tmp_path):
        from unittest.mock import patch
        # Patch modelling __init__ to avoid downloading the model
        with patch(
            "fiesta.pipeline.MultilingualE5LargeInstructModelling.__init__",
            return_value=None,
        ):
            from fiesta import MultilingualE5LargeInstructPipeline
            p = MultilingualE5LargeInstructPipeline(
                kb_id="test-kb", lang="EN", base_dir=str(tmp_path)
            )
            assert p.save_dir.exists()
            assert p.lang == "en"  # lowercased

    def test_save_dir_structure(self, tmp_path):
        from unittest.mock import patch
        with patch(
            "fiesta.pipeline.MultilingualE5LargeInstructModelling.__init__",
            return_value=None,
        ):
            from fiesta import MultilingualE5LargeInstructPipeline
            p = MultilingualE5LargeInstructPipeline(
                kb_id="MY-KB", lang="PT", base_dir=str(tmp_path)
            )
            expected = (
                tmp_path
                / "fiesta"
                / "pt"
                / "my-kb.pt"
            )
            assert p.save_dir == expected


# ---------------------------------------------------------------------------
# Tests: legacy static helpers
# ---------------------------------------------------------------------------

class TestLegacyStatics:
    def test_split_train_test_by_question(self, raw_data):
        import pandas as pd
        pp = MultilingualE5LargeInstructPreProcessing()
        df = pp._data_to_df(raw_data)
        train, test = MultilingualE5LargeInstructPreProcessing.split_train_test_by_question(df, test_size=0.2)
        assert len(train) + len(test) == len(df)
        assert len(test) <= 500

    def test_split_train_test_by_chunk_basic(self, raw_data):
        pp = MultilingualE5LargeInstructPreProcessing()
        df = pp._data_to_df(raw_data)
        train, test = MultilingualE5LargeInstructPreProcessing.split_train_test_by_chunk(df, test_size=0.2)
        # No chunk appears in both splits
        assert set(train["positive"]).isdisjoint(set(test["positive"]))

    def test_split_train_test_by_chunk_single_value_raises(self):
        import pandas as pd
        df = pd.DataFrame({
            "id": [0, 1],
            "anchor": ["q0", "q1"],
            "positive": ["same", "same"],
            "chunk_id": [0, 0],
            "doc_id": ["d", "d"],
        })
        with pytest.raises(ValueError, match="one unique value"):
            MultilingualE5LargeInstructPreProcessing.split_train_test_by_chunk(df)


# ---------------------------------------------------------------------------
# Tests: _split_by_question warning branch (chunk with too few samples)
# ---------------------------------------------------------------------------

class TestSplitWarningBranch:
    def test_small_chunk_goes_to_train(self):
        """A chunk with <= n_test_samples_per_chunk rows must land entirely in train."""
        import pandas as pd
        pp = MultilingualE5LargeInstructPreProcessing(n_test_samples_per_chunk=10)

        # chunk 0: 20 rows (enough for split)
        # chunk 1: 5 rows  (too few — all go to train, warning logged)
        rows = (
            [{"id": i, "anchor": f"q{i}", "positive": "chunk-A", "chunk_id": 0, "doc_id": "d0"} for i in range(20)]
            + [{"id": 20 + i, "anchor": f"q{20+i}", "positive": "chunk-B", "chunk_id": 1, "doc_id": "d1"} for i in range(5)]
        )
        df = pd.DataFrame(rows)

        train, test = pp._split_by_question(df)

        # All chunk-B rows must be in train only
        assert all(r == "chunk-B" for r in train[train["positive"] == "chunk-B"]["positive"])
        assert "chunk-B" not in set(test["positive"])


