#!/usr/bin/env python3
"""
Test the simplified NL → UPPAAL query pipeline on one spec.

  python scripts/run_nl_query_simple.py --spec-id 11
  python scripts/run_nl_query_simple.py --spec-id 6 --model gpt-4o-mini --verbose
  python scripts/run_nl_query_simple.py --spec-id 13 --gold-model datasets/gold_models/S13/model.xml

Reads NL queries from datasets/query_translation_eval.csv and translates each
one against the gold model, printing verdict + status.
"""

from __future__ import annotations

import argparse
import csv
import sys
from pathlib import Path

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

from frame.pipeline.model_checking_pipeline import ModelCheckingPipeline
from frame.pipeline.nl_query_simple import translate_nl_to_uppaal_query
from frame.pipeline.nl_uppaal_query import parse_verifyta_text_verdict
from frame.rag_component.llm import LLM

EVAL_CSV = ROOT / "datasets" / "query_translation_eval.csv"
GOLD_DIR = ROOT / "datasets" / "gold_models"

SYSTEM_PROMPT = """You are an expert in formal verification with UPPAAL timed automata.
Your only task is to produce valid UPPAAL CTL query formulas.
Follow the rules in the user message exactly.
Output exactly one line: the UPPAAL formula, nothing else."""


def load_queries(spec_id: int) -> list[dict]:
    rows = []
    with EVAL_CSV.open(encoding="utf-8", newline="") as f:
        for row in csv.DictReader(f):
            if int(row["spec_id"]) == spec_id:
                rows.append(row)
    return rows


def verdict_symbol(v: str) -> str:
    return {"satisfied": "T", "not_satisfied": "F"}.get(v, "?")


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--spec-id", type=int, default=11, help="Spec ID (1-20)")
    ap.add_argument(
        "--gold-model", type=Path, default=None,
        help="Override gold model path (default: datasets/gold_models/S{id:02d}/model.xml)",
    )
    ap.add_argument("--model", type=str, default=None, help="LLM model id")
    ap.add_argument("--verbose", action="store_true")
    args = ap.parse_args()

    xml_path = args.gold_model or GOLD_DIR / f"S{args.spec_id:02d}" / "model.xml"
    if not xml_path.is_file():
        print(f"Model not found: {xml_path}", file=sys.stderr)
        return 1

    rows = load_queries(args.spec_id)
    if not rows:
        print(f"No queries for spec_id={args.spec_id} in {EVAL_CSV}", file=sys.stderr)
        return 1

    llm = LLM(
        system_prompt=SYSTEM_PROMPT,
        model_name=args.model or "gpt-4o-mini",
        max_tokens=300,
    )

    checker = ModelCheckingPipeline(str(xml_path))

    print(f"\nSpec S{args.spec_id:02d} — {xml_path.name}  ({len(rows)} queries)\n")
    print(f"{'#':>2}  {'Status':<12} {'Got':>4} {'Exp':>4}  {'Formula':<60}  NL")
    print("-" * 120)

    correct = 0
    for i, row in enumerate(rows, 1):
        nl = row["nl_query"].strip()
        gold_q = row["ground_query"].strip()
        expected = row.get("answer", "").strip().upper()

        formula, status = translate_nl_to_uppaal_query(
            nl, xml_path, llm=llm, verbose=args.verbose
        )

        if formula:
            res, _t, errs = checker.verify(formula)
            v = parse_verifyta_text_verdict(res, errors=errs)
            got = verdict_symbol(v)
        else:
            got = "?"

        match = "✓" if got == expected else "✗"
        if got == expected:
            correct += 1

        print(
            f"{i:>2}  {status:<12} {got:>4} {expected:>4}  "
            f"{(formula or '(none)')[:60]:<60}  {nl[:60]}"
        )
        if args.verbose:
            print(f"    gold: {gold_q}")

    print("-" * 120)
    print(f"\nScore: {correct}/{len(rows)}  ({100*correct//len(rows)}%)\n")
    return 0


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