"""deg-suite v1 aggregation, validity gate, judge cache and training guard (CPU only)."""

from __future__ import annotations

import json

import pytest

from deg.bench import suite
from deg.bench.score import Prediction
from deg.bench.suite import SOURCES, Source, SuiteGuard, aggregate_suite, score_source
from deg.bench.suite_judge import Item, Judge, grade_score, parse_grade


def test_headline_is_mean_of_capability_means():
    comps = {"decide": {"a": 1.0, "b": 0.0, "c": 0.5, "d": 0.5},  # 4 sources -> 0.5
             "extract": {"a": 0.9},                                # 1 source  -> 0.9
             "generate": {"a": 0.2, "b": 0.4}}                     # 2 sources -> 0.3
    out = aggregate_suite(comps)
    assert out["capabilities"] == pytest.approx({"decide": 0.5, "extract": 0.9, "generate": 0.3})
    assert out["headline"] == pytest.approx((0.5 + 0.9 + 0.3) / 3)


def test_adding_a_source_cannot_swamp_other_capabilities():
    base = {"decide": {"a": 0.8}, "extract": {"a": 0.6}, "generate": {"a": 0.4}}
    more = {**base, "decide": {"a": 0.8, **{f"x{i}": 1.0 for i in range(50)}}}
    h0, h1 = aggregate_suite(base)["headline"], aggregate_suite(more)["headline"]
    # 50 perfect decide sources move only the DECIDE third of the headline
    assert h1 - h0 <= (1.0 - 0.8) / 3 + 1e-9
    assert aggregate_suite(more)["capabilities"]["extract"] == pytest.approx(0.6)


def test_missing_capability_gives_no_headline():
    out = aggregate_suite({"decide": {"a": 0.9, "b": None}, "extract": {}, "generate": {}})
    assert out["headline"] is None
    assert out["capabilities"]["decide"] == pytest.approx(0.9)
    assert out["missing_capabilities"] == ["extract", "generate"]


def test_component_out_of_range_rejected():
    with pytest.raises(ValueError):
        aggregate_suite({"decide": {"a": 1.5}, "extract": {"a": 0}, "generate": {"a": 0}})


def test_every_capability_has_sources():
    for cap in suite.CAPABILITIES:
        assert sum(cap in s.caps for s in SOURCES) >= 2


def test_jev_intelligence_chance_corrected_and_tier_weighted():
    # easy: 2/2 correct on noul (chance .5) -> 1.0; hard: 1/2 on noul -> 0.0
    rows = [("easy", True, 2), ("easy", True, 2), ("hard", True, 2), ("hard", False, 2)]
    w = suite.JEV_TIER_WEIGHTS
    assert suite.jev_intelligence(rows) == pytest.approx(w["easy"] / (w["easy"] + w["hard"]))
    assert suite.jev_intelligence([("standard", False, 4)]) == 0.0  # clipped at 0


SCHEMA = {"type": "object", "properties": {
    "kind": {"enum": ["a", "b"]},
    "name": {"type": "string"},
    "items": {"type": "array", "items": {"type": "string"}},
    "note": {"type": "string", "x-generate": True, "maxLength": 20}},
    "required": ["kind", "name", "items", "note"], "additionalProperties": False}
REC = {"uid": "r1", "context": "Name: Ada. Items: x, y.", "schema": SCHEMA,
       "value": {"kind": "a", "name": "Ada", "items": ["x", "y"], "note": "ok"}}


class FixedJudge:
    model = "fixed"
    calls = 0

    def score_many(self, items):
        return [0.5 for _ in items]


def test_fields_scores_and_validity_gate():
    src = Source("t", "fields", ("decide", "extract", "generate"))
    good = Prediction("r1", value={"kind": "a", "name": "Ada", "items": ["x", "y"],
                                   "note": "fine"})
    out = score_source(src, [REC], [good], ("decide", "extract", "generate"), FixedJudge())
    assert out["components"]["decide"] == 1.0
    assert out["components"]["extract"] == 1.0  # mean(norm 1.0, count 1.0)
    assert out["components"]["generate"] == 0.5
    # the same correct fields with one schema violation (note over maxLength): all zero
    bad = Prediction("r1", value={"kind": "a", "name": "Ada", "items": ["x", "y"],
                                  "note": "x" * 30})
    out = score_source(src, [REC], [bad], ("decide", "extract", "generate"), FixedJudge())
    assert out["valid_rate"] == 0.0
    assert out["components"] == {"decide": 0.0, "extract": 0.0, "generate": 0.0}
    err = Prediction("r1", error="OutOfMemoryError: ...")
    out = score_source(src, [REC], [err], ("decide", "extract", "generate"), FixedJudge())
    assert out["components"] == {"decide": 0.0, "extract": 0.0, "generate": 0.0}


