"""
Tests for babilong_eval — all offline, no API calls required.

Tests cover:
  1. Prompt templates (completeness, format)
  2. Metrics (label matching, preprocessing, edge cases)
  3. Preprocessor (anomaly scoring, chunking, filtering)
  4. Report (aggregation, formatting)
  5. Harness (config validation)
"""

import json
import math
import os
import tempfile

import pytest

from babilong_eval.prompts import (
    PromptTemplate,
    get_prompt_template,
    format_prompt,
)
from babilong_eval.metrics import (
    TASK_LABELS,
    preprocess_output,
    compare_answers,
    compute_accuracy,
)
from babilong_eval.preprocessor import (
    _compute_anomaly_score,
    _chunk_context,
    _compute_corpus_stats,
    _compute_general_anomaly_score,
    _CorpusStats,
    full_context,
    hcp_filtered,
    random_filtered,
    bm25_filtered,
    embedding_filtered,
)
from babilong_eval.report import (
    SampleResult,
    BabilongResult,
    print_report,
    save_results,
    _group_results,
    _compute_accuracy as report_accuracy,
)
from babilong_eval.harness import BabilongConfig


# ===================================================================
# 1. PROMPT TEMPLATES
# ===================================================================


class TestPromptTemplates:
    """Tests for BABILong prompt templates."""

    def test_all_five_tasks_exist(self):
        for task in ("qa1", "qa2", "qa3", "qa4", "qa5"):
            tmpl = get_prompt_template(task)
            assert tmpl.task == task

    def test_unknown_task_raises(self):
        with pytest.raises(ValueError, match="Unknown task"):
            get_prompt_template("qa99")

    def test_template_has_all_fields(self):
        tmpl = get_prompt_template("qa1")
        assert isinstance(tmpl.instruction, str) and len(tmpl.instruction) > 10
        assert isinstance(tmpl.examples, str) and len(tmpl.examples) > 10
        assert isinstance(tmpl.post_prompt, str) and len(tmpl.post_prompt) > 10

    def test_qa1_mentions_place(self):
        tmpl = get_prompt_template("qa1")
        assert "place" in tmpl.instruction.lower()

    def test_qa2_mentions_two_facts(self):
        tmpl = get_prompt_template("qa2")
        assert "two" in tmpl.instruction.lower()

    def test_qa3_mentions_three_facts(self):
        tmpl = get_prompt_template("qa3")
        assert "three" in tmpl.instruction.lower()

    def test_format_prompt_structure(self):
        prompt = format_prompt("qa1", "Some context here.", "Where is Mary?")
        assert "<context>" in prompt
        assert "</context>" in prompt
        assert "Question: Where is Mary?" in prompt
        assert "Some context here." in prompt

    def test_format_prompt_ordering(self):
        """Instruction comes before examples, before post_prompt, before context."""
        prompt = format_prompt("qa1", "CTX_HERE", "Q_HERE")
        tmpl = get_prompt_template("qa1")
        i_pos = prompt.index(tmpl.instruction[:20])
        e_pos = prompt.index(tmpl.examples[:20])
        p_pos = prompt.index(tmpl.post_prompt[:20])
        c_pos = prompt.index("CTX_HERE")
        q_pos = prompt.index("Q_HERE")
        assert i_pos < e_pos < p_pos < c_pos < q_pos

    def test_case_insensitive_task_lookup(self):
        t1 = get_prompt_template("QA1")
        t2 = get_prompt_template("qa1")
        assert t1.task == t2.task


# ===================================================================
# 2. METRICS
# ===================================================================


