"""Tests for rbsp.query.bm25f."""

from __future__ import annotations

import math

import pytest

from rbsp.query.bm25f import BM25FScorer, FieldConfig


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

def _make_scorer(
    num_docs: int = 100,
    fields: list[FieldConfig] | None = None,
    field_avg_lengths: dict[str, float] | None = None,
    doc_freqs: dict[str, int] | None = None,
    k1: float = 1.2,
) -> BM25FScorer:
    """Create a scorer with corpus stats pre-loaded."""
    scorer = BM25FScorer(k1=k1, fields=fields)
    scorer.update_stats(
        num_docs=num_docs,
        field_avg_lengths=field_avg_lengths or {"filename": 5.0, "content": 100.0},
        doc_freqs=doc_freqs or {},
    )
    return scorer


def _idf(n: int, df: int) -> float:
    """Reference IDF calculation matching the scorer."""
    return math.log((n - df + 0.5) / (df + 0.5) + 1.0)


# ---------------------------------------------------------------------------
# FieldConfig
# ---------------------------------------------------------------------------

class TestFieldConfig:
    def test_defaults(self):
        fc = FieldConfig("body")
        assert fc.name == "body"
        assert fc.weight == 1.0
        assert fc.b == 0.75

    def test_custom(self):
        fc = FieldConfig("title", weight=5.0, b=0.2)
        assert fc.name == "title"
        assert fc.weight == 5.0
        assert fc.b == 0.2

    def test_frozen(self):
        fc = FieldConfig("x")
        with pytest.raises(AttributeError):
            fc.name = "y"

    def test_hashable(self):
        fc = FieldConfig("x", weight=1.0, b=0.5)
        assert hash(fc) == hash(FieldConfig("x", weight=1.0, b=0.5))

    def test_equality(self):
        a = FieldConfig("title", weight=2.0, b=0.3)
        b = FieldConfig("title", weight=2.0, b=0.3)
        assert a == b


# ---------------------------------------------------------------------------
# BM25FScorer — construction and stats
# ---------------------------------------------------------------------------

class TestBM25FScorerConstruction:
    def test_default_fields(self):
        scorer = BM25FScorer()
        names = [f.name for f in scorer.fields]
        assert "filename" in names
        assert "content" in names

    def test_custom_fields(self):
        fields = [FieldConfig("a"), FieldConfig("b")]
        scorer = BM25FScorer(fields=fields)
        assert len(scorer.fields) == 2
        assert scorer.fields[0].name == "a"

    def test_default_k1(self):
        scorer = BM25FScorer()
        assert scorer.k1 == 1.2

    def test_custom_k1(self):
        scorer = BM25FScorer(k1=2.0)
        assert scorer.k1 == 2.0

    def test_num_docs_initial(self):
        scorer = BM25FScorer()
        assert scorer.num_docs == 0

    def test_update_stats(self):
        scorer = BM25FScorer()
        scorer.update_stats(
            num_docs=50,
            field_avg_lengths={"filename": 3.0, "content": 80.0},
            doc_freqs={"foo": 10},
        )
        assert scorer.num_docs == 50
        assert scorer.field_avg_lengths == {"filename": 3.0, "content": 80.0}

    def test_update_stats_clamps_zero_avgdl(self):
        scorer = BM25FScorer()
        scorer.update_stats(
            num_docs=10,
            field_avg_lengths={"filename": 0.0, "content": 50.0},
            doc_freqs={},
        )
        # Zero avg length is clamped to 1.0 to avoid division by zero
        assert scorer.field_avg_lengths["filename"] == 1.0
        assert scorer.field_avg_lengths["content"] == 50.0


# ---------------------------------------------------------------------------
# BM25FScorer — IDF
# ---------------------------------------------------------------------------

