#!/usr/bin/env python3
"""Small pluggable eval harness for before/after AI harness validation."""

import argparse
import json
import os
import re
import shlex
import shutil
import subprocess
import sys
import tempfile
import time
from pathlib import Path


def runner_binary(template):
    parts = shlex.split(template)
    return parts[0] if parts else ""


def manual_protocol(binary, tasks_path):
    return f"""Runner binary `{binary}` was not found.

Manual protocol:
1. Open the target repo in your agent tool.
2. For each task in `{tasks_path}`, run the prompt exactly once.
3. Record JSON with shape:
   {{"label":"manual","tasks":[{{"id":"task-id","passed":true,"duration_s":0.0,"exit_code":0,"stdout":"...","answer":"..."}}]}}
4. Compare two result files with: python3 scripts/eval_run.py --compare before.json after.json -o report.md
"""


def maybe_usage(stdout):
    try:
        data = json.loads(stdout)
    except Exception:
        return {}
    usage = {}
    for key in ["usage", "cost", "total_cost_usd", "tokens", "input_tokens", "output_tokens"]:
        if key in data:
            usage[key] = data[key]
    return usage


def timeout_output(value):
    if value is None:
        return ""
    if isinstance(value, bytes):
        return value.decode("utf-8", errors="replace")
    return value


def extract_answer(stdout):
    """Return the normalized text used for grading.

    Claude CLI `--output-format json` emits a JSON envelope whose `result` field is
    the agent answer. Older/manual results may be raw text, so fall back to stdout.
    """
    raw = timeout_output(stdout)
    answer = raw
    try:
        data = json.loads(raw)
    except Exception:
        data = None
    if isinstance(data, dict) and isinstance(data.get("result"), str):
        answer = data["result"]
    return answer.strip().strip("`").strip()


def regex_passes(pattern, answer):
    return re.search(pattern or "", answer or "") is not None


def _as_pattern_list(value):
    if value is None:
        return []
    if isinstance(value, str):
        return [value]
    if isinstance(value, list):
        return [str(v) for v in value]
    return []


_STOPWORDS = {
    "the", "and", "for", "with", "that", "this", "must", "should", "will", "have",
    "from", "into", "your", "when", "then", "than", "each", "also", "such", "which",
    "list", "explain", "describe", "mention", "answer", "correct", "correctly",
    "include", "including", "provide", "using", "used", "note", "about", "these",
    "there", "their", "does", "what", "where", "some", "only", "make", "sure",
}


def rubric_keywords(rubric):
    """Extract salient terms an answer should contain from a free-text rubric.

    Prefers explicit signals — backtick-quoted `code`, "double" / 'single' quoted
    phrases — and otherwise falls back to distinctive words (>= 4 chars, minus a
    small stopword list). Deterministic: no LLM required.
    """
    rubric = rubric or ""
    keywords = []
    seen = set()

    def _add(term):
        term = term.strip().lower()
        if term and term not in seen and term not in _STOPWORDS:
            seen.add(term)
            keywords.append(term)

    for m in re.findall(r"`([^`]+)`", rubric):
        _add(m)
    for m in re.findall(r"\"([^\"]+)\"|'([^']+)'", rubric):
        _add(m[0] or m[1])
    if not keywords:
        for word in re.findall(r"[A-Za-z0-9_][A-Za-z0-9_.-]{3,}", rubric):
            _add(word)
    return keywords


