#!/usr/bin/env python3
"""Score an OpenAI-compatible endpoint on held-out eval sets.

This is the reusable, comparable evaluator for the secondary/guardrail metrics:
run it against the base model now, and against every checkpoint's vLLM endpoint
as training goes on -- same sets, same parsing, directly comparable reports.

It complements (does not replace) the agentic CyberGym harness, which is the
primary metric and runs separately via OpenHands + Docker.

Supported eval sets (auto-detected by each row's `kind`, built by build_eval_sets.py):
  vuln_detection  -- row has {messages, gold_label in {vulnerable, not_vulnerable}}
  mcq             -- row has {question, choices, gold_index/gold_letter}

Usage (real endpoint, e.g. base vLLM):
  python training/scripts/eval_endpoint.py \
    --base-url http://127.0.0.1:8000/v1 --model qwen36-base --label base \
    --eval data/eval/vuln_detection_test.jsonl --eval data/eval/knowledge_mcq.jsonl \
    --report-dir reports/eval

Offline plumbing test:
  python training/scripts/eval_endpoint.py --mock --label smoke \
    --eval data/eval/vuln_detection_test.jsonl --report-dir /tmp/eval
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import urllib.request
from pathlib import Path
from typing import Any

MCQ_INSTRUCTION = "Answer with the single letter of the correct choice (e.g. 'Answer: C')."


def read_jsonl(path: Path) -> list[dict[str, Any]]:
    rows = []
    with path.open("r", encoding="utf-8") as fh:
        for line in fh:
            line = line.strip()
            if line:
                rows.append(json.loads(line))
    return rows


def call_endpoint(base_url: str, model: str, api_key: str, messages: list[dict[str, str]],
                  temperature: float, max_tokens: int, timeout: int = 180) -> str:
    payload = {"model": model, "messages": messages, "temperature": temperature, "max_tokens": max_tokens}
    req = urllib.request.Request(
        base_url.rstrip("/") + "/chat/completions",
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json", "Authorization": f"Bearer {api_key}"},
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        body = json.loads(resp.read())
    return body["choices"][0]["message"]["content"]


def strip_think(text: str) -> str:
    """Return the answer portion after the last </think> (if any)."""
    if "</think>" in text:
        return text.rsplit("</think>", 1)[-1].strip()
    return text.strip()


def parse_verdict(text: str) -> str | None:
    low = strip_think(text).lower()
    if "not vulnerable" in low or "no vulnerability" in low or "not a vulnerability" in low or "is secure" in low:
        return "not_vulnerable"
    if "vulnerable" in low or "vulnerability" in low:
        return "vulnerable"
    return None


def parse_letter(text: str, n_choices: int) -> int | None:
    body = strip_think(text)
    letters = [chr(ord("A") + i) for i in range(n_choices)]
    # Prefer an explicit "Answer: X"
    import re

    m = re.search(r"answer\s*[:\-]?\s*\(?([A-Z])\)?", body, flags=re.IGNORECASE)
    if m and m.group(1).upper() in letters:
        return ord(m.group(1).upper()) - ord("A")
    # else first standalone capital letter token within range
    for tok in re.findall(r"\b([A-Z])\b", body):
        if tok in letters:
            return ord(tok) - ord("A")
    return None


def mock_answer(row: dict[str, Any], idx: int) -> str:
    """Deterministic test stub: correct on even index, wrong on odd."""
    correct = idx % 2 == 0
    if row.get("kind") == "mcq":
        gold = row.get("gold_index", 0)
        n = len(row.get("choices", []))
        choose = gold if correct else (gold + 1) % max(1, n)
        return f"<think>mock</think>\nAnswer: {chr(ord('A') + choose)}"
    gold = row.get("gold_label", "not_vulnerable")
    if not correct:
        gold = "vulnerable" if gold == "not_vulnerable" else "not_vulnerable"
    verdict = "Vulnerable." if gold == "vulnerable" else "Not vulnerable."
    return f"<think>mock</think>\n{verdict}"


def eval_vuln(rows, args) -> dict[str, Any]:
    tp = fp = tn = fn = unparsed = 0
    for i, row in enumerate(rows):
        out = mock_answer(row, i) if args.mock else call_endpoint(
            args.base_url, args.model, args.api_key, row["messages"], args.temperature, args.max_tokens
        )
        pred = parse_verdict(out)
        gold = row["gold_label"]
        if pred is None:
            unparsed += 1
            continue
        if gold == "vulnerable":
            tp += pred == "vulnerable"
            fn += pred != "vulnerable"
        else:
            tn += pred == "not_vulnerable"
            fp += pred != "not_vulnerable"
    total = len(rows)
    correct = tp + tn
    acc = correct / total if total else 0.0
    prec = tp / (tp + fp) if (tp + fp) else 0.0
    rec = tp / (tp + fn) if (tp + fn) else 0.0
    f1 = 2 * prec * rec / (prec + rec) if (prec + rec) else 0.0
    return {
        "kind": "vuln_detection", "n": total, "accuracy": acc,
        "precision_vuln": prec, "recall_vuln": rec, "f1_vuln": f1,
        "tp": tp, "fp": fp, "tn": tn, "fn": fn, "unparsed": unparsed,
    }


def eval_mcq(rows, args) -> dict[str, Any]:
    correct = unparsed = 0
    for i, row in enumerate(rows):
        choices = row["choices"]
        user = row["question"] + "\n\n" + "\n".join(
            f"{chr(ord('A') + j)}. {c}" for j, c in enumerate(choices)
        ) + "\n\n" + MCQ_INSTRUCTION
        messages = [
            {"role": "system", "content": "You are a cybersecurity expert. Authorized security research context."},
            {"role": "user", "content": user},
        ]
        out = mock_answer(row, i) if args.mock else call_endpoint(
            args.base_url, args.model, args.api_key, messages, args.temperature, args.max_tokens
        )
        pred = parse_letter(out, len(choices))
        if pred is None:
            unparsed += 1
            continue
        correct += int(pred == row["gold_index"])
    total = len(rows)
    return {"kind": "mcq", "n": total, "accuracy": correct / total if total else 0.0,
            "correct": correct, "unparsed": unparsed}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--eval", action="append", required=True, help="Eval JSONL (repeatable).")
    parser.add_argument("--label", required=True, help="Run label, e.g. base / epoch1 / merged.")
    parser.add_argument("--model", default=os.getenv("EVAL_MODEL", "qwen36-base"))
    parser.add_argument("--base-url", default=os.getenv("EVAL_BASE_URL", "http://127.0.0.1:8000/v1"))
    parser.add_argument("--api-key", default=os.getenv("EVAL_API_KEY", "EMPTY"))
    parser.add_argument("--report-dir", default="reports/eval")
    parser.add_argument("--temperature", type=float, default=0.0)
    parser.add_argument("--max-tokens", type=int, default=2048)
    parser.add_argument("--limit", type=int, default=None, help="Only eval first N rows per set.")
    parser.add_argument("--mock", action="store_true", help="Offline deterministic stub (testing only).")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    results = []
    for path_str in args.eval:
        path = Path(path_str)
        if not path.is_file():
            results.append({"file": path_str, "error": "missing"})
            continue
        rows = read_jsonl(path)
        if args.limit:
            rows = rows[: args.limit]
        if not rows:
            results.append({"file": path_str, "error": "empty"})
            continue
        kind = rows[0].get("kind", "vuln_detection")
        res = eval_mcq(rows, args) if kind == "mcq" else eval_vuln(rows, args)
        res["file"] = path_str
        results.append(res)

    report_dir = Path(args.report_dir)
    report_dir.mkdir(parents=True, exist_ok=True)
    payload = {"label": args.label, "model": args.model, "mock": bool(args.mock), "results": results}
    (report_dir / f"{args.label}_eval.json").write_text(json.dumps(payload, indent=2), encoding="utf-8")

    lines = [f"# Endpoint Eval — {args.label}", "", f"- Model: `{args.model}`",
             f"- Endpoint: `{args.base_url}`", f"- Mock: {bool(args.mock)}", ""]
    for r in results:
        if "error" in r:
            lines.append(f"## {r['file']} — ERROR: {r['error']}")
            continue
        lines.append(f"## {r['kind']} ({r['file']})")
        lines.append(f"- n: {r['n']}")
        lines.append(f"- accuracy: **{r['accuracy']:.2%}**")
        if r["kind"] == "vuln_detection":
            lines.append(f"- precision(vuln): {r['precision_vuln']:.2%}  recall(vuln): {r['recall_vuln']:.2%}  f1: {r['f1_vuln']:.2%}")
            lines.append(f"- tp/fp/tn/fn: {r['tp']}/{r['fp']}/{r['tn']}/{r['fn']}  unparsed: {r['unparsed']}")
        else:
            lines.append(f"- correct: {r['correct']}  unparsed: {r['unparsed']}")
        lines.append("")
    (report_dir / f"{args.label}_eval.md").write_text("\n".join(lines) + "\n", encoding="utf-8")

    print(json.dumps(payload, indent=2))
    return 0


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