class TestBM25FIDF:
    def test_rare_term_high_idf(self):
        scorer = _make_scorer(num_docs=1000, doc_freqs={"rare": 1})
        assert scorer.idf("rare") > 5.0

    def test_common_term_low_idf(self):
        scorer = _make_scorer(num_docs=1000, doc_freqs={"common": 900})
        assert scorer.idf("common") < 1.0

    def test_unknown_term_max_idf(self):
        scorer = _make_scorer(num_docs=100, doc_freqs={})
        expected = _idf(100, 0)
        assert scorer.idf("unknown") == pytest.approx(expected)

    def test_idf_matches_reference(self):
        scorer = _make_scorer(num_docs=200, doc_freqs={"x": 50})
        assert scorer.idf("x") == pytest.approx(_idf(200, 50))

    def test_idf_non_negative(self):
        # Even when df == N, IDF should be non-negative due to +1 smoothing
        scorer = _make_scorer(num_docs=10, doc_freqs={"everywhere": 10})
        assert scorer.idf("everywhere") >= 0.0


# ---------------------------------------------------------------------------
# BM25FScorer — scoring
# ---------------------------------------------------------------------------

class TestBM25FScoring:
    def test_zero_score_no_match(self):
        scorer = _make_scorer(doc_freqs={"foo": 10})
        score = scorer.score(
            query_terms=["foo"],
            doc_fields={"filename": {}, "content": {}},
            doc_field_lengths={"filename": 3, "content": 50},
        )
        assert score == 0.0

    def test_positive_score_on_match(self):
        scorer = _make_scorer(doc_freqs={"foo": 10})
        score = scorer.score(
            query_terms=["foo"],
            doc_fields={"content": {"foo": 3}},
            doc_field_lengths={"content": 50},
        )
        assert score > 0.0

    def test_filename_match_scores_higher(self):
        scorer = _make_scorer(doc_freqs={"foo": 10})
        # Match in filename (weight=3.0)
        score_fn = scorer.score(
            query_terms=["foo"],
            doc_fields={"filename": {"foo": 1}},
            doc_field_lengths={"filename": 5},
        )
        # Match in content (weight=1.0)
        score_ct = scorer.score(
            query_terms=["foo"],
            doc_fields={"content": {"foo": 1}},
            doc_field_lengths={"content": 100},
        )
        assert score_fn > score_ct

    def test_multiple_field_matches_add(self):
        scorer = _make_scorer(doc_freqs={"foo": 10})
        # Match in both fields
        score_both = scorer.score(
            query_terms=["foo"],
            doc_fields={"filename": {"foo": 1}, "content": {"foo": 2}},
            doc_field_lengths={"filename": 5, "content": 100},
        )
        # Match in content only
        score_content = scorer.score(
            query_terms=["foo"],
            doc_fields={"content": {"foo": 2}},
            doc_field_lengths={"content": 100},
        )
        assert score_both > score_content

    def test_multi_term_query(self):
        scorer = _make_scorer(doc_freqs={"foo": 10, "bar": 20})
        score_one = scorer.score(
            query_terms=["foo"],
            doc_fields={"content": {"foo": 1, "bar": 1}},
            doc_field_lengths={"content": 50},
        )
        score_two = scorer.score(
            query_terms=["foo", "bar"],
            doc_fields={"content": {"foo": 1, "bar": 1}},
            doc_field_lengths={"content": 50},
        )
        assert score_two > score_one

    def test_higher_tf_increases_score(self):
        scorer = _make_scorer(doc_freqs={"foo": 10})
        score_low = scorer.score(
            query_terms=["foo"],
            doc_fields={"content": {"foo": 1}},
            doc_field_lengths={"content": 100},
        )
        score_high = scorer.score(
            query_terms=["foo"],
            doc_fields={"content": {"foo": 10}},
            doc_field_lengths={"content": 100},
        )
        assert score_high > score_low

    def test_tf_saturation(self):
        scorer = _make_scorer(doc_freqs={"foo": 10})
        score_100 = scorer.score(
            query_terms=["foo"],
            doc_fields={"content": {"foo": 100}},
            doc_field_lengths={"content": 200},
        )
        score_1000 = scorer.score(
            query_terms=["foo"],
            doc_fields={"content": {"foo": 1000}},
            doc_field_lengths={"content": 2000},
        )
        # Scores should converge — the ratio should be close to 1
        ratio = score_1000 / score_100 if score_100 > 0 else float("inf")
        assert 0.5 < ratio < 2.0

    def test_length_normalization(self):
        scorer = _make_scorer(
            doc_freqs={"foo": 10},
            field_avg_lengths={"content": 100.0},
            fields=[FieldConfig("content", weight=1.0, b=0.75)],
        )
        # Short document scores higher than long one (same TF)
        score_short = scorer.score(
            query_terms=["foo"],
            doc_fields={"content": {"foo": 1}},
            doc_field_lengths={"content": 20},
        )
        score_long = scorer.score(
            query_terms=["foo"],
            doc_fields={"content": {"foo": 1}},
            doc_field_lengths={"content": 500},
        )
        assert score_short > score_long

    def test_b_zero_disables_normalization(self):
        scorer = _make_scorer(
            doc_freqs={"foo": 10},
            field_avg_lengths={"content": 100.0},
            fields=[FieldConfig("content", weight=1.0, b=0.0)],
        )
        # With b=0, document length should not matter (same TF)
        score_short = scorer.score(
            query_terms=["foo"],
            doc_fields={"content": {"foo": 5}},
            doc_field_lengths={"content": 10},
        )
        score_long = scorer.score(
            query_terms=["foo"],
            doc_fields={"content": {"foo": 5}},
            doc_field_lengths={"content": 1000},
        )
        assert score_short == pytest.approx(score_long)

    def test_missing_field_skipped(self):
        scorer = _make_scorer(doc_freqs={"foo": 10})
        # Only provide content, not filename
        score = scorer.score(
            query_terms=["foo"],
            doc_fields={"content": {"foo": 1}},
            doc_field_lengths={"content": 100},
        )
        assert score > 0.0

    def test_empty_query(self):
        scorer = _make_scorer(doc_freqs={"foo": 10})
        score = scorer.score(
            query_terms=[],
            doc_fields={"content": {"foo": 1}},
            doc_field_lengths={"content": 100},
        )
        assert score == 0.0

    def test_empty_document(self):
        scorer = _make_scorer(doc_freqs={"foo": 10})
        score = scorer.score(
            query_terms=["foo"],
            doc_fields={},
            doc_field_lengths={},
        )
        assert score == 0.0


