"""Score a Vela 2.0 model (or the regex scanner) on the eval set.

  python scripts/run_baseline.py --model vllm-sr/Vela-2.0-0.8B
  python scripts/run_baseline.py --model regex          # secrets task only, no GPU
  python scripts/run_baseline.py --model vllm-sr/Vela-2.0-0.8B --heads runs/vela-0.8b-guard

Writes results/<slug>/predictions.jsonl and results/<slug>/metrics.json.
"""
import argparse
import json
import logging
import time
from pathlib import Path

from vela_guard import evaluate, regex_baseline

TASKS = ("secrets", "injection", "toolcall")

log = logging.getLogger("run_baseline")


def load_rows(data_dir: Path, tasks: tuple[str, ...]) -> dict[str, list[dict]]:
    rows = {}
    for task in tasks:
        path = data_dir / f"{task}.jsonl"
        if not path.exists():
            raise FileNotFoundError(f"{path} missing; run scripts/build_eval.py first")
        rows[task] = [json.loads(line) for line in path.read_text().splitlines() if line.strip()]
    return rows


def load_vela(name: str, device: str, heads_dir: Path | None):
    import torch
    from transformers import AutoModel

    if device == "cuda" and not torch.cuda.is_available():
        raise RuntimeError("CUDA not available; this script is meant to run on a RunPod GPU")
    if heads_dir is not None:
        from trainer.guard_model import GuardModel

        return GuardModel.load(name, heads_dir, device)
    model = AutoModel.from_pretrained(name, trust_remote_code=True)
    return model.to(device).eval()


def predict_vela(model, row: dict) -> tuple[dict, float]:
    state, questions = evaluate.build_request(row)
    start = time.perf_counter()
    response = model.system_one(state, questions)
    return evaluate.parse_response(row["task"], response), time.perf_counter() - start


def predict_regex(row: dict) -> tuple[dict, float]:
    start = time.perf_counter()
    spans = [{k: s[k] for k in ("label", "start", "end")} for s in regex_baseline.find_secrets(row["state"]["request"])]
    return {"spans": spans, "score": 1.0 if spans else 0.0}, time.perf_counter() - start


def run(args: argparse.Namespace) -> None:
    is_regex = args.model == "regex"
    tasks = ("secrets",) if is_regex else tuple(args.tasks)
    rows_by_task = load_rows(args.data, tasks)
    model = None if is_regex else load_vela(args.model, args.device, args.heads)

    out_dir = args.out / (args.model.replace("/", "__") + ("__guard" if args.heads else ""))
    out_dir.mkdir(parents=True, exist_ok=True)
    report, failures = {"model": args.model, "tasks": {}}, 0
    with (out_dir / "predictions.jsonl").open("w") as fh:
        for task, rows in rows_by_task.items():
            kept_rows, preds, latencies = [], [], []
            for row in rows:
                try:
                    pred, seconds = predict_regex(row) if is_regex else predict_vela(model, row)
                except Exception:  # keep going; the failure is logged and counted in the report
                    log.exception("prediction failed for %s", row["id"])
                    failures += 1
                    continue
                kept_rows.append(row)
                preds.append(pred)
                latencies.append(seconds)
                fh.write(json.dumps({"id": row["id"], "pred": pred, "seconds": seconds}) + "\n")
            task_report = evaluate.score_task(task, kept_rows, preds)
            task_report["median_ms"] = 1000 * sorted(latencies)[len(latencies) // 2] if latencies else None
            report["tasks"][task] = task_report
            log.info("%s done: %d/%d rows", task, len(kept_rows), len(rows))
    report["failures"] = failures
    (out_dir / "metrics.json").write_text(json.dumps(report, indent=2))
    log.info("%s", headline(report))


def headline(report: dict) -> str:
    t = report["tasks"]
    parts = [report["model"]]
    if "secrets" in t:
        parts.append(f"secrets exactF1={t['secrets']['all']['exact']['f1']:.3f} overlapF1={t['secrets']['all']['overlap']['f1']:.3f}")
    if "injection" in t:
        parts.append(f"injection acc={t['injection']['all']['accuracy']:.3f} auroc={t['injection']['all']['auroc']:.3f}")
    if "toolcall" in t:
        parts.append(f"toolcall acc={t['toolcall']['all']['accuracy']:.3f} block->allow={t['toolcall']['all']['block_passed_as_allow']:.3f}")
    return " | ".join(parts)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--model", required=True, help="HF repo id, or 'regex'")
    parser.add_argument("--data", type=Path, default=Path("data/eval"))
    parser.add_argument("--out", type=Path, default=Path("results"))
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--tasks", nargs="+", choices=TASKS, default=list(TASKS))
    parser.add_argument("--heads", type=Path, default=None, help="trained guard heads dir (trainer.train_heads --out)")
    run(parser.parse_args())


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s")
    main()