class TestMetrics:
    """Tests for BABILong scoring logic."""

    def test_task_labels_all_present(self):
        for task in ("qa1", "qa2", "qa3", "qa4", "qa5"):
            assert task in TASK_LABELS
            assert len(TASK_LABELS[task]) >= 6

    def test_qa1_labels(self):
        expected = {"bathroom", "bedroom", "garden", "hallway", "kitchen", "office"}
        assert set(TASK_LABELS["qa1"]) == expected

    def test_preprocess_lowercase(self):
        assert preprocess_output("KITCHEN") == "kitchen"

    def test_preprocess_first_sentence(self):
        assert preprocess_output("kitchen. Based on the context...") == "kitchen"

    def test_preprocess_strip_after_context(self):
        assert preprocess_output("The answer is kitchen context blah") == "the answer is kitchen"

    def test_preprocess_strip_after_question(self):
        assert preprocess_output("kitchen Question: what?") == "kitchen"

    def test_compare_answers_correct(self):
        assert compare_answers("kitchen", "kitchen", "Where is Mary?", "qa1")

    def test_compare_answers_with_surrounding_text(self):
        """Model says more than just the label — should still match."""
        assert compare_answers(
            "The answer is kitchen.",
            "kitchen",
            "Where is Mary?",
            "qa1",
        )

    def test_compare_answers_wrong_label(self):
        assert not compare_answers("bedroom", "kitchen", "Where is Mary?", "qa1")

    def test_compare_answers_multiple_labels_fails(self):
        """If model outputs two labels, should fail (too ambiguous)."""
        assert not compare_answers(
            "kitchen or bedroom",
            "kitchen",
            "Where is Mary?",
            "qa1",
        )

    def test_compare_answers_label_in_question_subtracted(self):
        """Labels in the question are subtracted from predictions."""
        # Question mentions "kitchen", output says "kitchen" — net is empty → fail
        assert not compare_answers(
            "kitchen",
            "kitchen",
            "Mary went to the kitchen. Where is she now?",
            "qa1",
        )

    def test_compare_answers_target_not_in_question(self):
        """Normal case: target not in question, model outputs target."""
        assert compare_answers(
            "hallway",
            "hallway",
            "Where is Bob?",
            "qa1",
        )

    def test_compare_answers_qa5_person_name(self):
        """QA5 can have person names as labels."""
        assert compare_answers("bill", "bill", "Who went first?", "qa5")

    def test_compare_answers_empty_output(self):
        assert not compare_answers("", "kitchen", "Where is Mary?", "qa1")

    def test_compute_accuracy_all_correct(self):
        results = [
            ("kitchen", "kitchen", "Where is Mary?", "qa1"),
            ("bedroom", "bedroom", "Where is Bob?", "qa1"),
        ]
        assert compute_accuracy(results) == 1.0

    def test_compute_accuracy_half_correct(self):
        results = [
            ("kitchen", "kitchen", "Where is Mary?", "qa1"),
            ("kitchen", "bedroom", "Where is Bob?", "qa1"),
        ]
        assert compute_accuracy(results) == 0.5

    def test_compute_accuracy_empty(self):
        assert compute_accuracy([]) == 0.0


# ===================================================================
# 3. PREPROCESSOR — Anomaly Scoring
# ===================================================================


class TestAnomalyScoring:
    """Tests for the anomaly score that detects bAbI facts in PG-19 filler."""

    def test_babi_fact_high_score(self):
        """A typical bAbI fact should score high."""
        score = _compute_anomaly_score("Mary went to the kitchen.")
        assert score > 0.4, f"bAbI fact scored too low: {score}"

    def test_pg19_prose_low_score(self):
        """PG-19 literary prose should score low."""
        prose = (
            "The old gentleman regarded his companion with a mixture of "
            "curiosity and satisfaction, as one who contemplated a remarkable "
            "specimen of humanity, and possessed, withal, a keen sense of "
            "the ridiculous aspects of the situation."
        )
        score = _compute_anomaly_score(prose)
        assert score < 0.35, f"PG-19 prose scored too high: {score}"

    def test_babi_fact_beats_prose(self):
        """bAbI facts should always score higher than PG-19 prose."""
        fact_score = _compute_anomaly_score("John journeyed to the bathroom.")
        prose_score = _compute_anomaly_score(
            "The mansion stood silently among the ancient oaks, its weathered "
            "facade bearing the accumulated patina of centuries, a monument "
            "to the enduring permanence of stone and mortar."
        )
        assert fact_score > prose_score

    def test_multiple_babi_facts_high_score(self):
        """Multiple bAbI facts in one chunk should score very high."""
        text = (
            "Mary went to the kitchen. Bob journeyed to the office. "
            "Sandra travelled to the garden."
        )
        score = _compute_anomaly_score(text)
        assert score > 0.4

    def test_empty_string_zero(self):
        assert _compute_anomaly_score("") == 0.0

    def test_score_bounded(self):
        """Score must be in [0, 1]."""
        texts = [
            "Mary went to the kitchen.",
            "A very long passage of literary text with many clauses, "
            "semicolons; commas, and complex syntactic structures.",
            "",
            "the",
        ]
        for text in texts:
            score = _compute_anomaly_score(text)
            assert 0.0 <= score <= 1.0, f"Score out of bounds: {score}"

    def test_direction_sentence_moderate_score(self):
        """QA4-style directional facts should also score above prose."""
        score = _compute_anomaly_score("The kitchen is north of the hallway.")
        prose_score = _compute_anomaly_score(
            "The elaborate ceremony proceeded with all due pomp and circumstance."
        )
        assert score > prose_score