# ---------------------------------------------------------------------------
# BM25FScorer — upper bound
# ---------------------------------------------------------------------------

class TestBM25FUpperBound:
    def test_upper_bound_exceeds_score(self):
        scorer = _make_scorer(doc_freqs={"foo": 10, "bar": 20})
        actual = scorer.score(
            query_terms=["foo", "bar"],
            doc_fields={"content": {"foo": 1, "bar": 2}},
            doc_field_lengths={"content": 100},
        )
        upper = scorer.score_upper_bound(["foo", "bar"])
        assert upper >= actual

    def test_upper_bound_positive(self):
        scorer = _make_scorer(num_docs=100, doc_freqs={"x": 5})
        assert scorer.score_upper_bound(["x"]) > 0.0

    def test_upper_bound_empty_query(self):
        scorer = _make_scorer()
        assert scorer.score_upper_bound([]) == 0.0

    def test_upper_bound_sum_of_idf(self):
        scorer = _make_scorer(num_docs=100, doc_freqs={"a": 10, "b": 20})
        expected = scorer.idf("a") + scorer.idf("b")
        assert scorer.score_upper_bound(["a", "b"]) == pytest.approx(expected)


# ---------------------------------------------------------------------------
# BM25FScorer — normalized score
# ---------------------------------------------------------------------------

