"""単体・単体+見直し・分担・分担+見直しを、同じ課題で比べる。使い方:
    python3 run.py <model> <reps> <out.jsonl> [E|R|ER]
道具は渡さない（テキストを読んで、テキストを返すだけ）。"""
import asyncio, json, os, re, sys, tempfile, time
import materials as m

SYSTEM = "あなたは正確に作業する担当者です。道具は使えません。指示された出力だけを返し、前置きや解説は書きません。"
CWD = tempfile.mkdtemp(prefix="maexp-")
SEM = asyncio.Semaphore(8)


async def call(model, prompt, retries=2):
    """claude -p を1回呼ぶ。トークンは、指定モデルの modelUsage から読む。"""
    env = dict(os.environ, CLAUDE_CODE_DISABLE_AUTO_MEMORY="1")
    for attempt in range(retries + 1):
        async with SEM:
            p = await asyncio.create_subprocess_exec(
                "claude", "-p", "--model", model, "--tools", "", "--setting-sources", "",
                "--strict-mcp-config", "--disable-slash-commands", "--no-session-persistence",
                "--system-prompt", SYSTEM, "--output-format", "json",
                cwd=CWD, env=env, stdin=asyncio.subprocess.PIPE,
                stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)
            out, err = await p.communicate(prompt.encode())
        try:
            j = json.loads(out)
            if j.get("is_error"):
                raise RuntimeError(j.get("result"))
            mu = [v for k, v in j["modelUsage"].items() if k.startswith(model.rsplit("-20", 1)[0])][0]
            return {
                "text": j["result"], "input": mu["inputTokens"], "output": mu["outputTokens"],
                "thinking": mu.get("thinkingTokens", 0),
                "cache_read": mu["cacheReadInputTokens"], "cache_write": mu["cacheCreationInputTokens"],
                "cost_usd": mu["costUSD"], "ms": j["duration_ms"], "model_key": [k for k in j["modelUsage"] if k.startswith(model.rsplit("-20", 1)[0])][0],
            }
        except Exception as e:  # noqa
            last = f"{e!r} {out[:200]!r} {err[:200]!r}"
            await asyncio.sleep(2 + attempt * 3)
    return {"text": "", "error": last, "input": 0, "output": 0, "thinking": 0, "cache_read": 0, "cache_write": 0, "cost_usd": 0, "ms": 0}


def strip_fence(t):
    t = t.strip()
    mm = re.match(r"^```[a-z]*\n(.*?)\n```$", t, re.S)
    return mm.group(1) if mm else t


def parse_json(t):
    t = strip_fence(t)
    i, k = t.find("{"), t.rfind("}")
    try:
        return json.loads(t[i:k + 1])
    except Exception:
        return None


# ------------------------------------------------------------------ 編集
def p_edit_single():
    return f"""次の文書を、変更依頼どおりに直してください。

# 変更依頼
{m.edit_request()}

# 文書
{m.EDIT_DOC_L}

修正後の文書の全文だけを、Markdownのまま出力してください。"""


def p_edit_review(draft):
    return f"""文書に変更を加える仕事の、見直しをします。元の文書、変更依頼、初稿を渡します。
依頼の6項目それぞれについて、初稿が元の文書と見比べて正しく直せているかを点検し、足りない・間違っている箇所は直してください。依頼に関係ない部分を変えていたら、元に戻してください。

# 変更依頼
{m.edit_request()}

# 元の文書
{m.EDIT_DOC_L}

# 初稿
{draft}

最終版の文書の全文だけを、Markdownのまま出力してください。"""


OWNER_NOTE = "\n新しい章（第3章の直後）を書いて挿入するのは、第4〜7章の担当者の仕事です。ほかの担当者は書きません。"


def p_edit_part(i, part, assigned=False):
    names = m.part_names("L")
    return f"""文書を3人で分担して、変更依頼どおりに直しています。あなたは{names[i]}の担当です。ほかの部分は別の担当者が同時に直していて、あなたには見えません。{OWNER_NOTE if assigned else ""}

# 変更依頼
{m.edit_request()}

# あなたの担当部分
{part}

あなたの担当部分を直して、その部分の全文だけを、Markdownのまま出力してください。ほかの章は書かないでください。"""


async def edit_condition(model, cond):
    calls = []
    if cond in ("E1", "E2", "E7"):
        c = await call(model, p_edit_single()); calls.append(c)
        text = strip_fence(c["text"])
        for _ in range({"E1": 0, "E2": 1, "E7": 2}[cond]):
            c2 = await call(model, p_edit_review(text)); calls.append(c2)
            text = strip_fence(c2["text"])
    else:
        parts = m.split_edit_doc("L")
        assigned = cond in ("E5", "E6")
        cs = await asyncio.gather(*[call(model, p_edit_part(i, p, assigned)) for i, p in enumerate(parts)])
        calls += cs
        text = "\n\n".join(strip_fence(c["text"]).strip() for c in cs) + "\n"
        if cond in ("E4", "E6"):
            c2 = await call(model, p_edit_review(text)); calls.append(c2)
            text = strip_fence(c2["text"])
    chk = m.grade_edit(text, "L")
    return {"calls": calls, "final": text, "checks": chk}