# ===================================================================
# 4. PREPROCESSOR — Chunking
# ===================================================================


class TestChunking:
    """Tests for context chunking."""

    def test_basic_chunking(self):
        text = "First sentence. Second sentence. Third sentence."
        chunks = _chunk_context(text, chunk_size=100)
        assert len(chunks) >= 1
        # All text should be preserved
        reconstructed = " ".join(c.content for c in chunks)
        assert "First sentence" in reconstructed

    def test_chunk_positions_sequential(self):
        text = "A. " * 50 + "B. " * 50  # Long enough to need multiple chunks
        chunks = _chunk_context(text, chunk_size=50)
        positions = [c.position for c in chunks]
        assert positions == list(range(len(chunks)))

    def test_small_chunk_size(self):
        """Small chunk_size produces more chunks."""
        text = "Mary went to the kitchen. Bob went to the office. Sandra went to the garden."
        small = _chunk_context(text, chunk_size=40)
        large = _chunk_context(text, chunk_size=200)
        assert len(small) >= len(large)

    def test_empty_context(self):
        assert _chunk_context("") == []


# ===================================================================
# 5. PREPROCESSOR — Filtering
# ===================================================================


class TestFiltering:
    """Tests for the three experimental conditions."""

    _CONTEXT = (
        "The old mansion stood silently among the ancient oaks. "
        "Its weathered facade bore the patina of centuries. "
        "Mary went to the kitchen. "
        "The gardens stretched endlessly toward the horizon, "
        "their carefully maintained hedgerows forming geometric patterns. "
        "Bob journeyed to the office. "
        "The library contained thousands of leather-bound volumes."
    )

    def test_full_context_passthrough(self):
        result = full_context(self._CONTEXT)
        assert result == self._CONTEXT

    def test_hcp_filtered_shorter(self):
        result = hcp_filtered(
            self._CONTEXT,
            question="Where is Mary?",
            task="qa1",
            retain_ratio=0.5,
        )
        assert len(result) < len(self._CONTEXT)

    def test_hcp_filtered_ratio_100_returns_full(self):
        result = hcp_filtered(
            self._CONTEXT,
            question="Where is Mary?",
            task="qa1",
            retain_ratio=1.0,
        )
        assert result == self._CONTEXT

    def test_random_filtered_shorter(self):
        result = random_filtered(self._CONTEXT, retain_ratio=0.5, seed=42)
        assert len(result) < len(self._CONTEXT)

    def test_random_filtered_deterministic(self):
        r1 = random_filtered(self._CONTEXT, retain_ratio=0.5, seed=42)
        r2 = random_filtered(self._CONTEXT, retain_ratio=0.5, seed=42)
        assert r1 == r2

    def test_random_filtered_different_seeds(self):
        r1 = random_filtered(self._CONTEXT, retain_ratio=0.5, seed=1)
        r2 = random_filtered(self._CONTEXT, retain_ratio=0.5, seed=2)
        # With different seeds, very likely different results
        # (not guaranteed but astronomically unlikely to be the same)
        # Use a weaker assertion: at least one should differ from full context
        assert len(r1) < len(self._CONTEXT)
        assert len(r2) < len(self._CONTEXT)

    def test_hcp_filtered_preserves_order(self):
        """Filtered chunks should maintain original order."""
        result = hcp_filtered(
            self._CONTEXT,
            question="Where is Mary?",
            task="qa1",
            retain_ratio=0.5,
        )
        # Result is a string of joined chunks in original order
        assert isinstance(result, str)


