"""Independent eval runner for DeepSeek V4 Flash q2 served by ds4-server.

Usage:
  python3 harness.py aime       # 30 AIME 2025 problems, thinking mode
  python3 harness.py humaneval  # 164 HumanEval tasks, non-thinking, greedy

Results append to results/<suite>.jsonl (resumable: done ids are skipped).
"""

import csv
import json
import random
import sys
import time
import urllib.request
from pathlib import Path

from graders import extract_code, extract_mcq_letter, grade_aime, run_humaneval_check

BASE = Path(__file__).parent
SERVER = "http://127.0.0.1:8000/v1/chat/completions"

AIME_INSTRUCTION = (
    "Please reason step by step, and put your final answer within \\boxed{}. "
    "The answer is an integer between 0 and 999."
)
HUMANEVAL_INSTRUCTION = (
    "Complete the following Python function. Reply with the complete, "
    "self-contained implementation (including the function signature and any "
    "imports it needs) in a single ```python code block. Do not include tests "
    "or example usage.\n\n```python\n{prompt}```"
)


def chat(payload, timeout=3600):
    req = urllib.request.Request(
        SERVER,
        data=json.dumps(payload).encode(),
        headers={"Content-Type": "application/json"},
    )
    start = time.monotonic()
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        body = json.load(resp)
    elapsed = time.monotonic() - start
    msg = body["choices"][0]["message"]
    return {
        "content": msg.get("content") or "",
        "reasoning": msg.get("reasoning_content") or "",
        "finish_reason": body["choices"][0].get("finish_reason"),
        "usage": body.get("usage", {}),
        "wall_s": round(elapsed, 2),
    }


def load_done(path):
    if not path.exists():
        return set()
    return {json.loads(line)["id"] for line in path.open() if line.strip()}


def run_aime(out_path, done):
    problems = json.load((BASE / "data" / "aime25.json").open())
    for p in problems:
        pid = f"aime25-{p['id']}"
        if pid in done:
            continue
        r = chat(
            {
                "model": "deepseek-v4-flash",
                "messages": [
                    {"role": "user", "content": p["problem"] + "\n\n" + AIME_INSTRUCTION}
                ],
                "max_tokens": 24576,
            }
        )
        graded_text = r["content"] if r["content"].strip() else r["reasoning"]
        passed = grade_aime(graded_text, p["answer"])
        yield out_path, {
            "id": pid,
            "passed": passed,
            "expected": p["answer"],
            "finish_reason": r["finish_reason"],
            "usage": r["usage"],
            "wall_s": r["wall_s"],
            "content": r["content"],
            "reasoning_chars": len(r["reasoning"]),
        }


def run_humaneval(out_path, done):
    tasks = [
        json.loads(line)
        for line in (BASE / "data" / "HumanEval.jsonl").open()
        if line.strip()
    ]
    for t in tasks:
        tid = t["task_id"]
        if tid in done:
            continue
        r = chat(
            {
                "model": "deepseek-v4-flash",
                "messages": [
                    {
                        "role": "user",
                        "content": HUMANEVAL_INSTRUCTION.format(prompt=t["prompt"]),
                    }
                ],
                "think": False,
                "temperature": 0,
                "max_tokens": 4096,
            }
        )
        code = extract_code(r["content"])
        passed, detail = run_humaneval_check(code, t["test"], t["entry_point"])
        yield out_path, {
            "id": tid,
            "passed": passed,
            "detail": detail if not passed else "ok",
            "finish_reason": r["finish_reason"],
            "usage": r["usage"],
            "wall_s": r["wall_s"],
            "content": r["content"],
        }


GPQA_INSTRUCTION = (
    "Answer the following multiple-choice question. Think briefly if needed, then "
    "give your final answer on the last line in the exact form 'Answer: X' where X "
    "is one of A, B, C, D."
)


def run_gpqa(out_path, done):
    with (BASE / "data" / "gpqa_diamond.csv").open(newline="") as f:
        rows = list(csv.DictReader(f))
    for i, row in enumerate(rows):
        qid = f"gpqa-diamond-{i}"
        if qid in done:
            continue
        choices = [
            ("correct", row["Correct Answer"].strip()),
            ("wrong", row["Incorrect Answer 1"].strip()),
            ("wrong", row["Incorrect Answer 2"].strip()),
            ("wrong", row["Incorrect Answer 3"].strip()),
        ]
        random.Random(i).shuffle(choices)
        letters = "ABCD"
        expected = letters[[k for k, _ in choices].index("correct")]
        options = "\n".join(f"{letters[j]}) {text}" for j, (_, text) in enumerate(choices))
        prompt = f"{GPQA_INSTRUCTION}\n\n{row['Question'].strip()}\n\n{options}"
        r = chat(
            {
                "model": "deepseek-v4-flash",
                "messages": [{"role": "user", "content": prompt}],
                "think": False,
                "temperature": 0,
                "max_tokens": 3072,
            }
        )
        got = extract_mcq_letter(r["content"])
        yield out_path, {
            "id": qid,
            "passed": got == expected,
            "expected": expected,
            "given": got,
            "finish_reason": r["finish_reason"],
            "usage": r["usage"],
            "wall_s": r["wall_s"],
            "content": r["content"],
        }


def main():
    suite = sys.argv[1] if len(sys.argv) > 1 else ""
    runners = {"aime": run_aime, "humaneval": run_humaneval, "gpqa": run_gpqa}
    if suite not in runners:
        sys.exit(f"usage: harness.py [{'|'.join(runners)}] [result-suffix]")
    suffix = f"-{sys.argv[2]}" if len(sys.argv) > 2 else ""
    out_path = BASE / "results" / f"{suite}{suffix}.jsonl"
    out_path.parent.mkdir(exist_ok=True)
    done = load_done(out_path)
    if done:
        print(f"resuming: {len(done)} items already done", flush=True)
    n_pass = n_run = 0
    for path, record in runners[suite](out_path, done):
        with path.open("a") as f:
            f.write(json.dumps(record) + "\n")
        n_run += 1
        n_pass += bool(record["passed"])
        usage = record["usage"]
        gen = usage.get("completion_tokens", 0)
        tps = round(gen / record["wall_s"], 2) if record["wall_s"] else 0
        print(
            f"[{record['id']}] {'PASS' if record['passed'] else 'FAIL'} "
            f"gen={gen}tok wall={record['wall_s']}s ({tps} t/s) "
            f"running: {n_pass}/{n_run}",
            flush=True,
        )
    print(f"suite={suite} new_pass={n_pass}/{n_run}", flush=True)


if __name__ == "__main__":
    main()