import anyio, asyncio, json, shutil, sys, os, time
from pathlib import Path
from claude_agent_sdk import ClaudeSDKClient, ClaudeAgentOptions, AssistantMessage, ToolUseBlock, ResultMessage

ROOT = Path("/tmp/exp-slide")
MAT = ROOT / "materials"
OUT = ROOT / "runs"
MODELS = {"haiku": "claude-haiku-4-5-20251001", "sonnet": "claude-sonnet-5-5", "opus": "claude-opus-5-5"}

BASE = ("同じフォルダにある資料（metrics.csv、regions.csv、meeting-notes.md）をもとに、"
        "経営会議で第2四半期の業績を報告するスライドを8枚作ってください。"
        "Marp形式のMarkdownにして、slides.md に保存します。スライドは --- で区切ります。")
RULES = ("\n\n資料の扱い：売上と粗利益の正本は metrics.csv です。地域別の内訳は regions.csv を使います。"
         "meeting-notes.md は速報・未承認の口頭メモなので、数字の正本にはしません。"
         "Q3の170百万円は提案値で、決まった目標ではありません。")
LEDGER = ("\n\n作る前に、スライドで使う数字をすべて numbers.md の台帳にまとめてください。"
          "1行に1つの数字で、値・単位・計算式・出典ファイル・区分（確定／計算／未承認／仮説）を書きます。"
          "スライドには台帳の数字だけを使い、区分が未承認・仮説のものは、そう分かる書き方にします。"
          "slides.md を保存したら、全行を台帳と元資料に照合し、誤りがあれば直してから終えてください。")
REVIEW_SELF = ("いま作った slides.md を、元資料（metrics.csv、regions.csv、meeting-notes.md）と1行ずつ照合してください。"
               "数字の誤り、比較の基準の取り違え、根拠のない断定があれば slides.md を直し、直した箇所を一覧で教えてください。")
REVIEW_FRESH = ("slides.md は、同じフォルダの資料（metrics.csv、regions.csv、meeting-notes.md）をもとに、"
                "経営会議向けの第2四半期報告として作られたスライドです。"
                "元資料と1行ずつ照合してください。数字の誤り、比較の基準の取り違え、根拠のない断定があれば slides.md を直し、直した箇所を一覧で教えてください。")

def opts(cwd, model):
    return ClaudeAgentOptions(
        model=MODELS[model], cwd=str(cwd), tools=["Read", "Write"], allowed_tools=["Read", "Write"],
        setting_sources=[], extra_args={"strict-mcp-config": None}, permission_mode="bypassPermissions", max_turns=30)

async def turn(client, prompt, log):
    await client.query(prompt)
    text = ""
    async for m in client.receive_response():
        if isinstance(m, AssistantMessage):
            for b in m.content:
                if isinstance(b, ToolUseBlock):
                    log.append({"tool": b.name, "input": b.input})
                elif hasattr(b, "text"):
                    text += b.text
        elif isinstance(m, ResultMessage):
            log.append({"result": {"cost": m.total_cost_usd, "turns": m.num_turns, "usage": m.usage}})
    return text

def mk(name):
    d = OUT / name
    if d.exists(): shutil.rmtree(d)
    d.mkdir(parents=True)
    for f in MAT.iterdir(): shutil.copy(f, d / f.name)
    return d

async def run_chain(model, i, sem):
    async with sem:
        try:
            # A (+D)
            d = mk(f"{model}-A{i:02d}"); log = []
            async with ClaudeSDKClient(opts(d, model)) as c:
                t1 = await turn(c, BASE, log)
                shutil.copy(d / "slides.md", d / "slides_v1.md")
                (d / "reply1.txt").write_text(t1)
                log.append({"marker": "review_self"})
                t2 = await turn(c, REVIEW_SELF, log)
                shutil.copy(d / "slides.md", d / "slides_v2.md")
                (d / "reply2.txt").write_text(t2)
            (d / "log.json").write_text(json.dumps(log, ensure_ascii=False, indent=1))
            # E: fresh reviewer on A's v1
            e = mk(f"{model}-E{i:02d}"); log = []
            shutil.copy(d / "slides_v1.md", e / "slides.md")
            async with ClaudeSDKClient(opts(e, model)) as c:
                t = await turn(c, REVIEW_FRESH, log)
            shutil.copy(e / "slides.md", e / "slides_v2.md")
            (e / "reply.txt").write_text(t); (e / "log.json").write_text(json.dumps(log, ensure_ascii=False, indent=1))
        except Exception as ex:
            print("FAIL chain", model, i, repr(ex)); 
        print("done chain", model, i, flush=True)

async def run_single(cond, model, i, sem):
    async with sem:
        try:
            d = mk(f"{model}-{cond}{i:02d}"); log = []
            prompt = BASE + RULES + (LEDGER if cond == "C" else "")
            async with ClaudeSDKClient(opts(d, model)) as c:
                t = await turn(c, prompt, log)
            (d / "reply.txt").write_text(t); (d / "log.json").write_text(json.dumps(log, ensure_ascii=False, indent=1))
        except Exception as ex:
            print("FAIL", cond, model, i, repr(ex))
        print("done", cond, model, i, flush=True)

async def main():
    models = sys.argv[1].split(","); n0, n1 = int(sys.argv[2]), int(sys.argv[3]); conc = int(sys.argv[4])
    sem = asyncio.Semaphore(conc)
    tasks = []
    for m in models:
        for i in range(n0, n1):
            tasks += [run_chain(m, i, sem), run_single("B", m, i, sem), run_single("C", m, i, sem)]
    await asyncio.gather(*tasks)
asyncio.run(main())