"""E1 token 经济学实验脚本（论文主实验，本地确定性可跑）。

对比四种策略下 Architect（大模型）单请求输入 token 量：
  A1 全量上下文  ：每轮把完整历史+工件全文发给 Architect（无压缩基线）
  A2 交流文本协议：只用 render_for_architect 压缩摘要（D7）
  A3 A2 + rollup ：先把已完成步骤折叠为 archive 摘要行再渲染
  A4 A3 + prefix ：记录可被 --cache-reuse 命中的稳定前缀 token（降低 prefill 成本）

北极星指标（方案 1.0）：A2/A3/A4 相对 A1 的 token 下降 ≥80%。

用法：
    python scripts/bench_tokens.py [--data eval/v2_sample.json] [--out research/v2_experiments]
本地模式：不调用真实 API，用 estimate_tokens 对策略做确定性测量，输出 CSV+MD。
--live 模式（可选，需 API key + 本地模型）：走真实管线记录 usage。
"""
from __future__ import annotations

import argparse
import csv
import json
import sys
from pathlib import Path

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8")
    sys.stderr.reconfigure(encoding="utf-8")

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from router_system.workspace import Workspace, estimate_tokens  # noqa: E402

# 每步模拟工件文本（本地模式用，代表真实产物体量）
_ARTIFACT_TEMPLATE = (
    "（工件）{domain} 步骤实现说明：这是第 {i} 步的完整实现细节与说明文本，"
    "包含关键逻辑、边界处理与可运行示例，长度适中以模拟真实产物。"
)


def _brief_for(query: str, domain: str, n_steps: int = 3) -> dict:
    return {
        "goal": query,
        "constraints": ["遵守领域规范", "输出可交付"],
        "tags": [domain],
        "acceptance": [{"id": "a1", "check": "满足用户需求", "machine_checkable": True}],
        "plan": [
            {"id": f"s{i+1}", "task": f"{domain} 步骤{i+1}：推进目标", "deps": [] if i == 0 else [f"s{i}"],
             "done_criteria": "达到步骤目标"}
            for i in range(n_steps)
        ],
    }


def build_workspace(query: str, domain: str, n_steps: int = 3, n_rounds: int = 3) -> Workspace:
    """构造一个模拟进行到中后期的交流文本（含 progress/issues/decisions）。"""
    ws = Workspace.new("bench" + query.encode("utf-8").hex()[:8], query,
                       api_token_cap=8000, rounds_cap=6)
    ws.apply_brief(_brief_for(query, domain, n_steps))
    # 已完成前 n_rounds 步（至少 1），最后一步待办
    done_steps = max(1, min(n_rounds, n_steps))
    for i in range(done_steps):
        ws.add_progress(f"s{i+1}", "done",
                        f"步骤{i+1}完成：{_ARTIFACT_TEMPLATE.format(domain=domain, i=i+1)[:60]}",
                        artifact=f"a://s{i+1}.py" if domain == "code" else f"a://s{i+1}.md")
    # 加入 issue + decision（模拟一轮裁决）
    if done_steps < n_steps:
        iid = ws.add_issue(f"s{done_steps+1}", f"a://s{done_steps+1}.py#L1",
                           "验证未通过", "达到目标", "已自修 2 次", "请裁决")
        ws.add_decision(iid, "按此方向继续推进", [{"id": f"s{done_steps+1}", "task": "按裁决修订"}])
    ws.mark_round()
    return ws


def _artifact_text(domain: str, i: int) -> str:
    return _ARTIFACT_TEMPLATE.format(domain=domain, i=i)


def measure(ws: Workspace, n_steps: int = 3):
    """测量四种策略的单请求 Architect 输入 token。"""
    domain = (ws.get("brief") or {}).get("tags", ["general"])[0]

    # A1 全量上下文：把完整历史逐字发送（query + brief 全文 + 全部工件全文 +
    #              全部 issues/decisions/progress 全文），无任何压缩。
    a1 = _full_context_tokens(ws, domain, n_steps)

    # A2 交流文本：render_for_architect
    a2 = estimate_tokens(ws.render_for_architect())

    # A3 A2 + rollup
    ws3 = Workspace(ws.data)
    ws3.rollup()
    a3 = estimate_tokens(ws3.render_for_architect())

    # A4 A3 + prefix：token 数同 A3；prefix_hit 为可复用稳定前缀
    prefix_hit = estimate_tokens(_prefix_region(ws))
    return {"a1": a1, "a2": a2, "a3": a3, "a4": a3, "prefix_hit": prefix_hit}


def _full_context_tokens(ws: Workspace, domain: str, n_steps: int) -> int:
    """A1 基线：完整逐字上下文的 token 数。"""
    d = ws.data
    total = estimate_tokens(d.get("query", ""))
    # brief 全文（含 goal/constraints/plan 全部字段）
    total += estimate_tokens(json.dumps(d.get("brief"), ensure_ascii=False))
    # 全部工件全文
    total += sum(estimate_tokens(_artifact_text(domain, i + 1)) for i in range(n_steps))
    # issues / decisions / progress 全文
    for iss in d.get("issues", []) or []:
        total += estimate_tokens(json.dumps(iss, ensure_ascii=False))
    for dec in d.get("decisions", []) or []:
        total += estimate_tokens(json.dumps(dec, ensure_ascii=False))
    for p in d.get("progress", []) or []:
        total += estimate_tokens(json.dumps(p, ensure_ascii=False))
    return total