# ===================================================================
# 6. REPORT
# ===================================================================


class TestReport:
    """Tests for result aggregation and reporting."""

    def _make_samples(self) -> list:
        return [
            SampleResult("qa1", "full", 1.0, "kitchen", "kitchen", "Where?", True, 1000, 1000),
            SampleResult("qa1", "full", 1.0, "bedroom", "kitchen", "Where?", False, 1000, 1000),
            SampleResult("qa1", "hcp", 0.25, "kitchen", "kitchen", "Where?", True, 1000, 250),
            SampleResult("qa1", "hcp", 0.25, "kitchen", "kitchen", "Where?", True, 1000, 250),
            SampleResult("qa1", "random", 0.25, "bedroom", "kitchen", "Where?", False, 1000, 250),
            SampleResult("qa1", "random", 0.25, "kitchen", "kitchen", "Where?", True, 1000, 250),
        ]

    def test_babilong_result_accuracy(self):
        samples = self._make_samples()
        config = BabilongConfig()
        result = BabilongResult(config=config, samples=samples)
        # 4 correct out of 6
        assert abs(result.overall_accuracy - 4 / 6) < 0.01

    def test_group_results(self):
        samples = self._make_samples()
        grouped = _group_results(samples)
        assert "qa1" in grouped
        assert "full" in grouped["qa1"]
        assert "hcp" in grouped["qa1"]
        assert "random" in grouped["qa1"]

    def test_report_accuracy_helper(self):
        samples = [
            SampleResult("qa1", "full", 1.0, "k", "k", "?", True, 100, 100),
            SampleResult("qa1", "full", 1.0, "b", "k", "?", False, 100, 100),
        ]
        assert report_accuracy(samples) == 0.5

    def test_print_report_no_crash(self, capsys):
        """print_report should run without errors."""
        samples = self._make_samples()
        config = BabilongConfig()
        result = BabilongResult(config=config, samples=samples)
        print_report(result)
        captured = capsys.readouterr()
        assert "BABILong Evaluation Report" in captured.out

    def test_save_results_json(self):
        samples = self._make_samples()
        config = BabilongConfig()
        result = BabilongResult(config=config, samples=samples)

        with tempfile.NamedTemporaryFile(suffix=".json", delete=False) as f:
            filepath = f.name

        try:
            save_results(result, filepath)
            with open(filepath, "r") as f:
                data = json.load(f)
            assert "summary" in data
            assert "per_task" in data
            assert "samples" in data
            assert data["summary"]["total_samples"] == 6
        finally:
            os.unlink(filepath)


# ===================================================================
# 7. CONFIG
# ===================================================================


class TestConfig:
    """Tests for BabilongConfig."""

    def test_default_config(self):
        config = BabilongConfig()
        assert config.n_samples == 50
        assert "qa1" in config.tasks
        assert config.model == "gpt-4o-mini"

    def test_custom_config(self):
        config = BabilongConfig(
            tasks=("qa1",),
            lengths=("4k",),
            n_samples=10,
        )
        assert config.tasks == ("qa1",)
        assert config.n_samples == 10

    def test_smoke_config(self):
        """Smoke test config should be minimal."""
        config = BabilongConfig(
            tasks=("qa1",),
            lengths=("4k",),
            n_samples=5,
            conditions=("full", "hcp"),
            retain_ratios=(0.25,),
        )
        assert config.n_samples == 5
        assert len(config.tasks) == 1
        assert len(config.lengths) == 1


