#!/usr/bin/env python3
"""Self-contained tests for raglite.py.

Creates its own fixtures in temp dirs, exercises the library and the CLI,
and exits 0 only if every check passes. No network, no external deps.

Run:  python3 /root/.hermes/wizard-outputs/null/test_raglite.py
"""

import json
import subprocess
import sys
import tempfile
from pathlib import Path

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))

import raglite  # noqa: E402

FAILURES = []


def check(name, cond, detail=""):
    if cond:
        print(f"  ok   {name}")
    else:
        print(f"  FAIL {name} {detail}")
        FAILURES.append(name)


def doc_of(pid: str) -> str:
    """'path/to/codeforge.md::3' -> 'codeforge.md'"""
    return pid.split("::")[0].rsplit("/", 1)[-1]


DOCS = {
    "codeforge.md": (
        "The wizard forged the spell in the codeforge.\n"
        "\n"
        "The codeforge hums with runes and ember-light.\n"
    ),
    "kibana.md": (
        "Kibana dashboards visualize logs from the elasticsearch cluster.\n"
        "\n"
        "Operators drill into slow queries here.\n"
    ),
    "bm25.md": (
        "BM25 ranks passages by term frequency and inverse document frequency.\n"
        "\n"
        "It is a bag-of-words retrieval function.\n"
    ),
    "standup.md": (
        "The troupe stands up at dawn.\n"
        "\n"
        "Each wizard reports their lane and blockers.\n"
    ),
}


def write_corpus(root):
    for name, body in DOCS.items():
        (root / name).write_text(body, encoding="utf-8")


def test_tokenize():
    print("[tokenize]")
    toks = raglite.tokenize("Hello, WORLD! It's a BM25-test.")
    check("punctuation stripped, lowercased",
          toks == ["hello", "world", "it's", "a", "bm25", "test"], repr(toks))
    check("empty input -> empty list", raglite.tokenize("") == [])


def test_split_passages():
    print("[split_passages]")
    parts = raglite.split_passages("a\n\nb\n\n\nc\n")
    check("blank-line split", parts == ["a", "b", "c"], repr(parts))
    check("whitespace-only block dropped", raglite.split_passages("a\n   \nb") == ["a", "b"])


def test_index_and_ranking(tmp):
    print("[index + ranking]")
    idx = raglite.build_index([str(tmp)])
    check("4 docs x 2 paragraphs -> 8 passages", idx["n_docs"] == 8, idx["n_docs"])
    check("avgdl is positive", idx["avgdl"] > 0, idx["avgdl"])

    h = raglite.search(idx, "codeforge", k=3)
    check("query 'codeforge' returns hits", len(h) > 0)
    check("top hit is the codeforge passage",
          bool(h) and doc_of(h[0]["id"]) == "codeforge.md"
          and "codeforge" in h[0]["text"],
          h[0]["id"] if h else "none")

    h = raglite.search(idx, "elasticsearch logs", k=3)
    check("top hit for 'elasticsearch logs' is kibana",
          bool(h) and doc_of(h[0]["id"]) == "kibana.md", h[0]["id"] if h else "none")

    h = raglite.search(idx, "bm25 ranking", k=3)
    check("top hit for 'bm25 ranking' is bm25.md",
          bool(h) and doc_of(h[0]["id"]) == "bm25.md", h[0]["id"] if h else "none")

    h = raglite.search(idx, "nonexistentterm", k=5)
    check("unknown term -> no hits", h == [], repr(h))

    h = raglite.search(idx, "the", k=5)
    check("scores strictly positive", all(x["score"] > 0 for x in h), repr(h))
    check("scores sorted descending",
          all(h[i]["score"] >= h[i + 1]["score"] for i in range(len(h) - 1)), repr(h))

    # Negative control: a kibana query must not surface the codeforge doc.
    h = raglite.search(idx, "kibana dashboards", k=5)
    check("negative control: codeforge not ranked for kibana query",
          all(doc_of(x["id"]) != "codeforge.md" for x in h),
          repr([x["id"] for x in h]))


def test_roundtrip(tmp):
    print("[json round-trip]")
    idx = raglite.build_index([str(tmp)])
    path = tmp / "index.json"
    raglite.save_index(idx, path)
    check("index file written", path.exists() and path.stat().st_size > 0)

    loaded = raglite.load_index(path)
    a = raglite.search(idx, "codeforge", k=5)
    b = raglite.search(loaded, "codeforge", k=5)
    check("loaded index gives identical hits",
          [(x["id"], x["score"]) for x in a] == [(x["id"], x["score"]) for x in b],
          f"{a} != {b}")


def test_exact_bm25():
    """Compare against BM25 values computed independently (bm25_probe.py)."""
    print("[exact bm25 constants]")
    with tempfile.TemporaryDirectory(prefix="raglite_tiny_") as d:
        tiny = Path(d)
        (tiny / "d1.md").write_text("a a b\n", encoding="utf-8")
        (tiny / "d2.md").write_text("a c\n", encoding="utf-8")
        idx = raglite.build_index([str(tiny)])

    # Independently recomputed with the same Robertson formula (bm25_probe.py).
    EXPECTED = {
        ("d1.md", "a"): 0.24472692187108,
        ("d2.md", "a"): 0.20035335911423582,
        ("d1.md", "b"): 0.6359148445504086,
        ("d2.md", "c"): 0.7617001984175223,
    }
    for (doc, term), want in EXPECTED.items():
        hits = raglite.search(idx, term, k=10)
        got = next((x["score"] for x in hits if doc_of(x["id"]) == doc), None)
        check(f"score[{doc},{term}] == {want:.12f}",
              got is not None and abs(got - want) < 1e-12,
              f"got {got!r}")

    # Higher term frequency must win for the same term.
    hits = raglite.search(idx, "a", k=2)
    check("tf ordering: d1 (tf=2) ranks above d2 (tf=1)",
          len(hits) == 2 and doc_of(hits[0]["id"]) == "d1.md", repr(hits))


def test_cli(tmp):
    print("[cli]")
    out = tmp / "cli_index.json"
    r = subprocess.run([sys.executable, str(HERE / "raglite.py"),
                        "index", str(tmp), "--out", str(out)],
                       capture_output=True, text=True)
    check("cli index exits 0", r.returncode == 0, r.stderr[-400:])
    check("cli index reports 8 passages", "8 passages" in r.stdout, r.stdout.strip())

    r = subprocess.run([sys.executable, str(HERE / "raglite.py"),
                        "query", str(out), "elasticsearch logs", "-k", "2", "--json"],
                       capture_output=True, text=True)
    check("cli query exits 0", r.returncode == 0, r.stderr[-400:])
    try:
        hits = json.loads(r.stdout)
    except json.JSONDecodeError as e:
        check("cli query emits valid JSON", False, f"{e}: {r.stdout[:200]}")
        return
    check("cli query emits valid JSON", True)
    check("cli top hit is kibana",
          bool(hits) and doc_of(hits[0]["id"]) == "kibana.md", repr(hits))


def main():
    print("raglite test suite")
    test_tokenize()
    test_split_passages()
    with tempfile.TemporaryDirectory(prefix="raglite_test_") as d:
        tmp = Path(d)
        write_corpus(tmp)
        test_index_and_ranking(tmp)
        test_roundtrip(tmp)
        test_cli(tmp)
    test_exact_bm25()
    print()
    if FAILURES:
        print(f"FAILED ({len(FAILURES)}): " + ", ".join(FAILURES))
        return 1
    print("ALL CHECKS PASSED")
    return 0


if __name__ == "__main__":
    sys.exit(main())