#!/usr/bin/env python3
"""
Run :class:`~frame.pipeline.formal_model_building_pipeline.FormalModelBuildingPipeline`
for a range of ``spec_id`` rows in ``datasets/spec.csv``.

Uses UPPAAL IR KB for the Builder **unless** ``--no-rag`` is set.

Writes:
  artifacts/formal_build/batch_formal_<kind>_<stamp>/report.jsonl  (one JSON object per spec)
  artifacts/formal_build/batch_formal_<kind>_<stamp>/S{nn}_{label}/... per spec

Usage:
  python scripts/run_formal_batch_kb.py --start 1 --end 20 --no-rag
  python scripts/run_formal_batch_kb.py --start 1 --end 20 --rebuild-kb-index   # KB on
"""
from __future__ import annotations

import argparse
import csv
import json
import re
import subprocess
import sys
import time
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Tuple

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))

from frame.pipeline.formal_model_building_pipeline import FormalModelBuildingPipeline

SPEC_CSV = ROOT / "datasets" / "spec.csv"


def _safe_label(s: str) -> str:
    s = (s or "").strip() or "spec"
    return re.sub(r"[^\w\-]+", "_", s)[:80]


def load_all_specs(spec_csv: Path) -> List[Dict[str, str]]:
    rows: List[Dict[str, str]] = []
    with spec_csv.open(encoding="utf-8", newline="") as f:
        for row in csv.DictReader(f):
            rows.append(row)
    rows.sort(key=lambda r: int(r["spec_id"]))
    return rows


def maybe_rebuild_kb_index() -> None:
    script = ROOT / "scripts" / "index_uppaal_kb.py"
    if not script.is_file():
        return
    subprocess.run([sys.executable, str(script), "--rebuild"], cwd=str(ROOT), check=False)


