"""Unit tests for src/generate.py — gate, context packing, think-splitting, and the
refusal/answer/stream paths with the LLM and retrieval monkeypatched out."""
import pytest

import src.generate as gen


# ---------- pure helpers ----------

def test_split_think_no_tags():
    ans, think = gen._split_think("just an answer")
    assert ans == "just an answer" and think == ""


def test_split_think_extracts_and_strips():
    raw = "<think>step 1\nstep 2</think>The answer is $\\lambda$."
    ans, think = gen._split_think(raw)
    assert ans == "The answer is $\\lambda$."
    assert "step 1" in think and "step 2" in think


def test_split_think_multiple_blocks():
    raw = "<think>a</think>mid<think>b</think>end"
    ans, think = gen._split_think(raw)
    assert ans == "midend"
    assert think == "a\nb"


def test_truncate_under_cap_unchanged():
    assert gen._truncate("short", 100) == "short"


def test_truncate_over_cap_marks_and_bounds():
    text = "line\n" * 2000
    out = gen._truncate(text, 100)
    assert out.endswith("... [passage truncated]")
    assert len(out) <= 100 * gen._CHARS_PER_TOK + len("\n... [passage truncated]")


def _hit(i, *, page=None, score=0.5, text="x" * 200):
    return {"parent_text": text, "chunk_id": f"ck{i}", "source": "book",
            "chapter": f"Chapter {i}", "section": "Sec", "page": page,
            "rerank_score": score, "parent_id": f"p{i}"}


def test_build_context_tags_and_metadata():
    ctx, sources = gen._build_context([_hit(1, page=42), _hit(2)])
    assert "[S1] (Chapter 1 / Sec)" in ctx and "[S2]" in ctx
    assert sources[0] == {"tag": "S1", "chunk_id": "ck1", "location": "Chapter 1 / Sec",
                          "page": 42, "score": 0.5}
    assert sources[1]["page"] is None


def test_build_context_location_falls_back_to_source():
    h = _hit(1)
    h["chapter"] = h["section"] = None
    _, sources = gen._build_context([h])
    assert sources[0]["location"] == "book"


def test_build_context_stops_at_token_budget():
    big = "y" * (gen.MAX_PARENT_TOKENS * gen._CHARS_PER_TOK)  # each hit ~MAX_PARENT_TOKENS
    hits = [_hit(i, text=big) for i in range(10)]
    ctx, sources = gen._build_context(hits)
    assert 0 < len(sources) < 10                     # budget cut it off
    est = sum(len(b) for b in ctx.split("\n\n")) // gen._CHARS_PER_TOK
    assert est <= gen.CTX_TOKEN_BUDGET + gen.MAX_PARENT_TOKENS  # never wildly over


# ---------- decompose gate ----------

@pytest.mark.parametrize("q", [
    "Compare the binomial and geometric variance",
    "What is the difference between dense and sparse?",
    "Poisson versus binomial mean",
    "Derive the variance of a binomial from its mgf",
])
def test_gate_matches_multi_part_questions(q):
    assert gen.DECOMPOSE_GATE.search(q)


@pytest.mark.parametrize("q", [
    "What is the variance of the Poisson distribution?",
    "State Bayes' theorem.",
    "Derive the variance of the binomial.",   # derive WITHOUT from/via stays single-hop
])
def test_gate_ignores_single_hop_questions(q):
    assert not gen.DECOMPOSE_GATE.search(q)


def test_decompose_skips_llm_when_gate_misses(monkeypatch):
    def boom(messages, entry):
        raise AssertionError("LLM must not be called for single-hop questions")
    monkeypatch.setattr(gen, "_chat", boom)
    assert gen.decompose("What is the variance of the Poisson?") == []


def test_decompose_strips_bullets_and_caps_at_four(monkeypatch):
    raw = "1. mean of A\n- variance of A\n\n2) mean of B\nvariance of B\nextra query five\n"
    monkeypatch.setattr(gen, "_chat", lambda m, entry: raw)
    subs = gen.decompose("compare A and B")
    assert subs == ["mean of A", "variance of A", "mean of B", "variance of B"]


def test_decompose_fails_open_on_llm_error(monkeypatch):
    def boom(messages, entry):
        raise OSError("ollama down")
    monkeypatch.setattr(gen, "_chat", boom)
    assert gen.decompose("compare A and B") == []


# ---------- answer / stream paths ----------

def test_answer_refuses_on_no_hits(monkeypatch):
    monkeypatch.setattr(gen, "_retrieve_for", lambda c, q, k: [])
    out = gen.answer(None, "q")
    assert out["refused"] and out["sources"] == []
    assert out["answer"] == "I don't know based on the provided material."


def test_answer_refuses_below_threshold(monkeypatch):
    monkeypatch.setattr(gen, "_retrieve_for",
                        lambda c, q, k: [_hit(1, score=gen.MIN_RERANK_SCORE / 2)])
    assert gen.answer(None, "q")["refused"]


def test_answer_grounds_and_splits_thinking(monkeypatch):
    monkeypatch.setattr(gen, "_retrieve_for", lambda c, q, k: [_hit(1, page=7)])
    monkeypatch.setattr(gen, "_chat", lambda m, entry: "<think>hmm</think>It is $\\mu$. [S1]")
    out = gen.answer(None, "q", show_thinking=True)
    assert not out["refused"]
    assert out["answer"] == "It is $\\mu$. [S1]"
    assert out["thinking"] == "hmm"
    assert out["sources"][0]["tag"] == "S1"
    # without the flag, thinking never leaks
    out2 = gen.answer(None, "q")
    assert "thinking" not in out2


def test_stream_event_order_and_final_answer(monkeypatch):
    monkeypatch.setattr(gen, "_retrieve_for", lambda c, q, k: [_hit(1)])
    monkeypatch.setattr(gen, "_chat_stream", lambda m, entry: iter(["The ", "answer."]))
    events = list(gen.stream(None, "q"))
    assert [e["type"] for e in events] == ["sources", "token", "token", "done"]
    assert events[-1]["answer"] == "The answer."
    assert not events[-1]["refused"]


def test_stream_refusal_is_single_done_event(monkeypatch):
    monkeypatch.setattr(gen, "_retrieve_for", lambda c, q, k: [])
    events = list(gen.stream(None, "q"))
    assert len(events) == 1 and events[0]["type"] == "done" and events[0]["refused"]


def test_stream_style_selects_prompt(monkeypatch):
    seen = {}
    monkeypatch.setattr(gen, "_retrieve_for", lambda c, q, k: [_hit(1)])

    def fake_stream(messages, entry):
        seen["system"] = messages[0]["content"]
        return iter(["ok"])
    monkeypatch.setattr(gen, "_chat_stream", fake_stream)
    list(gen.stream(None, "q", style="eli5"))
    assert seen["system"] == gen.ELI5_SYSTEM
    list(gen.stream(None, "q", style="nonsense"))
    assert seen["system"] == gen.SYSTEM          # unknown style falls back to default