# ------------------------------------------------------------------ 調べもの
def pages_block(names=None):
    pg = m.source_pages("L")
    return "\n\n---\n\n".join(pg[n] for n in (names or pg))


JSON_FMT = '{"rows": {"製品名": {"yearly_yen": 0, "max_file_mb": 0, "region": "", "free_seats": 0}, ...}, "q1": "製品名", "q2": "製品名"}'


def p_res_single():
    return f"""20の製品ガイドを読んで答えてください。

{m.FIELDS_DESC}

{m.CROSS_Q}

# 製品ガイド
{pages_block()}

JSONだけを出力してください。形式: {JSON_FMT}"""


def p_res_review(draft):
    return f"""調べものの見直しをします。製品ガイドと、初稿の答え（JSON）を渡します。
80個の値と q1・q2 を、ガイドの原文と照らして一つずつ確かめ、間違いがあれば直してください。

{m.FIELDS_DESC}

{m.CROSS_Q}

# 製品ガイド
{pages_block()}

# 初稿
{json.dumps(draft, ensure_ascii=False)}

最終版のJSONだけを出力してください。形式: {JSON_FMT}"""


def p_res_worker(name):
    return f"""製品ガイドを読んで答えてください。

{m.FIELDS_DESC}

# 製品ガイド（{name}）
{pages_block([name])}

JSONだけを出力してください。形式: {{"yearly_yen": 0, "max_file_mb": 0, "region": "", "free_seats": 0}}"""


def p_res_integrate(rows):
    return f"""20の製品について、担当者が調べた結果の表があります。この表だけを使って答えてください。

{m.CROSS_Q}

# 表
{json.dumps(rows, ensure_ascii=False)}

JSONだけを出力してください。形式: {{"q1": "製品名", "q2": "製品名"}}"""


def grade_res(ans):
    gold = m.research_gold("L")
    cells = {}
    rows = (ans or {}).get("rows") or {}
    for n in [p[0] for p in m.products("L")]:
        r = rows.get(n) or {}
        for f, g in gold[n].items():
            v = r.get(f)
            if isinstance(g, str):
                v = str(v).replace("リージョン", "").strip() if v is not None else None
            else:
                try:
                    v = int(round(float(str(v).replace(",", ""))))
                except Exception:
                    v = None
            cells[f"{n}.{f}"] = (v == g)
    cells["q1"] = (ans or {}).get("q1") == gold["_q1"]
    cells["q2"] = (ans or {}).get("q2") == gold["_q2"]
    return cells


async def res_condition(model, cond):
    calls = []
    if cond in ("R1", "R2", "R5"):
        c = await call(model, p_res_single()); calls.append(c)
        ans = parse_json(c["text"])
        for _ in range({"R1": 0, "R2": 1, "R5": 2}[cond]):
            c2 = await call(model, p_res_review(ans)); calls.append(c2)
            ans = parse_json(c2["text"])
    else:
        names = [p[0] for p in m.products("L")]
        cs = await asyncio.gather(*[call(model, p_res_worker(n)) for n in names])
        calls += cs
        rows = {n: (parse_json(c["text"]) or {}) for n, c in zip(names, cs)}
        c2 = await call(model, p_res_integrate(rows)); calls.append(c2)
        q = parse_json(c2["text"]) or {}
        ans = {"rows": rows, "q1": q.get("q1"), "q2": q.get("q2")}
        if cond == "R4":
            c3 = await call(model, p_res_review(ans)); calls.append(c3)
            ans = parse_json(c3["text"])
    return {"calls": calls, "final": ans, "checks": grade_res(ans)}


async def one(model, cond, rep, out):
    t0 = time.time()
    r = await (edit_condition if cond.startswith("E") else res_condition)(model, cond)
    r.update(model=model, cond=cond, rep=rep, wall_s=round(time.time() - t0, 1))
    with open(out, "a") as f:
        f.write(json.dumps(r, ensure_ascii=False) + "\n")
    ok = sum(r["checks"].values()), len(r["checks"])
    print(model, cond, rep, ok, flush=True)


async def main():
    model, reps, out = sys.argv[1], int(sys.argv[2]), sys.argv[3]
    which = sys.argv[4] if len(sys.argv) > 4 else "ER"
    conds = (["E1", "E2", "E3", "E4", "E5", "E6", "E7"] if "E" in which else []) + (["R1", "R2", "R3", "R4", "R5"] if "R" in which else [])
    done = set()
    if os.path.exists(out):
        for l in open(out):
            d = json.loads(l); done.add((d["cond"], d["rep"]))
    await asyncio.gather(*[one(model, c, r, out) for r in range(reps) for c in conds if (c, r) not in done])


asyncio.run(main())