from __future__ import annotations

import json
from pathlib import Path
from typing import Any

from bayesian_langgraph import BayesianIndex, BayesianRAGGraph
from evaluation.rag_benchmark import CLEAN_CASES, CLEAN_DOCS, CONFLICT_CASES, CONFLICT_DOCS


def _case_ok(out: dict[str, Any], case: dict[str, Any]) -> bool:
    if case["abstain"]:
        return bool(out.get("abstained")) and not out.get("citations")
    return (
        not out.get("abstained")
        and case["doc"] in out.get("citations", [])
        and str(case["gold"]).lower() in str(out.get("answer", "")).lower()
    )


def _run_suite(docs, cases) -> dict[str, Any]:
    rag = BayesianRAGGraph(BayesianIndex(docs), k=4)
    rows = []
    for case in cases:
        out = rag.invoke(case["q"])
        rows.append({
            "id": case["id"],
            "ok": _case_ok(out, case),
            "answer": out["answer"],
            "citations": out["citations"],
            "abstained": out["abstained"],
            "top_doc": out["hits"][0]["id"] if out["hits"] else None,
        })
    return {
        "n_cases": len(rows),
        "passed": sum(r["ok"] for r in rows),
        "accuracy": sum(r["ok"] for r in rows) / len(rows),
        "rows": rows,
    }


def run_framework_benchmark(output_path: str | Path | None = None) -> dict[str, Any]:
    rag = BayesianRAGGraph(BayesianIndex(CLEAN_DOCS), k=4)
    graph = rag.graph.get_graph()
    result = {
        "framework": "langgraph",
        "stategraph_nodes": sorted(graph.nodes.keys()),
        "stategraph_edges": sorted([e.source, e.target] for e in graph.edges),
        "clean": _run_suite(CLEAN_DOCS, CLEAN_CASES),
        "conflict": _run_suite(CONFLICT_DOCS, CONFLICT_CASES),
    }
    if output_path is not None:
        Path(output_path).write_text(json.dumps(result, indent=2), encoding="utf-8")
    return result


if __name__ == "__main__":
    print(json.dumps(run_framework_benchmark(), indent=2))