def builtin_judge(answer, check):
    """Deterministic, dependency-free default judge.

    Grades ``answer`` against the task's ``check`` without any external command:
      - ``expect``: regex pattern(s) that MUST all match (case-insensitive).
      - ``reject``: regex pattern(s) that must NOT match.
      - otherwise: keyword coverage derived from the ``rubric`` / ``criteria`` text,
        passing when coverage >= ``min_score`` (default 0.5).
    Returns the same verdict shape as the external judge, tagged ``judge="builtin"``.
    """
    answer = answer or ""
    expect = _as_pattern_list(check.get("expect"))
    reject = _as_pattern_list(check.get("reject"))
    rubric = check.get("rubric") or check.get("criteria") or ""
    try:
        threshold = float(check.get("min_score", 0.5))
    except (TypeError, ValueError):
        threshold = 0.5

    criteria = []  # (label, satisfied)
    for pat in expect:
        criteria.append((f"expect:{pat}", re.search(pat, answer, re.I) is not None))
    for pat in reject:
        criteria.append((f"reject:{pat}", re.search(pat, answer, re.I) is None))

    keyword_mode = not expect and not reject
    if keyword_mode:
        for kw in rubric_keywords(rubric):
            criteria.append((f"keyword:{kw}", kw in answer.lower()))

    if not criteria:
        return {
            "passed": False, "score": 0.0, "judge": "builtin",
            "reason": "builtin judge needs `expect`/`reject`/`rubric` to grade; none provided",
        }

    satisfied = sum(1 for _, ok in criteria if ok)
    total = len(criteria)
    score = round(satisfied / total, 4)
    if keyword_mode:
        passed = score >= threshold
    else:
        passed = all(ok for _, ok in criteria)
    unmet = [label for label, ok in criteria if not ok]
    reason = f"builtin judge: {satisfied}/{total} criteria satisfied"
    if unmet:
        reason += "; unmet: " + ", ".join(unmet[:5])
    return {"passed": bool(passed), "score": score, "reason": reason, "judge": "builtin"}


HEALTH_GRADES = [(90, "A"), (80, "B"), (70, "C"), (60, "D"), (0, "F")]


def health_grade(score):
    for threshold, letter in HEALTH_GRADES:
        if score >= threshold:
            return letter
    return "F"


def _collect_records(data):
    """Return the flat list of task records from a run-tasks or matrix result dict."""
    if not isinstance(data, dict):
        return []
    if isinstance(data.get("agents"), dict):
        records = []
        for agent in data["agents"].values():
            if isinstance(agent, dict):
                records.extend(agent.get("tasks", []) or [])
        return records
    return data.get("tasks", []) or []


def compute_health(data):
    """Compute an automated efficacy health score (0-100) with a letter grade.

    Works on both a single ``run`` result (``{"tasks": [...]}``) and a ``matrix``
    result (``{"agents": {...}}``); the score is the pass rate across every task
    record. Timeouts count as failures.
    """
    records = _collect_records(data)
    total = len(records)
    passed = sum(1 for r in records if r.get("passed"))
    timed_out = sum(1 for r in records if r.get("timed_out"))
    pass_rate = (passed / total) if total else 0.0
    score = round(100 * pass_rate)
    return {
        "score": score,
        "grade": health_grade(score),
        "passed": passed,
        "total": total,
        "timed_out": timed_out,
        "pass_rate": round(pass_rate, 4),
    }


def parse_judge_output(stdout):
    """Parse an LLM-as-judge verdict.

    The judge command MUST print a single JSON object with shape
    ``{"passed": bool, "score": number, "reason": "..."}``. If ``passed`` is
    omitted but a numeric ``score`` is present, a score >= 0.5 counts as a pass.
    A best-effort attempt is made to locate a JSON object embedded in noisier
    output so simple echo-based judges keep working.
    """
    raw = timeout_output(stdout).strip()
    data = None
    try:
        data = json.loads(raw)
    except Exception:
        match = re.search(r"\{.*\}", raw, re.S)
        if match:
            try:
                data = json.loads(match.group(0))
            except Exception:
                data = None
    if not isinstance(data, dict):
        return {"passed": False, "score": None, "reason": "judge output was not valid JSON", "raw": raw}
    passed = data.get("passed")
    score = data.get("score")
    if passed is None and isinstance(score, (int, float)) and not isinstance(score, bool):
        passed = score >= 0.5
    return {"passed": bool(passed), "score": score, "reason": data.get("reason", ""), "raw": raw}