# ===================================================================
# 8. GENERAL ANOMALY SCORING (domain-agnostic)
# ===================================================================


class TestGeneralAnomalyScoring:
    """Tests for the domain-agnostic general anomaly scorer."""

    _FILLER = (
        "The old gentleman regarded his companion with curiosity. "
        "He possessed a keen sense of the ridiculous aspects of life. "
        "The mansion stood silently among the ancient oaks. "
        "Its weathered facade bore the accumulated patina of centuries. "
        "The gardens stretched endlessly toward the horizon with hedgerows. "
    )
    _FACT = "Mary went to the kitchen."
    _MIXED = _FILLER + " " + _FACT

    def _make_chunks(self, text: str):
        from hcp.types import ContextChunk
        return _chunk_context(text, chunk_size=200)

    def test_corpus_stats_basic(self):
        chunks = self._make_chunks(self._MIXED)
        stats = _compute_corpus_stats(chunks)
        assert stats.mean_sentence_len > 0
        assert stats.std_sentence_len > 0
        assert stats.total_words > 0
        assert len(stats.word_freq) > 0
        assert stats.rare_threshold >= 1

    def test_corpus_stats_empty(self):
        from hcp.types import ContextChunk
        chunks = [ContextChunk(content="", position=0)]
        stats = _compute_corpus_stats(chunks)
        assert stats.total_words >= 1  # floor at 1

    def test_general_score_bounded(self):
        chunks = self._make_chunks(self._MIXED)
        stats = _compute_corpus_stats(chunks)
        for chunk in chunks:
            score = _compute_general_anomaly_score(chunk.content, stats)
            assert 0.0 <= score <= 1.0, f"Score out of bounds: {score}"

    def test_general_score_empty_text(self):
        from hcp.types import ContextChunk
        chunks = self._make_chunks(self._FILLER)
        stats = _compute_corpus_stats(chunks)
        assert _compute_general_anomaly_score("", stats) == 0.0

    def test_general_score_with_query_words(self):
        chunks = self._make_chunks(self._MIXED)
        stats = _compute_corpus_stats(chunks)
        query = frozenset({"mary", "kitchen"})
        score_with = _compute_general_anomaly_score(self._FACT, stats, query)
        score_without = _compute_general_anomaly_score(self._FACT, stats, None)
        assert score_with >= score_without, "Query-term overlap should boost score"

    def test_general_score_detects_anomaly(self):
        """Short injected fact should score differently from long prose."""
        # Use a corpus with long-ish sentences so short facts are anomalous
        long_filler = (
            "The old gentleman regarded his companion with a mixture of "
            "curiosity and satisfaction, as one who contemplated a remarkable "
            "specimen of humanity, and possessed withal a keen sense of "
            "the ridiculous aspects of the situation at hand. "
            "The mansion stood silently among the ancient oaks, its weathered "
            "facade bearing the accumulated patina of centuries, a monument "
            "to the enduring permanence of stone and mortar, and a testament "
            "to the architectural ambitions of a bygone era. "
            "The gardens stretched endlessly toward the distant horizon, their "
            "carefully maintained hedgerows forming intricate geometric patterns "
            "that delighted the eye and demonstrated the gardener's consummate "
            "skill in the ancient art of topiary. "
        )
        corpus_text = long_filler + " Bob went to the office."
        chunks = self._make_chunks(corpus_text)
        stats = _compute_corpus_stats(chunks)
        fact_score = _compute_general_anomaly_score(
            "Bob went to the office.", stats
        )
        prose_score = _compute_general_anomaly_score(
            "The elaborate ceremony proceeded with all due pomp and circumstance, "
            "gathering the assembled dignitaries in a magnificent display of "
            "unity and purpose, while the orchestra played a solemn overture.",
            stats,
        )
        # Fact should score higher (more anomalous) because it's much shorter
        # than the corpus mean sentence length
        assert fact_score > prose_score

    def test_corpus_stats_frozen(self):
        """_CorpusStats is frozen (immutable)."""
        chunks = self._make_chunks(self._FILLER)
        stats = _compute_corpus_stats(chunks)
        with pytest.raises(AttributeError):
            stats.mean_sentence_len = 999.0


