"""Offline regression evaluation of fact retention, budgets and local latency.

Run: .venv/bin/python -m scripts.evaluate_context
These scores measure retained model inputs, not generated-answer accuracy.
"""

import json
import time
from pathlib import Path

from langchain_core.messages import AIMessage, HumanMessage, ToolMessage

from blueprints.conversation.context import message_tokens, trim_history

FIXTURES = Path(__file__).resolve().parents[1] / "tests/fixtures/context/scenarios.json"


def check_facts(payload: dict, case: dict) -> tuple[list[str], int]:
    missing = []
    for path, expected in case.get("required_fields", {}).items():
        actual = payload
        for key in path.split("."):
            actual = actual.get(key) if isinstance(actual, dict) else None
        if actual != expected:
            missing.append(path)
    ids = set()

    def collect(value):
        if isinstance(value, dict):
            for key, item in value.items():
                if key == "id" and isinstance(item, str):
                    ids.add(item)
                collect(item)
        elif isinstance(value, list):
            for item in value:
                collect(item)

    collect(payload)
    for source_id in case.get("required_source_ids", []):
        if source_id not in ids:
            missing.append(f"source_id:{source_id}")
    headings = (
        payload.get("data", {}).get("markdown_compaction", {}).get("headings", [])
    )
    for heading in case.get("required_headings", []):
        if heading not in headings:
            missing.append(f"heading:{heading['heading']}")
    total = sum(
        len(case.get(key, []))
        for key in ("required_fields", "required_source_ids", "required_headings")
    )
    return missing, total


def evaluate(path: Path = FIXTURES) -> list[dict]:
    results = []
    for case in json.loads(path.read_text()):
        data = dict(case["data"])
        if case["kind"] == "obsidian.note":
            data["markdown"] = (
                "# Plan\n"
                + "background details\n" * 2000
                + "## Unresolved decisions\nDiscuss savings target."
            )
        else:
            field = "transactions" if case["kind"] == "ynab.transactions" else "rows"
            data[field] = [
                {
                    "id": f"transaction-{i}",
                    "date": "2026-09-15",
                    "memo": "synthetic detail " * 100,
                }
                for i in range(100)
            ]
        artifact = {"kind": case["kind"], "data": data}
        messages = [
            HumanMessage(content="Compare this evidence with my plan"),
            AIMessage(
                content="",
                tool_calls=[{"id": "call-1", "name": "read_evidence", "args": {}}],
            ),
            ToolMessage(
                content=json.dumps(artifact), artifact=artifact, tool_call_id="call-1"
            ),
        ]
        before = sum(message_tokens(m) for m in messages)
        started = time.perf_counter()
        projected = trim_history(messages, token_budget=case["budget"])
        elapsed = (time.perf_counter() - started) * 1000
        projected_payload = json.loads(projected[-1].content)
        missing, required_count = check_facts(projected_payload, case)
        after = sum(message_tokens(m) for m in projected)
        results.append(
            {
                "scenario": case["name"],
                "input_estimated_tokens": before,
                "output_estimated_tokens": after,
                "budget": case["budget"],
                "retained_required_facts": required_count - len(missing),
                "total_required_facts": required_count,
                "missing_facts": missing,
                "local_compaction_ms": round(elapsed, 2),
                "passed": not missing and after <= case["budget"],
            }
        )
    return results


if __name__ == "__main__":
    report = evaluate()
    print(json.dumps(report, indent=2))
    raise SystemExit(0 if all(case["passed"] for case in report) else 1)