"""Fases 0 (tarea 2) y 4 de docs/PLAN_MEJORA_RAG.md — la batería CLOSE sobre el goal raíz de cada
problema, 0 llamadas LLM: misma función (`close_goal`) y misma selección de premisas
(`select_premises`) que usa el grafo. Reanudable; una fila por (problema, variante).

Variantes:
  battery_plain       sin premisas (lo que corren `no_rag` y `baseline_auto`)
  battery_rag         premisas por goal: prior + índice denso (fusión RRF, canal C encendido) + hints + `grind +suggestions`
  battery_rag_prior   igual, sólo con el canal A (la configuración por defecto, `close.dense_channel: false`)
Una prueba cuenta si, además de cerrar el goal en el REPL, el archivo final pasa `LeanChecker.check`
(validador + kernel + auditoría de axiomas).

Uso: RAGSOLVER_MATHLIB_TABLE=mathlib_full python scripts/probe_close.py [--out data/eval/results/probe_close/battery.jsonl]
"""

from __future__ import annotations

import argparse
import json
import time
from pathlib import Path


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--problems", default="data/eval/problems/eval_v1.jsonl")
    ap.add_argument("--out", default="data/eval/results/probe_close/battery.jsonl")
    ap.add_argument("--variants", default="battery_plain,battery_rag,battery_rag_prior")
    args = ap.parse_args()

    from ragsolver.agents.agents import load_agents_config
    from ragsolver.lean.verify import LeanChecker
    from ragsolver.orchestrator.graph import close_goal, select_premises
    from ragsolver.rag import premises as P
    from ragsolver.rag.retrieve import Retriever

    cfg = load_agents_config()
    ccfg, lean = cfg["close"], cfg["lean"]
    problems = [json.loads(l) for l in open(args.problems, encoding="utf-8") if l.strip()]
    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    done = {(r["problem_id"], r["variant"]) for r in map(json.loads, out.open())} if out.exists() else set()
    variants = args.variants.split(",")

    retriever = Retriever() if "battery_rag" in variants else None
    if retriever is not None:
        retriever.trace_dir = str(out.parent / "traces")
    checker = LeanChecker(imports=lean["imports"], timeout_s=600)
    with out.open("a", encoding="utf-8") as f:
        for p in problems:
            todo = [v for v in variants if (p["problem_id"], v) not in done]
            if not todo:
                continue
            goals, errors = checker.goals(p["formal_statement"], "sorry", p.get("header", ""), lean["max_heartbeats"])
            for v in todo:
                t0 = time.time()
                block, log, names = None, [], []
                if len(goals) == 1 and not errors:
                    g = goals[0]
                    if v == "battery_plain":
                        block, log = close_goal(checker, g, ccfg["tactics"], False, [], [], ccfg["timeout_s"])
                    else:
                        names = P.prior_channel(g["goal"], k=ccfg["k_premises"]) if v == "battery_rag_prior" else select_premises(retriever, g["goal"], ccfg["k_premises"], ctx={"problem_id": p["problem_id"], "agent": "probe"}, dense_channel=True)[0]
                        block, log = close_goal(checker, g, ccfg["tactics"], True, names, P.nlinarith_hints(g["goal"]), ccfg["timeout_s"])
                ok = bool(block) and checker.check(p["formal_statement"], block, p.get("header", ""), lean["max_heartbeats"]).ok
                f.write(json.dumps({"problem_id": p["problem_id"], "variant": v, "ok": ok, "block": block, "premises": names, "tried": log, "s": round(time.time() - t0, 1)}, ensure_ascii=False) + "\n")
                f.flush()
                print(p["problem_id"], v, ok, flush=True)
    checker.close()
    rows = [json.loads(l) for l in out.open()]
    print(json.dumps({v: sum(r["ok"] for r in rows if r["variant"] == v) for v in variants}), flush=True)


if __name__ == "__main__":
    main()