def test_decide_only_system_is_not_gated_by_validity():
    src = Source("t", "fields", ("decide", "extract", "generate"))
    out = score_source(src, [REC], [Prediction("r1", value={"kind": "a"})], ("decide",), None)
    assert out["components"] == {"decide": 1.0}


def test_squad_balances_answerable_and_unanswerable():
    rec = {"uid": "s", "context": "Paris is in France.",
           "schema": {"type": "object", "properties": {
               "q1": {"type": ["string", "null"]}, "q2": {"type": ["string", "null"]}},
               "required": ["q1", "q2"], "additionalProperties": False},
           "value": {"q1": "France", "q2": None}, "alternatives": {"/q1": ["France"], "/q2": []}}
    src = Source("squad2", "squad", ("extract",))
    always_null = Prediction("s", value={"q1": None, "q2": None})
    out = score_source(src, [rec], [always_null], ("decide", "extract", "generate"), None)
    assert out["components"]["extract"] == pytest.approx(0.5)
    right = Prediction("s", value={"q1": "france.", "q2": None})
    out = score_source(src, [rec], [right], ("decide", "extract", "generate"), None)
    assert out["components"]["extract"] == pytest.approx(1.0)


def test_judge_cache_is_deterministic(tmp_path):
    calls = []

    def call(prompt):
        calls.append(prompt)
        return '{"faithful": 2, "responsive": 1}'

    item = Item("ctx", "note", "a note", 40, "text")
    j = Judge(tmp_path / "c.jsonl", call=call, model="m")
    assert j.score_many([item, item]) == [0.75, 0.75]
    j2 = Judge(tmp_path / "c.jsonl", call=lambda p: pytest.fail("cache miss"), model="m")
    assert j2.score_many([item]) == [0.75]
    assert len(calls) <= 2  # concurrent duplicates may both call; the cache keeps the result
    other = Item("ctx", "note", "a note", 40, "other text")
    assert other.key("m") != item.key("m") and item.key("m") != item.key("m2")


def test_parse_grade():
    assert grade_score(parse_grade('ok {"faithful": 1, "responsive": 2}')) == 0.75
    with pytest.raises(ValueError):
        parse_grade('{"faithful": 3, "responsive": 2}')


def test_suite_guard_and_path_guard():
    long_ctx = ("The quarterly report says revenue grew eleven percent while operating "
                "margins narrowed because of freight costs and a weaker euro this year")
    g = SuiteGuard(texts=[long_ctx, "How do I locate my card?"])
    assert g.match("prefix text. " + long_ctx + " suffix") is not None
    assert g.match("how do I locate my card") is not None
    assert g.match("Where can I locate my new card after moving house last week?") is None
    assert g.match("A completely unrelated document about gardening tools.") is None
    with pytest.raises(ValueError):
        suite.guard_path("/data/deg/suite-v1/mixed.jsonl")
    suite.guard_path("/data/deg/datagen/full-off/records.jsonl")


def test_frozen_files_match_manifest_when_present():
    if not (suite.RESULTS / "MANIFEST.json").exists():
        pytest.skip("suite not built on this host")
    for name in json.loads((suite.RESULTS / "MANIFEST.json").read_text())["sources"]:
        if (suite.SUITE_DIR / f"{name}.jsonl").exists():
            assert suite.load_source(name)


def test_missing_source_withholds_headline(tmp_path, monkeypatch):
    """A built source without predictions must not silently reweight its capability."""
    recs = {"datagen-test": [REC]}
    monkeypatch.setattr(suite, "manifest", lambda: {"sources": {
        "datagen-test": {"sha256": "x"}, "extractbench": {"sha256": "y"}}})
    monkeypatch.setattr(suite, "load_source", lambda name, verify=True: recs[name])
    sdir = tmp_path / "sys"
    sdir.mkdir()
    good = Prediction("r1", value={"kind": "a", "name": "Ada", "items": ["x", "y"],
                                   "note": "fine"})
    (sdir / "datagen-test.preds.jsonl").write_text(json.dumps(good.to_json()) + "\n")
    card = suite.score_system(tmp_path, "sys", ("decide", "extract", "generate"), FixedJudge())
    assert card["incomplete_sources"] == ["extractbench"]
    assert card["headline"] is None and card["capabilities"]["extract"] is None
    assert card["capabilities"]["decide"] == 1.0