"""
=============================================================================
第 17 讲 · 代码 05（可选）：让真实的 Claude 在迷你 WebShop 里买东西
=============================================================================

对应讲义：lectures/lecture17_frontiers.md §6

【这个脚本要回答的问题】
    代码 01 里的智能体都是规则式的：按原词匹配属性，读不懂换了说法的页面。真实的模型读得懂"water resistant"就是"waterproof"吗？
    它会不会比较几件商品、记得选选项、注意价格上限？
    和 WebShop 一样的文本接口：每一步只能 search[查询] 或 choose[按钮]（商品编号、选项值、Buy、Back to Search）。
    默认 6 条指令，每条最多 12 步（每一步一次调用）。

【需要什么】pip install anthropic；设置 ANTHROPIC_API_KEY（或 `ant auth login`）。**会花钱**：最多 72 次调用。
    没有 SDK 或凭证时打印说明后正常退出。
【运行】
    python3 05_real_llm.py           # 默认 claude-opus-5-5，effort=low
    python3 05_real_llm.py --toy     # 不联网：用一个简单的规则式策略（检查脚本本身）
"""

import argparse
import os
import re
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "lecture01"))
from console import print_header, print_observations, print_table  # noqa: E402
from frontier_lab import SearchEngine, make_catalog, make_instructions, reward  # noqa: E402
from real_llm import DEFAULT_MODEL, ClaudeLM, sdk_available  # noqa: E402

SYSTEM = ("You are shopping on a website for a user. Each turn, reply with exactly one action: search[query] (only on the search page) "
          "or choose[button] where button is a product id (like P0123), an option value, Buy, or Back to Search. "
          "Buy the item that best matches the instruction (type, attributes, options, price).")
MAX_STEPS = 12


class Env:
    """文本版的 WebShop 页面：search → 结果页（10 条：编号、标题、价格）→ 商品页（描述、选项、Buy）。"""

    def __init__(self, instr, engine, catalog):
        self.instr, self.engine = instr, engine
        self.by_id = {p.pid: p for p in catalog}
        self.page, self.results, self.product, self.chosen = "search", [], None, {}

    def observe(self):
        head = f"Instruction: {self.instr.text}\n"
        if self.page == "search":
            return head + "[Search page] Use search[...]"
        if self.page == "results":
            return head + "[Results]\n" + "\n".join(f"{p.pid}: {p.title} — {p.price} dollars" for p in self.results) + "\n[Back to Search]"
        p = self.product
        opts = "\n".join(f"{name}: {', '.join(vals)}" for name, vals in p.options)
        return (head + f"[Product {p.pid}] {p.title}\nPrice: {p.price} dollars\n{p.description}\nOptions:\n{opts}\n"
                f"Selected: {self.chosen or 'none'}\n[Buy] [Back to Search]")

    def step(self, action):
        """【返回】(结束了吗, 奖励)"""
        m = re.match(r"\s*(search|choose)\[(.*)\]", action or "")
        if not m:
            return False, 0.0
        kind, arg = m.group(1), m.group(2).strip()
        if kind == "search" and self.page == "search":
            self.results, self.page = self.engine.search(arg, 10), "results"
        elif kind == "choose" and arg.lower() == "back to search":
            self.page, self.product, self.chosen = "search", None, {}
        elif kind == "choose" and self.page == "results" and arg in self.by_id:
            self.product, self.page, self.chosen = self.by_id[arg], "product", {}
        elif kind == "choose" and self.page == "product" and arg.lower() == "buy":
            return True, reward(self.instr, self.product, self.chosen)
        elif kind == "choose" and self.page == "product":
            for name, vals in self.product.options:
                if arg in vals:
                    self.chosen[name] = arg
        return False, 0.0


def toy_policy(env):
    """--toy 用的规则式策略：搜类型 + 属性，点第一个，选指令里的选项，买。"""
    ins = env.instr
    if env.page == "search":
        return f"search[{' '.join(sorted(ins.attributes))} {ins.category}]"
    if env.page == "results":
        return f"choose[{env.results[0].pid}]"
    for name, v in ins.options:
        if env.chosen.get(name) != v and v in dict(env.product.options).get(name, ()):
            return f"choose[{v}]"
    return "choose[Buy]"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default=DEFAULT_MODEL)
    ap.add_argument("--effort", default="low")
    ap.add_argument("--toy", action="store_true")
    args = ap.parse_args()
    cat = make_catalog()
    eng = SearchEngine(cat)
    instrs = make_instructions(cat)[:6]
    if not args.toy:
        ok, why = sdk_available()
        if not ok:
            print(why)
            return
        lm = ClaudeLM(model=args.model, effort=args.effort, max_tokens=2000, stop_sequences=())
        name = lm.name
    else:
        name = "toy（规则式策略）"
    rows, total = [], 0.0
    try:
        for ins in instrs:
            env = Env(ins, eng, cat)
            history = [{"role": "system", "content": SYSTEM}]
            r, steps, actions = 0.0, 0, []
            for steps in range(1, MAX_STEPS + 1):
                if args.toy:
                    act = toy_policy(env)
                else:
                    history.append({"role": "user", "content": env.observe()})
                    act = lm(history)
                    history.append({"role": "assistant", "content": act})
                actions.append(act.strip()[:30])
                done, r = env.step(act)
                if done:
                    break
            total += r
            rows.append([ins.iid, ins.text[:55] + "…", steps, f"{r:.2f}", " → ".join(actions[:4]) + (" …" if len(actions) > 4 else "")])
    except Exception as e:  # noqa: BLE001
        if not args.toy and type(e).__name__ in ("AuthenticationError", "APIConnectionError"):
            print("凭证无效 / 连不上 API：设置 ANTHROPIC_API_KEY 或运行 ant auth login，并检查网络")
            return
        raise
    print_header(f"模型 {name}：{len(instrs)} 条购物指令")
    print_table(["指令", "内容", "步数", "奖励", "前几个动作"], rows, aligns=["<", "<", ">", ">", "<"])
    print(f"\n平均奖励（分数）：{total / len(instrs) * 100:.1f}")
    print_observations([
        "它有没有读懂换了说法的属性（water resistant / waterproof）？有没有打开几件商品比较？",
        "它记得选选项吗？注意价格上限了吗？没拿满分的那几条，差在类型、属性、选项还是价格？",
        "6 条指令太少，下不了统计结论；对照代码 01 的规则式智能体，看看它的行为和哪一个最像。",
    ])


if __name__ == "__main__":
    main()