def run_judge(judge_cmd, answer, rubric, workdir, timeout):
    """Invoke a configurable judge command and return its parsed verdict.

    Contract (documented in SKILL.md): the judge command is run through the
    shell with the following inputs available so authors can pick whatever fits:
      - env ``JUDGE_ANSWER``: the produced answer text
      - env ``JUDGE_RUBRIC``: the rubric / criteria string (may be empty)
      - env ``JUDGE_INPUT``: path to a temp JSON file ``{"answer":..., "rubric":...}``
      - template placeholders ``{answer}`` / ``{rubric}`` / ``{input}`` (shell-quoted)
    The command MUST print ``{"passed": bool, "score": number, "reason": "..."}``.
    Raises ``subprocess.TimeoutExpired`` if the judge exceeds ``timeout`` seconds.
    """
    answer = answer or ""
    rubric = rubric or ""
    env = dict(os.environ)
    env["JUDGE_ANSWER"] = answer
    env["JUDGE_RUBRIC"] = rubric
    handle = tempfile.NamedTemporaryFile("w", suffix=".json", delete=False, encoding="utf-8")
    try:
        handle.write(json.dumps({"answer": answer, "rubric": rubric}, ensure_ascii=False))
        handle.close()
        env["JUDGE_INPUT"] = handle.name
        command = (
            judge_cmd
            .replace("{answer}", shlex.quote(answer))
            .replace("{rubric}", shlex.quote(rubric))
            .replace("{input}", shlex.quote(handle.name))
        )
        proc = subprocess.run(
            command, cwd=str(workdir) if workdir else None, text=True,
            capture_output=True, shell=True, env=env, timeout=timeout,
        )
        verdict = parse_judge_output(proc.stdout)
        verdict["exit_code"] = proc.returncode
        return verdict
    finally:
        try:
            os.unlink(handle.name)
        except OSError:
            pass


def grade_answer(task, answer, workdir, judge_cmd=None, default_judge=True):
    """Grade an answer for a task, returning ``(passed, judge_info)``.

    Supports ``regex``, ``command`` and ``judge`` check types. ``judge_info`` is
    ``None`` for non-judge checks. Timeouts count as a failure (no exception is
    propagated) so a matrix run keeps going.

    For ``judge`` checks an external ``--judge-cmd`` takes priority; when none is
    supplied the deterministic built-in judge (:func:`builtin_judge`) grades the
    answer so judge checks work out of the box. Pass ``default_judge=False`` to
    restore the legacy "requires --judge-cmd" behavior.
    """
    check = task.get("check", {})
    ctype = check.get("type")
    timeout = task.get("timeout_s", 60)
    if ctype == "regex":
        return regex_passes(check.get("value", ""), answer), None
    if ctype == "command":
        try:
            cproc = subprocess.run(check.get("value", ""), cwd=str(workdir), text=True, capture_output=True, shell=True, timeout=timeout)
            return cproc.returncode == 0, None
        except subprocess.TimeoutExpired:
            return False, None
    if ctype == "judge":
        rubric = check.get("rubric") or check.get("criteria") or ""
        if judge_cmd:
            try:
                verdict = run_judge(judge_cmd, answer, rubric, workdir, timeout)
            except subprocess.TimeoutExpired:
                return False, {"passed": False, "score": None, "reason": "judge timed out"}
            return verdict["passed"], verdict
        if default_judge:
            verdict = builtin_judge(answer, check)
            return verdict["passed"], verdict
        return False, {"passed": False, "score": None, "reason": "no --judge-cmd provided"}
    return False, None