class TestBM25FNormalized:
    def test_in_unit_range(self):
        scorer = _make_scorer(doc_freqs={"foo": 10})
        norm = scorer.score_normalized(
            query_terms=["foo"],
            doc_fields={"content": {"foo": 3}},
            doc_field_lengths={"content": 100},
        )
        assert 0.0 <= norm <= 1.0

    def test_perfect_match_near_one(self):
        scorer = _make_scorer(doc_freqs={"foo": 1})
        norm = scorer.score_normalized(
            query_terms=["foo"],
            doc_fields={"filename": {"foo": 100}, "content": {"foo": 100}},
            doc_field_lengths={"filename": 100, "content": 100},
        )
        # With very high TF in all fields, should be close to 1.0
        assert norm > 0.5

    def test_zero_for_no_match(self):
        scorer = _make_scorer(doc_freqs={"foo": 10})
        norm = scorer.score_normalized(
            query_terms=["foo"],
            doc_fields={"content": {}},
            doc_field_lengths={"content": 100},
        )
        assert norm == 0.0

    def test_zero_for_empty_query(self):
        scorer = _make_scorer()
        norm = scorer.score_normalized(
            query_terms=[],
            doc_fields={"content": {"foo": 1}},
            doc_field_lengths={"content": 100},
        )
        assert norm == 0.0


# ---------------------------------------------------------------------------
# BM25FScorer — weighted TF edge cases
# ---------------------------------------------------------------------------

class TestBM25FWeightedTF:
    def test_custom_weights_affect_ranking(self):
        fields = [
            FieldConfig("title", weight=10.0, b=0.3),
            FieldConfig("body", weight=1.0, b=0.75),
        ]
        scorer = _make_scorer(
            fields=fields,
            field_avg_lengths={"title": 5.0, "body": 200.0},
            doc_freqs={"foo": 10},
        )
        # Same TF in title vs body — title should dominate
        score_title = scorer.score(
            query_terms=["foo"],
            doc_fields={"title": {"foo": 1}},
            doc_field_lengths={"title": 5},
        )
        score_body = scorer.score(
            query_terms=["foo"],
            doc_fields={"body": {"foo": 1}},
            doc_field_lengths={"body": 200},
        )
        assert score_title > score_body  # title weight dominates

    def test_single_field_scorer(self):
        fields = [FieldConfig("text", weight=1.0, b=0.75)]
        scorer = _make_scorer(
            fields=fields,
            field_avg_lengths={"text": 50.0},
            doc_freqs={"hello": 5},
        )
        score = scorer.score(
            query_terms=["hello"],
            doc_fields={"text": {"hello": 2}},
            doc_field_lengths={"text": 50},
        )
        assert score > 0.0

    def test_three_fields(self):
        fields = [
            FieldConfig("title", weight=5.0, b=0.3),
            FieldConfig("body", weight=1.0, b=0.75),
            FieldConfig("tags", weight=3.0, b=0.0),
        ]
        scorer = _make_scorer(
            fields=fields,
            field_avg_lengths={"title": 5.0, "body": 100.0, "tags": 3.0},
            doc_freqs={"foo": 10},
        )
        score = scorer.score(
            query_terms=["foo"],
            doc_fields={
                "title": {"foo": 1},
                "body": {"foo": 2},
                "tags": {"foo": 1},
            },
            doc_field_lengths={"title": 5, "body": 100, "tags": 3},
        )
        assert score > 0.0


# ---------------------------------------------------------------------------
# BM25FScorer — consistency with BM25
# ---------------------------------------------------------------------------

class TestBM25FConsistency:
    def test_single_field_matches_bm25_formula(self):
        """With a single field of weight=1.0, BM25F should behave like BM25."""
        k1 = 1.2
        b = 0.75
        fields = [FieldConfig("text", weight=1.0, b=b)]
        scorer = _make_scorer(
            k1=k1,
            fields=fields,
            num_docs=100,
            field_avg_lengths={"text": 50.0},
            doc_freqs={"foo": 10},
        )

        tf = 3
        dl = 60
        avgdl = 50.0

        actual = scorer.score(
            query_terms=["foo"],
            doc_fields={"text": {"foo": tf}},
            doc_field_lengths={"text": dl},
        )

        # Compute BM25F by hand for single field
        idf_val = _idf(100, 10)
        norm = 1.0 - b + b * dl / avgdl
        tf_hat = 1.0 * tf / norm  # weight=1.0
        expected = idf_val * tf_hat / (k1 + tf_hat)

        assert actual == pytest.approx(expected)
