#!/usr/bin/env python3
"""额度/多模型分流教程实测。通过 ofox OpenAI 兼容接口真实调用，记录 usage、耗时、得分。

用法：runner.py e1|e2|e3|e4
  e1 会议纪要 → 行动项 JSON（8 个真值，自动比对 owner+日期）
  e2 修 bug（8 个隐藏测试，子进程执行）
  e3 长线程 vs 摘要重开：每轮真实 prompt_tokens
  e4 同一修 bug 任务，reasoning_effort low vs high
费用 = usage × catalog 快照牌价（pricing 字段），见 evidence/catalog-*.json。
"""
import glob
import json
import os
import re
import subprocess
import sys
import time
import urllib.request
from datetime import datetime, timezone
from pathlib import Path

BASE = Path(__file__).resolve().parent.parent
RUNS = BASE / "runs"
INP = BASE / "inputs"
KEY = os.environ["OFOX_API_KEY"]
CATALOG = sorted(glob.glob(str(BASE / "evidence" / "catalog-*.json")))[-1]
PRICES = {m["id"]: m["pricing"] for m in json.load(open(CATALOG))["data"]}

E1_MODELS = ["anthropic/claude-opus-5.5", "anthropic/claude-sonnet-5.5", "openai/gpt-6-sol",
             "openai/gpt-6-luna", "deepseek/deepseek-v4.1-flash", "qwen/qwen3.8-flash",
             "anthropic/claude-haiku-5.5"]
E2_MODELS = E1_MODELS
REPS = 2


def call(model, messages, **extra):
    body = {"model": model, "messages": messages, "max_tokens": 4000, **extra}
    req = urllib.request.Request("https://api.ofox.ai/v1/chat/completions",
                                 data=json.dumps(body).encode(),
                                 headers={"Authorization": f"Bearer {KEY}",
                                          "Content-Type": "application/json"})
    t = time.time()
    with urllib.request.urlopen(req, timeout=300) as r:
        d = json.load(r)
    return d, round(time.time() - t, 2)


def cost(model, usage):
    p = PRICES[model]
    pt = usage.get("prompt_tokens", 0)
    ct = usage.get("completion_tokens", 0)
    cached = (usage.get("prompt_tokens_details") or {}).get("cached_tokens", 0) or 0
    inp = float(p["input"]) * (pt - cached) + float(p.get("input_cache_read", p["input"])) * cached
    return round(inp + float(p["output"]) * ct, 6)


def log(exp, row):
    RUNS.mkdir(exist_ok=True)
    with open(RUNS / f"{exp}.jsonl", "a") as f:
        f.write(json.dumps(row, ensure_ascii=False) + "\n")
    print(json.dumps({k: row[k] for k in row if k not in ("output", "raw")}, ensure_ascii=False))


def base_row(exp, model, rep, d, secs):
    u = d.get("usage", {})
    return {"exp": exp, "model": model, "returned_model": d.get("model"), "rep": rep,
            "ts": datetime.now(timezone.utc).isoformat(), "id": d.get("id"), "seconds": secs,
            "usage": u, "cost_usd": cost(model, u),
            "output": d["choices"][0]["message"].get("content") or ""}


# ---------- e1 ----------
E1_PROMPT = """Extract every action item from the meeting notes below.
Only include items that were actually assigned to a named person WITH a due date. Do not include decisions, ideas that were not assigned, or items with no date.
Return ONLY a JSON array, no prose. Each element: {{"owner": "<first name>", "task": "<short task>", "due": "MM-DD"}}.

Meeting notes:
{notes}"""


def score_e1(text):
    gt = json.load(open(INP / "ground_truth.json"))["action_items"]
    m = re.search(r"\[.*\]", text, re.S)
    try:
        items = json.loads(m.group(0)) if m else []
    except Exception:
        return {"valid_json": False, "hit": 0, "extra": None}
    used, hit = set(), 0
    for g in gt:
        for i, it in enumerate(items):
            if i in used:
                continue
            due = re.sub(r"^\d{4}-", "", str(it.get("due", "")))
            if str(it.get("owner", "")).strip().lower() == g["owner"].lower() and due == g["due"]:
                used.add(i)
                hit += 1
                break
    return {"valid_json": True, "hit": hit, "of": len(gt), "returned": len(items), "extra": len(items) - len(used)}