def quality_grade(rec: Dict[str, Any]) -> str:
    ok = bool(rec.get("summary_ok"))
    val = bool(rec.get("validator_ok"))
    deg = bool(rec.get("degraded_export"))
    smoke = rec.get("smoke_ok")
    ir_n = int(rec.get("ir_error_count_final") or 0)
    if ok and val and smoke is True and not deg:
        return "excellent"
    if ok and smoke is True and (not deg or ir_n <= 3):
        return "good"
    if ok and smoke is True:
        return "acceptable"
    if ok and smoke is False:
        return "xml_ok_verifyta_bad"
    if not ok:
        return "failed_export"
    return "unknown"


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--spec-csv", type=Path, default=SPEC_CSV)
    ap.add_argument("--start", type=int, default=1)
    ap.add_argument("--end", type=int, default=20)
    ap.add_argument("--batch-root", type=Path, default=None, help="Reuse an existing batch root instead of creating a new timestamped one")
    ap.add_argument("--rebuild-kb-index", action="store_true")
    ap.add_argument(
        "--no-rag",
        action="store_true",
        help="Disable UPPAAL IR KB retrieval for the Builder (baseline / no-RAG run).",
    )
    ap.add_argument("--analyze-model", type=str, default="gpt-4.1")
    ap.add_argument("--builder-model", type=str, default="gpt-4o-mini")
    ap.add_argument(
        "--builder-two-phase",
        action="store_true",
        help="IR extractor Phase 1 structure + Phase 2 guards/sync (often better timing/sync)",
    )
    ap.add_argument("--max-validation-rounds", type=int, default=None)
    ap.add_argument("--disable-verifyta-smoke", action="store_true")
    args = ap.parse_args()

    if args.batch_root:
        batch_root = Path(args.batch_root).resolve()
    else:
        stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        batch_kind = "norag" if args.no_rag else "kb"
        batch_root = ROOT / "artifacts" / "formal_build" / f"batch_formal_{batch_kind}_{stamp}"
    batch_root.mkdir(parents=True, exist_ok=True)
    report_path = batch_root / "report.jsonl"

    if args.rebuild_kb_index and not args.no_rag:
        maybe_rebuild_kb_index()

    all_rows = {int(r["spec_id"]): r for r in load_all_specs(args.spec_csv)}
    indices = sorted(
        sid for sid in all_rows.keys() if args.start <= sid <= args.end
    )

    pipe_kw: Dict[str, Any] = {
        "stub_analyzer": False,
        "stub_agents": False,
        "llm_model": args.analyze_model,
        "builder_model_name": args.builder_model,
        "builder_two_phase": bool(args.builder_two_phase),
        "use_rag_kb": not args.no_rag,
        "kb_rebuild": False,
        "enable_verifyta_smoke": not args.disable_verifyta_smoke,
    }
    if args.max_validation_rounds is not None:
        pipe_kw["max_validation_or_refine_rounds"] = args.max_validation_rounds

    pipe = FormalModelBuildingPipeline(**pipe_kw)

    print(f"Batch root: {batch_root}")
    kb_on = pipe_kw["use_rag_kb"]
    print(f"Spec range: {args.start}–{args.end}  KB: {'on' if kb_on else 'off'}  verifyta: {pipe_kw['enable_verifyta_smoke']}")
    sys.stdout.flush()

    for sid in indices:
        row = all_rows[sid]
        label = (row.get("system_name") or "").strip() or f"spec_{sid}"
        spec_text = (row.get("nl_spec") or "").strip()
        out = batch_root / f"S{sid:02d}_{_safe_label(label)}"
        print(f"\n========== SPEC {sid}  {label!r} ==========")
        sys.stdout.flush()
        t0 = time.monotonic()

        summary = pipe.run(spec_text, out)

        elapsed = time.monotonic() - t0

        tpl_n = sch_n = None
        sch_path = out / "formal_schema.json"
        if sch_path.is_file():
            try:
                sch = json.loads(sch_path.read_text(encoding="utf-8"))
                tpls = sch.get("templates") or []
                tpl_n = len(tpls)
                sch_n = sum(len(t.get("transitions") or []) for t in tpls if isinstance(t, dict))
            except Exception:
                pass

        rec: Dict[str, Any] = {
            "spec_id": sid,
            "system_name": label,
            "elapsed_sec": round(elapsed, 1),
            "out_dir": str(out.relative_to(ROOT)),
            "summary_ok": summary.ok,
            "validator_ok": summary.validator_ok,
            "degraded_export": summary.degraded_export,
            "repair_rounds_used": summary.repair_rounds_used,
            "ir_error_count_final": summary.ir_error_count_final,
            "schema_error_count_final": summary.schema_error_count_final,
            "smoke_ok": summary.smoke_ok,
            "model_check_skipped": summary.model_check_skipped,
            "error": summary.error,
            "template_count": tpl_n,
            "transition_count_total": sch_n,
        }

        tp_path = out / "formal_build_trace.json"
        if tp_path.is_file():
            try:
                tp = json.loads(tp_path.read_text(encoding="utf-8"))
                rec["ir_errors_final"] = tp.get("ir_errors_final") or []
            except Exception:
                rec["ir_errors_final"] = []

        rec["quality_grade"] = quality_grade(rec)

        with report_path.open("a", encoding="utf-8") as rf:
            rf.write(json.dumps(rec, ensure_ascii=False) + "\n")

        print(
            f"  grade={rec['quality_grade']:18} ok={summary.ok} validator_ok={summary.validator_ok} "
            f"degraded={summary.degraded_export} smoke={summary.smoke_ok} "
            f"ir_err={summary.ir_error_count_final} rounds={summary.repair_rounds_used} "
            f"tpl/tr={tpl_n}/{sch_n}  {elapsed:.0f}s"
        )
        if summary.error:
            print(f"  error: {summary.error}")
        ie = rec.get("ir_errors_final") or []
        if isinstance(ie, list) and ie:
            preview = ie[0][:160] + ("..." if len(ie[0]) > 160 else "")
            print(f"  ir head: {preview!r}")
        sys.stdout.flush()

    print(f"\nWrote report: {report_path}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())