def run_tasks(args):
    tasks_path = Path(args.tasks)
    tasks = json.loads(tasks_path.read_text(encoding="utf-8"))
    binary = runner_binary(args.runner)
    if binary and shutil.which(binary) is None:
        print(manual_protocol(binary, args.tasks), file=sys.stderr)
        return 127
    workdir = Path(args.workdir).resolve()
    results = {"label": args.label, "workdir": str(workdir), "tasks": []}
    for task in tasks:
        prompt = task["prompt"]
        command = args.runner.replace("{prompt}", shlex.quote(prompt))
        start = time.time()
        try:
            proc = subprocess.run(command, cwd=str(workdir), text=True, capture_output=True, shell=True, timeout=task.get("timeout_s", 60))
        except subprocess.TimeoutExpired as exc:
            duration = round(time.time() - start, 3)
            stdout = timeout_output(exc.stdout)
            results["tasks"].append({
                "id": task["id"], "passed": False, "timed_out": True, "duration_s": duration,
                "exit_code": None, "stdout": stdout, "answer": extract_answer(stdout), "stderr": timeout_output(exc.stderr),
                "usage": {},
            })
            continue
        duration = round(time.time() - start, 3)
        answer = extract_answer(proc.stdout)
        passed, judge_info = grade_answer(task, answer, workdir, args.judge_cmd, not args.no_default_judge)
        record = {
            "id": task["id"], "passed": passed, "timed_out": False, "duration_s": duration,
            "exit_code": proc.returncode, "stdout": proc.stdout, "answer": answer, "stderr": proc.stderr,
            "usage": maybe_usage(proc.stdout),
        }
        if judge_info is not None:
            record["judge"] = judge_info
        results["tasks"].append(record)
    results["health"] = compute_health(results)
    output = Path(args.output) if args.output else Path(f"results-{args.label}.json")
    output.write_text(json.dumps(results, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"wrote {output}")
    health = results["health"]
    print(f"health score: {health['score']}/100 (grade {health['grade']}), "
          f"{health['passed']}/{health['total']} tasks passed")
    if args.fail_under is not None and health["score"] < args.fail_under:
        print(f"health score {health['score']} is below --fail-under {args.fail_under}", file=sys.stderr)
        return 5
    return 0


def regrade(args):
    tasks = json.loads(Path(args.tasks).read_text(encoding="utf-8"))
    task_map = {task["id"]: task for task in tasks}
    results = json.loads(Path(args.regrade).read_text(encoding="utf-8"))
    for record in results.get("tasks", []):
        task = task_map.get(record.get("id"))
        answer = extract_answer(record.get("stdout", ""))
        record["answer"] = answer
        if not task:
            record["regraded"] = False
            continue
        check = task.get("check", {})
        if check.get("type") == "regex":
            record["passed"] = regex_passes(check.get("value", ""), answer)
            record["regraded"] = True
        elif check.get("type") == "command":
            record["regraded"] = False
        else:
            record["regraded"] = False
    output = Path(args.output) if args.output else Path(args.regrade)
    output.write_text(json.dumps(results, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"wrote {output}")
    return 0


def compare(args):
    before = json.loads(Path(args.compare[0]).read_text(encoding="utf-8"))
    after = json.loads(Path(args.compare[1]).read_text(encoding="utf-8"))
    bmap = {t["id"]: t for t in before.get("tasks", [])}
    amap = {t["id"]: t for t in after.get("tasks", [])}
    ids = sorted(set(bmap) | set(amap))
    lines = ["# Phase 3 — Efficacy Comparison Report", "", f"Before label: `{before.get('label')}`", f"After label: `{after.get('label')}`", "", "| Task | Before | After | Duration Δ(s) | Usage |", "|---|---:|---:|---:|---|"]
    before_pass = after_pass = 0
    for tid in ids:
        b, a = bmap.get(tid, {}), amap.get(tid, {})
        before_pass += 1 if b.get("passed") else 0
        after_pass += 1 if a.get("passed") else 0
        delta = round((a.get("duration_s") or 0) - (b.get("duration_s") or 0), 3)
        usage = json.dumps(a.get("usage") or b.get("usage") or {}, ensure_ascii=False)
        lines.append(f"| `{tid}` | {b.get('passed')} | {a.get('passed')} | {delta} | `{usage}` |")
    lines.extend(["", f"Summary: pass rate {before_pass}/{len(ids)} → {after_pass}/{len(ids)}; delta {after_pass - before_pass} tasks."])
    content = "\n".join(lines) + "\n"
    if args.output:
        Path(args.output).write_text(content, encoding="utf-8")
    else:
        print(content, end="")
    return 0


def parse_named_runners(pairs):
    """Parse repeatable ``NAME=CMD`` runner specs into an ordered dict."""
    runners = {}
    for pair in pairs or []:
        if "=" not in pair:
            raise SystemExit(f"--runner-cmd must be NAME=CMD, got: {pair}")
        name, cmd = pair.split("=", 1)
        name = name.strip()
        if not name or not cmd.strip():
            raise SystemExit(f"--runner-cmd must be NAME=CMD, got: {pair}")
        runners[name] = cmd
    return runners


def load_matrix(args):
    """Build an ordered {agent_name: runner_template} mapping from --matrix + --runner-cmd."""
    runners = {}
    if args.matrix:
        data = json.loads(Path(args.matrix).read_text(encoding="utf-8"))
        if isinstance(data, dict) and isinstance(data.get("agents"), dict):
            data = data["agents"]
        if not isinstance(data, dict):
            raise SystemExit("--matrix file must map agent name -> runner command template")
        for name, cmd in data.items():
            runners[str(name)] = cmd
    runners.update(parse_named_runners(args.runner_cmd))
    return runners


def run_task_with_runner(runner, task, workdir, judge_cmd, default_judge=True):
    """Run a single task under a single runner and return a graded record."""
    prompt = task["prompt"]
    command = runner.replace("{prompt}", shlex.quote(prompt))
    start = time.time()
    try:
        proc = subprocess.run(command, cwd=str(workdir), text=True, capture_output=True, shell=True, timeout=task.get("timeout_s", 60))
    except subprocess.TimeoutExpired as exc:
        duration = round(time.time() - start, 3)
        stdout = timeout_output(exc.stdout)
        return {
            "id": task["id"], "passed": False, "timed_out": True, "duration_s": duration,
            "exit_code": None, "answer": extract_answer(stdout), "usage": {},
        }
    duration = round(time.time() - start, 3)
    answer = extract_answer(proc.stdout)
    passed, judge_info = grade_answer(task, answer, workdir, judge_cmd, default_judge)
    record = {
        "id": task["id"], "passed": passed, "timed_out": False, "duration_s": duration,
        "exit_code": proc.returncode, "answer": answer, "usage": maybe_usage(proc.stdout),
    }
    if judge_info is not None:
        record["judge"] = judge_info
    return record


def render_matrix(tasks, agent_records, runners):
    """Render a markdown matrix report: rows=tasks, cols=agents, plus pass-rate summary."""
    names = list(agent_records.keys())
    header = "| Task | " + " | ".join(names) + " |"
    divider = "|---|" + "".join(["---|" for _ in names])
    lines = ["# Eval Matrix Report", "", header, divider]
    maps = {name: {r["id"]: r for r in recs} for name, recs in agent_records.items()}
    for task in tasks:
        tid = task["id"]
        cells = []
        for name in names:
            record = maps[name].get(tid, {})
            if record.get("timed_out"):
                mark = "⏱ timeout"
            elif record.get("passed"):
                mark = "✅ pass"
            else:
                mark = "❌ fail"
            cells.append(f"{mark} ({record.get('duration_s', 0)}s)")
        lines.append(f"| `{tid}` | " + " | ".join(cells) + " |")
    lines.extend(["", "## Per-agent pass rate", "", "| Agent | Runner | Pass | Total | Pass rate |", "|---|---|---:|---:|---:|"])
    for name in names:
        recs = agent_records[name]
        total = len(recs)
        passed = sum(1 for r in recs if r.get("passed"))
        rate = round(100.0 * passed / total, 1) if total else 0.0
        runner = runners.get(name, "")
        lines.append(f"| `{name}` | `{runner}` | {passed} | {total} | {rate}% |")
    return "\n".join(lines) + "\n"


def run_matrix(args, runners):
    tasks = json.loads(Path(args.tasks).read_text(encoding="utf-8"))
    workdir = Path(args.workdir).resolve()
    agent_records = {}
    for name, runner in runners.items():
        binary = runner_binary(runner)
        if binary and shutil.which(binary) is None:
            print(manual_protocol(binary, args.tasks), file=sys.stderr)
            return 127
        agent_records[name] = [run_task_with_runner(runner, task, workdir, args.judge_cmd, not args.no_default_judge) for task in tasks]
    summary = {}
    for name, recs in agent_records.items():
        total = len(recs)
        passed = sum(1 for r in recs if r.get("passed"))
        summary[name] = {"passed": passed, "total": total, "pass_rate": round(passed / total, 4) if total else 0.0}
    matrix = {
        "workdir": str(workdir),
        "agents": {name: {"runner": runners[name], "tasks": recs, "summary": summary[name]} for name, recs in agent_records.items()},
        "summary": summary,
    }
    matrix["health"] = compute_health(matrix)
    json_out = Path(args.matrix_json) if args.matrix_json else Path("matrix-results.json")
    json_out.write_text(json.dumps(matrix, ensure_ascii=False, indent=2), encoding="utf-8")
    report_out = Path(args.matrix_report) if args.matrix_report else (Path(args.output) if args.output else Path("matrix-report.md"))
    report_out.write_text(render_matrix(tasks, agent_records, runners), encoding="utf-8")
    print(f"wrote {json_out}")
    print(f"wrote {report_out}")
    health = matrix["health"]
    print(f"health score: {health['score']}/100 (grade {health['grade']}), "
          f"{health['passed']}/{health['total']} tasks passed")
    if args.fail_under is not None and health["score"] < args.fail_under:
        print(f"health score {health['score']} is below --fail-under {args.fail_under}", file=sys.stderr)
        return 5
    return 0


def score_report(args):
    """One-click health score from an existing results / matrix JSON file."""
    data = json.loads(Path(args.score).read_text(encoding="utf-8"))
    health = data.get("health") if isinstance(data, dict) else None
    if not isinstance(health, dict) or "score" not in health:
        health = compute_health(data)
    if args.as_json:
        print(json.dumps(health, ensure_ascii=False, indent=2))
    else:
        print(f"Health score: {health['score']}/100 (grade {health['grade']})")
        print(f"{health['passed']}/{health['total']} tasks passed; {health.get('timed_out', 0)} timed out")
    if args.fail_under is not None and health["score"] < args.fail_under:
        print(f"health score {health['score']} is below --fail-under {args.fail_under}", file=sys.stderr)
        return 5
    return 0


def main(argv=None):
    parser = argparse.ArgumentParser(description="Run or compare AI harness eval tasks.")
    parser.add_argument("--tasks")
    parser.add_argument("--label")
    parser.add_argument("--workdir")
    parser.add_argument("--runner", default="claude -p {prompt} --output-format json")
    parser.add_argument("-o", "--output")
    parser.add_argument("--compare", nargs=2)
    parser.add_argument("--regrade")
    parser.add_argument("--judge-cmd", dest="judge_cmd", help="Command template for LLM-as-judge checks (see SKILL.md for the contract).")
    parser.add_argument("--matrix", help="JSON file mapping agent name -> runner command template.")
    parser.add_argument("--runner-cmd", dest="runner_cmd", action="append", default=[], help="Repeatable NAME=CMD runner for the eval matrix.")
    parser.add_argument("--matrix-report", dest="matrix_report", help="Output path for the markdown matrix report.")
    parser.add_argument("--matrix-json", dest="matrix_json", help="Output path for the matrix JSON results.")
    parser.add_argument("--no-default-judge", dest="no_default_judge", action="store_true",
                        help="Disable the built-in judge; judge checks then require --judge-cmd (legacy behavior).")
    parser.add_argument("--score", help="Print an automated health score for an existing results / matrix JSON file.")
    parser.add_argument("--fail-under", dest="fail_under", type=float, default=None,
                        help="Exit 5 when the health score is below this value (0-100).")
    parser.add_argument("--json", dest="as_json", action="store_true", help="Emit machine-readable JSON (used by --score).")
    args = parser.parse_args(argv)
    if args.score:
        return score_report(args)
    if args.compare:
        return compare(args)
    if args.regrade:
        if not args.tasks:
            parser.error("--tasks is required with --regrade")
        return regrade(args)
    if args.matrix or args.runner_cmd:
        if not (args.tasks and args.workdir):
            parser.error("--tasks and --workdir are required for a matrix run")
        runners = load_matrix(args)
        if not runners:
            parser.error("no runners defined; provide --matrix FILE and/or --runner-cmd NAME=CMD")
        return run_matrix(args, runners)
    required = [args.tasks, args.label, args.workdir]
    if not all(required):
        parser.error("--tasks, --label and --workdir are required unless --compare is used")
    return run_tasks(args)


if __name__ == "__main__":
    sys.exit(main())