def e1():
    notes = (INP / "meeting_notes.md").read_text()
    for model in E1_MODELS:
        for rep in range(REPS):
            try:
                d, s = call(model, [{"role": "user", "content": E1_PROMPT.format(notes=notes)}])
                row = base_row("e1", model, rep, d, s)
                row["score"] = score_e1(row["output"])
            except Exception as e:
                row = {"exp": "e1", "model": model, "rep": rep, "error": str(e)[:300]}
            log("e1", row)


# ---------- e2 ----------
E2_PROMPT = """The Python function below has bugs. Read the docstring carefully and fix ALL bugs so it fully matches the docstring.
Return ONLY the complete fixed function in one ```python code block. Keep the same function name and signature.

```python
{code}
```"""


def score_e2(text):
    m = re.search(r"```(?:python)?\n(.*?)```", text, re.S)
    code = m.group(1) if m else text
    prog = code + "\n\nimport sys\nsys.path.insert(0, %r)\nimport hidden_tests\nprint(hidden_tests.run(merge_bookings))\n" % str(INP)
    try:
        r = subprocess.run([sys.executable, "-I", "-c", prog], capture_output=True, text=True, timeout=20)
        mm = re.search(r"\((\d+), (\d+)\)", r.stdout)
        return {"passed": int(mm.group(1)), "of": int(mm.group(2))} if mm else {"passed": 0, "of": 8, "err": r.stderr[-200:]}
    except Exception as e:
        return {"passed": 0, "of": 8, "err": str(e)[:200]}


def e2(models=E2_MODELS, exp="e2", extra=None, reps=REPS):
    code = (INP / "buggy_code.py").read_text()
    for model in models:
        for rep in range(reps):
            try:
                d, s = call(model, [{"role": "user", "content": E2_PROMPT.format(code=code)}], **(extra or {}))
                row = base_row(exp, model, rep, d, s)
                row["score"] = score_e2(row["output"])
                if extra:
                    row["params"] = extra
            except Exception as e:
                row = {"exp": exp, "model": model, "rep": rep, "params": extra, "error": str(e)[:300]}
            log(exp, row)


# ---------- e3 ----------
FOLLOWUPS = [
    "Summarize these meeting notes in 5 bullet points.",
    "Rewrite the summary for an executive audience in 3 sentences.",
    "List the risks mentioned in the notes.",
    "Draft a short Slack message announcing the v3.4 decision.",
    "Make the Slack message friendlier.",
    "Translate the Slack message into Spanish.",
    "Draft an email to leadership about the WAU correction.",
    "Shorten that email to under 80 words.",
]


def e3(model="openai/gpt-6-luna"):
    notes = (INP / "meeting_notes.md").read_text()
    # A：同一线程一直聊（每轮带完整历史）
    hist = [{"role": "user", "content": "Here are meeting notes:\n\n" + notes + "\n\n" + FOLLOWUPS[0]}]
    for turn, q in enumerate(FOLLOWUPS):
        if turn > 0:
            hist.append({"role": "user", "content": q})
        d, s = call(model, hist)
        row = base_row("e3", model, turn, d, s)
        row["strategy"] = "single_long_thread"
        log("e3", row)
        hist.append({"role": "assistant", "content": row["output"]})
    # B：先做一次摘要，之后每个任务新开对话，只带摘要
    d, s = call(model, [{"role": "user", "content": "Write a compact factual brief (max 200 words) of these meeting notes, keeping all names, dates and decisions:\n\n" + notes}])
    brief_row = base_row("e3", model, -1, d, s)
    brief_row["strategy"] = "brief_then_fresh"
    log("e3", brief_row)
    brief = brief_row["output"]
    for turn, q in enumerate(FOLLOWUPS):
        d, s = call(model, [{"role": "user", "content": "Context brief:\n" + brief + "\n\nTask: " + q}])
        row = base_row("e3", model, turn, d, s)
        row["strategy"] = "brief_then_fresh"
        log("e3", row)


