"""Context-engineering mini experiments with Claude Agent SDK (no tools, empty cwd)."""
import asyncio, json, random, re, sys, tempfile, os

MODELS = {"haiku": "claude-haiku-4-5-20251001", "sonnet": "claude-sonnet-5-5"}
N = int(os.environ.get("N", "10"))
LEN = int(os.environ.get("T4_LEN", "120"))
ONLY = os.environ.get("ONLY")
OUT = os.environ.get("OUT", "results.json")
SYSTEM = "あなたは社内の問い合わせに答えるアシスタントです。短く、答えだけを日本語で返してください。"

# ---- 架空の社内ルール（モデルが知りえない情報）----
FACT_OLD = "【経費精算ルール 2025年版（2025-04-01施行）】経費精算の申請締め日は毎月20日です。"
FACT_NEW = "【経費精算ルール 2026年版（2026-04-01施行）】経費精算の申請締め日は毎月25日です。2025年版は廃止されました。"
FACT_NEW_NODATE = "【経費精算ルール】経費精算の申請締め日は毎月25日です。"
FACT_OLD_NODATE = "【経費精算ルール】経費精算の申請締め日は毎月20日です。"
Q = "経費精算の申請締め日は毎月何日ですか？ 分からなければ「分からない」と答えてください。"

def noise_docs(n, seed=1):
    rnd = random.Random(seed)
    topics = ["出張", "備品購入", "在宅勤務手当", "社内勉強会", "来客対応", "名刺発注", "会議室予約", "セキュリティ研修", "健康診断", "年末調整"]
    out = []
    for i in range(n):
        t = rnd.choice(topics)
        out.append(f"【{t}ルール その{i+1}】{t}に関する社内手順です。申請は社内ポータルから行い、承認者は直属の上長です。"
                   f"処理の目安は{rnd.randint(2,9)}営業日、問い合わせ先は総務部の{t}担当です。" + "詳細は別紙を参照してください。" * 3)
    return out

def build_tasks():
    tasks = []
    # T1 情報不足: 渡さない / 渡す
    tasks.append(("T1_missing", "no_context", Q, lambda a: a))
    tasks.append(("T1_missing", "with_fact", FACT_NEW + "\n\n" + Q, lambda a: a))
    # T2 埋没: 事実1件のみ / 事実+ノイズ約40件 / 事実+ノイズ約400件
    for k, label in [(0, "fact_only"), (40, "noise_40"), (400, "noise_400")]:
        docs = noise_docs(k)
        docs.insert(len(docs) // 2, FACT_NEW)
        tasks.append(("T2_noise", label, "\n".join(docs) + "\n\n" + Q, lambda a: a))
    # T3 矛盾: 新旧が混在。日付なし(旧→新の順 / 新→旧の順) / 日付+優先ルール
    tasks.append(("T3_clash", "nodate_old_then_new", FACT_OLD_NODATE + "\n" + FACT_NEW_NODATE + "\n\n" + Q, lambda a: a))
    tasks.append(("T3_clash", "nodate_new_then_old", FACT_NEW_NODATE + "\n" + FACT_OLD_NODATE + "\n\n" + Q, lambda a: a))
    tasks.append(("T3_clash", "dated_new_then_old", FACT_NEW + "\n" + FACT_OLD + "\n\n新旧で内容が違う場合は、施行日が新しい版を正としてください。\n" + Q, lambda a: a))
    # T4 能力限界: 文字数カウント。素のまま / 「丁寧に」「手順」を足したコンテキスト
    rnd = random.Random(7)
    s = "".join(rnd.choice("abcdefgh") for _ in range(LEN))
    q4 = f"次の文字列に、小文字の「e」は何個ありますか？ 数字だけ答えてください。\n{s}"
    q4b = ("あなたは正確さを最優先する校正者です。見落としは許されません。何度も見直してください。"
           "この作業の目的は品質検査の集計で、1個の誤差でも不合格になります。\n" + q4)
    truth = s.count("e")
    tasks.append(("T4_limit", "plain", q4, truth))
    tasks.append(("T4_limit", "more_context", q4b, truth))
    return tasks

def score(task, cond, ans, truth):
    a = ans.strip()
    if task == "T4_limit":
        m = re.findall(r"\d+", a)
        return {"correct": bool(m) and int(m[-1]) == truth}
    has25, has20 = "25日" in a, "20日" in a
    unknown = bool(re.search(r"分からない|分かりません|わからない|わかりません|不明", a))
    cat = "unknown" if unknown and not (has25 != has20) else ("25" if has25 and not has20 else ("20" if has20 and not has25 else "other"))
    if unknown and (has25 or has20): cat = "unknown"
    want = "unknown" if (task == "T1_missing" and cond == "no_context") or (task == "T3_clash" and cond.startswith("nodate")) else "25"
    return {"category": cat, "correct": cat == want}

async def one(task, cond, prompt, truth, mname, i, sem, cwd):
    from claude_agent_sdk import ClaudeAgentOptions, query, ResultMessage, AssistantMessage, TextBlock
    async with sem:
        opts = ClaudeAgentOptions(model=MODELS[mname], system_prompt=SYSTEM, tools=[], allowed_tools=[], setting_sources=[], cwd=cwd, extra_args={"strict-mcp-config": None})
        text, usage, cost = "", {}, None
        async for m in query(prompt=prompt, options=opts):
            if isinstance(m, AssistantMessage):
                for b in m.content:
                    if isinstance(b, TextBlock): text += b.text
            if isinstance(m, ResultMessage):
                usage, cost = m.usage or {}, m.total_cost_usd
        sc = score(task, cond, text, truth)
        r = {"task": task, "condition": cond, "model": mname, "i": i, "answer": text.strip(), "usage": usage, "cost_usd": cost, **sc}
        print(task, cond, mname, i, sc, flush=True)
        return r

async def main():
    if os.path.exists(OUT):
        raise SystemExit("既存データを上書きしないため、未使用のOUTを指定してください")
    sem = asyncio.Semaphore(8)
    cwd = tempfile.mkdtemp(prefix="ce-empty-")
    jobs = []
    async def collect(task, cond, prompt, truth, mname, i):
        try:
            return await one(task, cond, prompt, truth, mname, i, sem, cwd)
        except Exception as exc:
            print(f"FAILED {task}/{cond}/{mname}/{i}: {type(exc).__name__}: {exc}", file=sys.stderr, flush=True)
            return None

    for task, cond, prompt, extra in build_tasks():
        if ONLY and task != ONLY: continue
        truth = extra if isinstance(extra, int) else None
        for mname in MODELS:
            for i in range(N):
                jobs.append((task, cond, prompt, truth, mname, i))
    res, failed = [], 0
    with open(OUT, "x") as output:
        json.dump(res, output)
        output.flush()
        for job in asyncio.as_completed(collect(*args) for args in jobs):
            row = await job
            if row is None:
                failed += 1
                continue
            res.append(row)
            output.seek(0)
            json.dump(res, output, ensure_ascii=False, indent=1)
            output.truncate()
            output.flush()
    if failed:
        raise SystemExit(f"{failed}問い合わせが失敗しました。完了した{len(res)}回答を{OUT}に保存しました")

if __name__ == "__main__":
    asyncio.run(main())