#!/usr/bin/env python3
"""In-memory held-out evaluation used by run_sft.py during training.

On a single GPU you cannot serve vLLM while training holds VRAM, so the per-epoch
benchmark runs on the *training* model directly (model.generate) on a small sample
of the same held-out sets eval_endpoint.py uses. Scoring/parse logic is imported
from eval_endpoint so in-training numbers are comparable to the base baseline.

Robustness: callers should wrap this in try/except — a benchmark hiccup must never
kill a training run.
"""

from __future__ import annotations

import json
import sys
from pathlib import Path
from typing import Any

sys.path.insert(0, str(Path(__file__).resolve().parent))
from eval_endpoint import MCQ_INSTRUCTION, parse_letter, parse_verdict  # noqa: E402


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


def _generate(model, tokenizer, messages: list[dict[str, str]], max_new_tokens: int,
              enable_thinking: bool = False) -> str:
    import torch

    try:
        prompt = tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True, enable_thinking=enable_thinking
        )
    except TypeError:
        prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        out = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
        )
    return tokenizer.decode(out[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=False)


def _score_vuln(model, tokenizer, rows, max_new_tokens, enable_thinking=False) -> dict[str, Any]:
    tp = fp = tn = fn = unparsed = 0
    for r in rows:
        pred = parse_verdict(_generate(model, tokenizer, r["messages"], max_new_tokens, enable_thinking))
        gold = r["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"
    n = len(rows)
    prec = tp / (tp + fp) if (tp + fp) else 0.0
    rec = tp / (tp + fn) if (tp + fn) else 0.0
    return {
        "kind": "vuln_detection", "n": n,
        "accuracy": (tp + tn) / n if n else 0.0,
        "precision_vuln": prec, "recall_vuln": rec,
        "f1_vuln": 2 * prec * rec / (prec + rec) if (prec + rec) else 0.0,
        "unparsed": unparsed,
    }


def _score_mcq(model, tokenizer, rows, max_new_tokens, enable_thinking=False) -> dict[str, Any]:
    correct = unparsed = 0
    for r in rows:
        choices = r["choices"]
        user = r["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},
        ]
        pred = parse_letter(_generate(model, tokenizer, messages, max_new_tokens, enable_thinking), len(choices))
        if pred is None:
            unparsed += 1
            continue
        correct += int(pred == r["gold_index"])
    n = len(rows)
    return {"kind": "mcq", "n": n, "accuracy": correct / n if n else 0.0, "unparsed": unparsed}


def run_eval_sets(model, tokenizer, eval_files: list[str], sample_per_set: int = 80,
                  max_new_tokens: int = 256, enable_thinking: bool = False) -> dict[str, Any]:
    """Return {set_name: metrics} for each eval file (sampled for speed).

    enable_thinking defaults to False: this tracking eval measures direct-answer
    accuracy so generations stay short, fast and parseable. The base baseline and
    every per-epoch run use the same setting, so the numbers stay comparable. (The
    primary CyberGym metric exercises the full agentic thinking flow separately.)
    """
    was_training = model.training
    model.eval()
    # generate() needs the kv-cache, which gradient checkpointing disables
    prev_use_cache = getattr(model.config, "use_cache", None)
    try:
        model.config.use_cache = True
    except Exception:
        pass

    results: dict[str, Any] = {}
    try:
        for f in eval_files:
            path = Path(f)
            if not path.is_file():
                results[path.stem] = {"error": "missing"}
                continue
            rows = _read_jsonl(path, sample_per_set)
            if not rows:
                results[path.stem] = {"error": "empty"}
                continue
            kind = rows[0].get("kind", "vuln_detection")
            scorer = _score_mcq if kind == "mcq" else _score_vuln
            results[path.stem] = scorer(model, tokenizer, rows, max_new_tokens, enable_thinking)
    finally:
        if prev_use_cache is not None:
            try:
                model.config.use_cache = prev_use_cache
            except Exception:
                pass
        if was_training:
            model.train()
    return results