# ===================================================================
# 9. BM25 FILTERING
# ===================================================================


class TestBM25Filtering:
    """Tests for BM25 retrieval baseline."""

    _CONTEXT = (
        "The old mansion stood silently among the ancient oaks. "
        "Its weathered facade bore the patina of centuries. "
        "Mary went to the kitchen. "
        "The gardens stretched endlessly toward the horizon, "
        "their carefully maintained hedgerows forming geometric patterns. "
        "Bob journeyed to the office. "
        "The library contained thousands of leather-bound volumes."
    )

    def test_bm25_returns_shorter(self):
        result = bm25_filtered(self._CONTEXT, "Where is Mary?", retain_ratio=0.5)
        assert len(result) < len(self._CONTEXT)

    def test_bm25_ratio_100_returns_full(self):
        result = bm25_filtered(self._CONTEXT, "Where is Mary?", retain_ratio=1.0)
        assert result == self._CONTEXT

    def test_bm25_returns_string(self):
        result = bm25_filtered(self._CONTEXT, "Where is Mary?", retain_ratio=0.5)
        assert isinstance(result, str)
        assert len(result) > 0

    def test_bm25_empty_context(self):
        result = bm25_filtered("", "Where is Mary?", retain_ratio=0.5)
        assert result == ""

    def test_bm25_deterministic(self):
        """BM25 should be deterministic (no randomness)."""
        r1 = bm25_filtered(self._CONTEXT, "Where is Mary?", retain_ratio=0.25)
        r2 = bm25_filtered(self._CONTEXT, "Where is Mary?", retain_ratio=0.25)
        assert r1 == r2

    def test_bm25_different_questions_differ(self):
        """Different questions should produce different selections."""
        r1 = bm25_filtered(self._CONTEXT, "Where is Mary?", retain_ratio=0.5)
        r2 = bm25_filtered(self._CONTEXT, "Where is Bob?", retain_ratio=0.5)
        # Not guaranteed to differ with short context, but should at least run
        assert isinstance(r1, str)
        assert isinstance(r2, str)


# ===================================================================
# 10. EMBEDDING FILTERING
# ===================================================================


class TestEmbeddingFiltering:
    """Tests for embedding retrieval baseline.

    NOTE: First run downloads all-MiniLM-L6-v2 (~80MB).
    """

    _CONTEXT = (
        "The old mansion stood silently among the ancient oaks. "
        "Its weathered facade bore the patina of centuries. "
        "Mary went to the kitchen. "
        "The gardens stretched endlessly toward the horizon, "
        "their carefully maintained hedgerows forming geometric patterns. "
        "Bob journeyed to the office. "
        "The library contained thousands of leather-bound volumes."
    )

    def test_embedding_returns_shorter(self):
        result = embedding_filtered(self._CONTEXT, "Where is Mary?", retain_ratio=0.5)
        assert len(result) < len(self._CONTEXT)

    def test_embedding_ratio_100_returns_full(self):
        result = embedding_filtered(self._CONTEXT, "Where is Mary?", retain_ratio=1.0)
        assert result == self._CONTEXT

    def test_embedding_returns_string(self):
        result = embedding_filtered(self._CONTEXT, "Where is Mary?", retain_ratio=0.5)
        assert isinstance(result, str)
        assert len(result) > 0

    def test_embedding_empty_context(self):
        result = embedding_filtered("", "Where is Mary?", retain_ratio=0.5)
        assert result == ""

    def test_embedding_deterministic(self):
        """Embedding retrieval should be deterministic."""
        r1 = embedding_filtered(self._CONTEXT, "Where is Mary?", retain_ratio=0.25)
        r2 = embedding_filtered(self._CONTEXT, "Where is Mary?", retain_ratio=0.25)
        assert r1 == r2