# ---------- e5 难版：改派、改期、取消、相对日期 ----------
def score_e5(text):
    gt = json.load(open(INP / "ground_truth_hard.json"))["action_items"]
    m = re.search(r"\[.*\]", text, re.S)
    try:
        items = json.loads(m.group(0)) if m else []
    except Exception:
        return {"valid_json": False, "hit": 0, "of": len(gt), "extra": None}
    used, hit, missed = set(), 0, []
    for g in gt:
        ok = False
        for i, it in enumerate(items):
            if i in used:
                continue
            due = re.sub(r"^\d{4}-", "", str(it.get("due", "")))
            if str(it.get("owner", "")).strip().lower() == g["owner"].lower() and due in g["due"]:
                used.add(i)
                hit += 1
                ok = True
                break
        if not ok:
            missed.append(f'{g["owner"]} {g["due"][0]}')
    extra = [f'{it.get("owner")} {it.get("due")}' for i, it in enumerate(items) if i not in used]
    return {"valid_json": True, "hit": hit, "of": len(gt), "returned": len(items),
            "extra": len(extra), "missed": missed, "extra_items": extra,
            "perfect": hit == len(gt) and not extra}


E5B_CONFIGS = [  # 2026-10-20 实际运行过的 e5b 配置（推理模型放宽上限 / 低推理档）
    ("deepseek/deepseek-v4.1-flash", {"max_tokens": 16000}),
    ("qwen/qwen3.8-flash", {"max_tokens": 16000}),
    ("deepseek/deepseek-v4.1-flash", {"reasoning_effort": "low"}),
    ("qwen/qwen3.8-flash", {"reasoning_effort": "low"}),
]
E5C_CONFIGS = [("openai/gpt-6-sol", {"reasoning_effort": "low"})]  # 路由表 hard 档的证据


def e5_configs(exp, configs, reps=3):
    notes = (INP / "meeting_notes_hard.md").read_text()
    for model, extra in configs:
        for rep in range(reps):
            for attempt in range(3):
                try:
                    d, s = call(model, [{"role": "user", "content": E1_PROMPT.format(notes=notes)}], **extra)
                    row = base_row(exp, model, rep, d, s)
                    row["params"] = extra
                    row["score"] = score_e5(row["output"])
                    break
                except Exception as e:
                    row = {"exp": exp, "model": model, "rep": rep, "params": extra, "error": str(e)[:300]}
                    time.sleep(20)
            log(exp, row)


def e5(reps=3):
    notes = (INP / "meeting_notes_hard.md").read_text()
    for model in E1_MODELS:
        for rep in range(reps):
            for attempt in range(3):
                try:
                    d, s = call(model, [{"role": "user", "content": E1_PROMPT.format(notes=notes)}])
                    row = base_row("e5", model, rep, d, s)
                    row["score"] = score_e5(row["output"])
                    break
                except Exception as e:
                    row = {"exp": "e5", "model": model, "rep": rep, "error": str(e)[:300]}
                    time.sleep(20)
            log("e5", row)


if __name__ == "__main__":
    which = sys.argv[1]
    if which == "e5":
        e5()
        sys.exit()
    if which == "e5b":
        e5_configs("e5b", E5B_CONFIGS)
        sys.exit()
    if which == "e5c":
        e5_configs("e5c", E5C_CONFIGS)
        sys.exit()
    if which == "e2retry":
        e2(models=["qwen/qwen3.8-flash"], exp="e2", reps=1)
        sys.exit()
    if which == "e1":
        e1()
    elif which == "e2":
        e2()
    elif which == "e3":
        e3()
    elif which == "e4":
        for eff in ("low", "high"):
            e2(models=["openai/gpt-6-sol"], exp="e4", extra={"reasoning_effort": eff}, reps=3)