def _prefix_region(ws: Workspace) -> str:
    """稳定前缀（可被 prefix cache 命中）的文本。"""
    d = ws.data
    stable = {"version": d.get("version"), "request_id": d.get("request_id"),
              "query": d.get("query"), "brief": d.get("brief")}
    return json.dumps(stable, ensure_ascii=False)


def run(data_path: str, out_dir: str, n_steps: int = 3) -> None:
    items = json.loads(Path(data_path).read_text(encoding="utf-8"))
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    rows = []
    for it in items:
        ws = build_workspace(it["query"], it.get("domain", "general"), n_steps)
        m = measure(ws, n_steps)
        rows.append({
            "id": it["id"], "domain": it.get("domain", "general"),
            "a1_full": m["a1"], "a2_ws": m["a2"], "a3_rollup": m["a3"],
            "a4_prefix": m["a4"], "prefix_hit": m["prefix_hit"],
            "reduction_a2": round(1 - m["a2"] / m["a1"], 4) if m["a1"] else 0,
            "reduction_a4": round(1 - m["a4"] / m["a1"], 4) if m["a1"] else 0,
        })

    # CSV
    csv_path = out / "E1_token_economics.csv"
    with open(csv_path, "w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
        w.writeheader()
        w.writerows(rows)

    # 聚合
    n = len(rows)
    avg = {k: round(sum(r[k] for r in rows) / n, 2) for k in
           ("a1_full", "a2_ws", "a3_rollup", "a4_prefix", "prefix_hit")}
    red_a2 = round(1 - avg["a2_ws"] / avg["a1_full"], 4)
    red_a4 = round(1 - avg["a4_prefix"] / avg["a1_full"], 4)

    md = _render_md(rows, avg, red_a2, red_a4)
    (out / "E1_token_economics.md").write_text(md, encoding="utf-8")
    print(f"写入: {csv_path}")
    print(f"写入: {out / 'E1_token_economics.md'}")
    print(f"汇总: A1={avg['a1_full']} A2={avg['a2_ws']} A3={avg['a3_rollup']} "
          f"A4={avg['a4_prefix']} prefix_hit={avg['prefix_hit']}")
    print(f"token 下降: A2 相对 A1 = {red_a2*100:.1f}%  |  A4 相对 A1 = {red_a4*100:.1f}%")


def _render_md(rows, avg, red_a2, red_a4) -> str:
    lines = [
        "# E1 token 经济学（本地确定性测量）",
        "",
        "> 模式：本地 estimate_tokens 测量（不调用真实 API）。真实数据需 --live + API key + 本地模型。",
        "",
        f"- 样例数：{len(rows)}",
        f"- A1 全量上下文均值：**{avg['a1_full']} token**",
        f"- A2 交流文本均值：**{avg['a2_ws']} token**",
        f"- A3 A2+rollup 均值：**{avg['a3_rollup']} token**",
        f"- A4 A3+prefix 均值：**{avg['a4_prefix']} token**（prefix 可命中 {avg['prefix_hit']} token）",
        "",
        f"## 北极星指标（token 下降 ≥80%）",
        "",
        f"- A2 相对 A1：**{red_a2*100:.1f}%**",
        f"- A4 相对 A1：**{red_a4*100:.1f}%**",
        "",
        "### 说明（诚实解读）",
        "",
        "1. 本报告为本地确定性测量（estimate_tokens），未调用真实 API。",
        "2. A3（rollup）收益为规模相关：小样例下 archive 增量可能抵消收益，长会话才显现。",
        "3. 前缀稳定性（T10）已验证，配合 llama-server --cache-reuse 可复用稳定前缀。",
        "4. 北极星 ≥80% 需在 --live 模式（API key + 本地模型）下由 E1 实验确认。",
        "",
        "## 明细",
        "",
        "| id | domain | A1 | A2 | A3 | A4 | prefix_hit |",
        "|----|--------|----|----|----|----|----|",
    ]
    for r in rows:
        lines.append(f"| {r['id']} | {r['domain']} | {r['a1_full']} | {r['a2_ws']} | "
                     f"{r['a3_rollup']} | {r['a4_prefix']} | {r['prefix_hit']} |")
    return "\n".join(lines)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", default="eval/v2_sample.json")
    ap.add_argument("--out", default="research/v2_experiments")
    ap.add_argument("--steps", type=int, default=3)
    ap.add_argument("--live", action="store_true", help="真实 API（需 key + 本地模型）")
    args = ap.parse_args()
    if args.live:
        print("[warn] --live 需 API key + 本地 llama-server；当前未实现自动跑数，请接入后使用。")
    run(args.data, args.out, args.steps)
    return 0


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