"""
=============================================================================
CMU 17-630 · 前置讲 P2 共享模块：llm_api.py —— 全课统一的 LLM 调用接口
=============================================================================

【这个模块是做什么的】
    一次 LLM 调用，从工程上看就是：把一个**请求**（消息列表 + 参数）交给一个**服务**，
    拿回一个**响应**（文字 + 停止原因 + 用量）。这个模块把它写成一个统一的接口，
    后面各讲的实验都通过它调用模型：

        llm = SimLLM()                                   # 或 TinyGPTLLM()、ClaudeLLM()
        resp = llm.complete(LLMRequest(messages=[{"role": "user", "content": "What is the capital of volel?"}]))
        resp.text, resp.stop_reason, resp.usage.input_tokens

    同一个接口背后可以换不同的"后端"：
        SimLLM      模拟的 chat 模型：回答第 1 讲虚构世界的问题；知识覆盖率、JSON 格式错误率、
                    "客套话"概率、温度带来的随机性……都是显式的旋钮。不联网、不花钱、可复现。
        TinyGPTLLM  第 1 讲训练的迷你 GPT（一个 base 模型），套一个"chat 模板"假装成 chat 模型。
        ClaudeLLM   真实的 Claude（可选；需要 anthropic SDK 和凭证；会花钱）。
        HttpLLM     （在 mock_server.py 里）通过 HTTP 调用本机的模拟服务器——看清楚网络上传的是什么。

    以及可以层层套上去的"包装器"（每一层只做一件事，组合起来就是一个可靠的调用栈）：
        FlakyLLM    故意注入失败（429 限速、500 服务器错误、超时），用来测试下面几层
        RetryingLLM 失败时按策略重试（指数退避 + 随机抖动，尊重 retry-after）
        CachedLLM   完全相同的请求直接返回上次的结果（省钱、可复现）
        MeteredLLM  记账（按 token 计费）+ 预算上限
        LoggedLLM   把每一次调用写进 JSONL 日志，事后可以只凭日志重算实验结果

【时间】所有"等待"都走一个可替换的时钟（VirtualClock）：演示时不真的睡觉，几秒钟就能模拟
    几分钟的重试过程，而且结果可复现。真实系统里把它换成 time.sleep / time.monotonic 即可。

【运行】python3 llm_api.py     # 自检：用 SimLLM 跑几个请求，套上全部包装器
【依赖】只用 Python 标准库；token 计数用前置讲 P1 的 BPE 分词器（../prereq1/text_lab.py），
    SimLLM 的知识来自第 1 讲的虚构世界（../lecture01/pretrain_world.py）。
    TinyGPTLLM 需要 PyTorch；ClaudeLLM 需要 anthropic SDK——两者都是用到时才导入。
"""

from __future__ import annotations

import hashlib
import json
import math
import os
import random
import re
import sys
from dataclasses import asdict, dataclass, field
from typing import Optional

HERE = os.path.dirname(os.path.abspath(__file__))
for _d in ("prereq1", "lecture01"):
    _p = os.path.join(HERE, "..", _d)
    if _p not in sys.path:
        sys.path.insert(0, _p)

from pretrain_world import RELATIONS, build_world  # noqa: E402
from text_lab import load_or_train_bpe  # noqa: E402


# ---------------------------------------------------------------------------
# 1. 请求、响应、错误
# ---------------------------------------------------------------------------

@dataclass
class LLMRequest:
    """
    一次调用的全部输入。

    【字段】
        messages   ：消息列表 [{"role": "user" / "assistant", "content": 文字}, ...]，必须从 user 开始、两种角色交替
        system     ：系统提示（单独的字段，不放在 messages 里——和 Anthropic Messages API 的约定一致）
        model      ：模型名
        max_tokens ：这次回答最多生成多少 token；写满了就被截断（stop_reason = "max_tokens"）
        temperature：采样温度（0 = 每次取最可能的 token）。注意：有些真实模型不接受这个参数（见 ClaudeLLM）
        stop       ：停止序列：生成的文字一旦出现其中之一，就在它之前停下（stop_reason = "stop_sequence"）
        seed       ：采样用的随机种子；支持它的后端在 seed 相同时给出相同的采样结果
        metadata   ：**不发给模型**的附加信息（例如这次调用属于哪个提示模板、第几道题），用于记账和日志
    【Python 小知识】@dataclass 自动生成 __init__；field(default_factory=dict) 让每个请求有自己的空字典
        （如果写成 metadata: dict = {}，所有请求会共享同一个字典——经典的 Python 陷阱）。
    """
    messages: list
    system: str = ""
    model: str = "sim-small"
    max_tokens: int = 256
    temperature: float = 0.0
    stop: tuple = ()
    seed: Optional[int] = None
    metadata: dict = field(default_factory=dict)

    def payload(self):
        """真正"发给模型服务"的那部分（不含 metadata）。"""
        return {"model": self.model, "system": self.system, "messages": self.messages,
                "max_tokens": self.max_tokens, "temperature": self.temperature,
                "stop": list(self.stop), "seed": self.seed}

    def key(self):
        """
        请求的指纹：payload 的规范化 JSON 的 sha256 前 16 位。两个请求的指纹相同 ⟺ 发给模型的内容完全相同。

        【为什么要 sort_keys=True】字典的键顺序不同，JSON 字符串就不同，指纹也就不同；
            排序之后，"内容相同"才等价于"字符串相同"。缓存命中率低的一个常见原因就是没排序。
        """
        blob = json.dumps(self.payload(), sort_keys=True, ensure_ascii=False)
        return hashlib.sha256(blob.encode("utf-8")).hexdigest()[:16]


@dataclass
class Usage:
    """用量：输入 token 数、输出 token 数（API 按这两个数计费）。"""
    input_tokens: int
    output_tokens: int


@dataclass
class LLMResponse:
    """
    一次调用的结果。

    【字段】text 回答文字；stop_reason 为什么停下（"end_turn" 自然结束 / "max_tokens" 被截断 /
        "stop_sequence" 碰到停止序列 / "refusal" 拒答）；usage 用量；model 实际回答的模型；
        latency_s 耗时（秒；SimLLM 给的是模拟值）；cached 是否来自缓存；attempts 一共尝试了几次（含重试）
    """
    text: str
    stop_reason: str
    usage: Usage
    model: str
    latency_s: float = 0.0
    cached: bool = False
    attempts: int = 1

    def to_dict(self):
        """转成可以写进 JSON 的字典。【Python 小知识】dataclasses.asdict 递归地把 dataclass 变成 dict。"""
        return asdict(self)


class LLMError(Exception):
    """所有 LLM 调用错误的基类。retryable 表示"原样重试有可能成功"。"""
    retryable = False
    status = None


class BadRequestError(LLMError):
    """请求本身有问题（400）：格式错、参数不合法。重试没有用，要改请求。"""
    status = 400


class AuthError(LLMError):
    """认证失败（401）：没带 API key 或 key 不对。重试没有用。"""
    status = 401


class NotFoundError(LLMError):
    """要找的东西不存在（404）：最常见的是模型名写错了。重试没有用。"""
    status = 404


class RateLimitError(LLMError):
    """
    被限速（429）：短时间内请求太多。等一会儿重试通常就好了。

    【属性】retry_after：服务端建议等待的秒数（来自 HTTP 头 retry-after；没有时为 None）
    """
    retryable = True
    status = 429

    def __init__(self, msg, retry_after=None):
        super().__init__(msg)
        self.retry_after = retry_after


class ServerError(LLMError):
    """服务端出错或过载（500 / 503 / 529……）。可以重试。"""
    retryable = True
    status = 500


class LLMTimeout(LLMError):
    """等太久没有回应（客户端超时）。可以重试，但要注意：服务端可能其实已经处理了这个请求。"""
    retryable = True


class BudgetExceeded(LLMError):
    """再调用一次就会超出预算。不会自动重试——这是故意停下来。"""


def validate_messages(messages):
    """
    检查消息列表是否合法：非空、从 user 开始、user / assistant 交替、内容是非空字符串。

    【边界】不合法时抛 BadRequestError（和真实 API 返回 400 一致：这种错误重试没有用）。
    【为什么要求交替】chat 模型是在"对话记录"这种文档上训练的（第 1 讲 §9.4），
        两条连续的 user 消息、或者以 assistant 开头，都不是它熟悉的格式。真实 API 对此的规定各家不同，
        这里取最严格的版本。
    """
    if not messages:
        raise BadRequestError("messages 不能为空")
    for i, m in enumerate(messages):
        if m.get("role") not in ("user", "assistant"):
            raise BadRequestError(f"messages[{i}] 的 role 必须是 user 或 assistant（系统提示请放在 system 字段）")
        if not isinstance(m.get("content"), str) or not m["content"].strip():
            raise BadRequestError(f"messages[{i}] 的 content 必须是非空字符串")
        expected = "user" if i % 2 == 0 else "assistant"
        if m["role"] != expected:
            raise BadRequestError(f"messages[{i}] 应该是 {expected}（必须从 user 开始、两种角色交替）")


# ---------------------------------------------------------------------------
# 2. token 计数与 chat 模板
# ---------------------------------------------------------------------------

_TOKENIZER = None


def tokenizer():
    """全课统一的"计数用"分词器：前置讲 P1 的字节级 BPE（英文 90%，1600 次合并），第一次调用时加载。"""
    global _TOKENIZER
    if _TOKENIZER is None:
        _TOKENIZER = load_or_train_bpe()
    return _TOKENIZER


def count_tokens(text):
    """一段文字的 token 数（用 P1 的分词器；真实模型的数字会不同，但量级和趋势一致）。"""
    return len(tokenizer().encode(text)) if text else 0


def render_chat(system, messages, add_generation_prompt=True):
    """
    把 system + messages 渲染成模型真正读到的那段文字（一种**示意性**的 chat 模板）。

    【返回】字符串，例如
        <|system|>Answer in one word.<|end|>
        <|user|>Capital of volel?<|end|>
        <|assistant|>
    【思路】chat 模型的"消息列表"在模型内部仍然是一段文档（第 1 讲 §9.4）：每条消息被特殊标记包起来，
        最后留一个 "<|assistant|>"，模型从那里开始续写。真实模型的特殊标记各不相同，这里只是示意。
    【提示工程联系】这解释了为什么"消息的条数"也要花 token（每条消息都有包装标记），
        也解释了为什么多轮对话越来越贵：每次调用都要把整段对话记录重新发一遍。
    """
    parts = []
    if system:
        parts.append(f"<|system|>{system}<|end|>")
    for m in messages:
        parts.append(f"<|{m['role']}|>{m['content']}<|end|>")
    if add_generation_prompt:
        parts.append("<|assistant|>")
    return "\n".join(parts)


def count_request_tokens(req):
    """一个请求的输入 token 数 = 渲染后的整段 chat 模板的 token 数。"""
    return count_tokens(render_chat(req.system, req.messages))


def _hash_float(*parts):
    """
    把若干个值确定性地映射到 [0, 1) 之间的一个数（同样的输入永远得到同样的数）。

    【为什么不用 Python 内置的 hash()】hash() 对字符串的结果每次启动 Python 都不一样（为了安全，
        有一个随机的"盐"，见 PYTHONHASHSEED），会让实验不可复现。sha256 每次都一样。
    """
    h = hashlib.sha256("|".join(str(p) for p in parts).encode("utf-8")).hexdigest()
    return int(h[:12], 16) / 16 ** 12


# ---------------------------------------------------------------------------
# 3. 后端
# ---------------------------------------------------------------------------

class LLM:
    """
    所有后端和包装器的共同父类。子类实现 complete(req) → LLMResponse。

    【用法】llm.complete(LLMRequest(...))，或者简写 llm.chat("问题", system="...", **参数)
    """
    name = "llm"

    def complete(self, req):
        raise NotImplementedError

    def chat(self, user_text, system="", **kwargs):
        """单轮对话的简写：llm.chat("What is the capital of volel?", max_tokens=50)。"""
        model = kwargs.pop("model", getattr(self, "name", "llm"))
        return self.complete(LLMRequest(messages=[{"role": "user", "content": user_text}],
                                        system=system, model=model, **kwargs))


# 识别问题的正则：从用户的话里找出"问的是哪个国家的哪种知识"
_QUESTION_PATTERNS = [
    ("capital", re.compile(r"capital of (\w+)", re.I)),
    ("language", re.compile(r"language .*?in (\w+)", re.I)),
    ("export", re.compile(r"(?:what is |is )?(\w+) known for", re.I)),
    ("leader", re.compile(r"who leads (\w+)", re.I)),
]


class SimLLM(LLM):
    """
    模拟的 chat 模型：回答第 1 讲虚构世界里 24 个国家的四种知识，也能"介绍"一个国家（长回答）。

    【旋钮】
        coverage       ：它"知道"多少比例的事实（其余的会**自信地答错**——编一个别国的答案，也就是幻觉）
        json_error_rate：要求输出 JSON 时，写出不合法 JSON（单引号）的概率
        chatty         ：在答案前面加一句客套话（"Sure! "）的概率；系统提示里有 "only" 时不加
        （指令里有 "one word" / "single word" / "answer word" 时只写答案本身；有 "JSON" 时写成 {"answer": ...}）
        flip_per_temp  ：温度每升高 1，"抽到一个错误答案"的概率增加多少（上限 0.6）
        latency        ：(基础延迟秒数, 每个输入 token 的秒数, 每秒输出多少 token)——只用来**报告**模拟的耗时，不真的等待
    【确定性】
        温度 0：同样的请求永远得到同样的回答。
        温度 > 0：给了 seed 时，"请求 + seed"决定结果（可复现）；没给 seed 时，用一个内部计数器，
        每次调用都不一样——和真实的采样一样。
    【这个模拟器能 / 不能说明什么】它能演示工程层面的现象（截断、停止序列、JSON 解析、重试、缓存、记账），
        **不能**说明真实模型的能力。后面几讲需要模拟"提示怎样影响模型"时，会用专门设计的模拟器。
    """

    def __init__(self, name="sim-small", coverage=0.9, json_error_rate=0.05, chatty=0.3,
                 flip_per_temp=0.25, latency=(0.25, 0.0002, 50.0)):
        self.name = name
        self.coverage, self.json_error_rate, self.chatty = coverage, json_error_rate, chatty
        self.flip_per_temp, self.latency = flip_per_temp, latency
        self.world = build_world()
        self.by_name = {c.name: c for c in self.world.countries}
        self._counter = 0

    # ---------------------------------------------------------------- 理解问题
    def parse(self, text):
        """从用户的话里找出 (国家, 关系)；找不到返回 (None, None)。"介绍 X"返回 (X, "about")。"""
        m = re.search(r"(?:tell me about|describe|explain) (\w+)", text, re.I)
        if m and m.group(1).lower() in self.by_name:
            return m.group(1).lower(), "about"
        for rel, pat in _QUESTION_PATTERNS:
            m = pat.search(text)
            if m and m.group(1).lower() in self.by_name:
                return m.group(1).lower(), rel
        return None, None

    def knows(self, country, relation):
        """这个模型"知不知道"某个事实：由 (模型名, 国家, 关系) 的哈希决定，所以每次都一样。"""
        return _hash_float(self.name, country, relation) < self.coverage

    def believed_answer(self, country, relation):
        """
        模型"心里认为"的答案：知道就是正确答案；不知道就从别的国家那里借一个同类的答案（自信的幻觉）。
        """
        c = self.by_name[country]
        if self.knows(country, relation):
            return c.attr(relation)
        others = [x for x in self.world.countries if x.name != country]
        return others[int(_hash_float("halluc", self.name, country, relation) * len(others))].attr(relation)

    # ---------------------------------------------------------------- 写回答
    def _about(self, country):
        """一段关于某国的介绍（长回答，用来演示 max_tokens 截断和流式输出）。"""
        c = self.by_name[country]
        return (f"{country.capitalize()} is a country in the {c.region}. Its capital is {c.capital}, "
                f"and most people there speak {c.language}. The country is best known for its {c.export}, "
                f"which it sells to its neighbours. It is led by {c.leader}. Many visitors come in summer, "
                f"and the old roads around {c.capital} are worth a walk.")

    def _draw(self, req):
        """这次调用的随机数生成器：温度 0 不需要随机；有 seed 用 seed；没有 seed 用内部计数器。"""
        if req.seed is not None:
            return random.Random(f"{req.key()}|{req.seed}")
        self._counter += 1
        return random.Random(f"{req.key()}|call{self._counter}")

    def generate_text(self, req):
        """不考虑截断和停止序列时，模型"想写"的完整回答。"""
        user = req.messages[-1]["content"]
        instructions = (req.system + " " + user).lower()
        wants_json = "json" in instructions
        country, rel = self.parse(user)
        rng = self._draw(req)
        if country is None:
            return "I can answer questions about the countries of this world, for example their capitals."
        if rel == "about":
            return self._about(country)
        answer = self.believed_answer(country, rel)
        if req.temperature > 0:
            p_flip = min(0.6, self.flip_per_temp * req.temperature)
            if rng.random() < p_flip:
                pool = sorted({getattr(c, rel) for c in self.world.countries} - {answer})
                answer = rng.choice(pool)
        terse = any(k in instructions for k in ("answer word", "one word", "single word"))
        if wants_json:
            if _hash_float("jsonerr", self.name, req.key()) < self.json_error_rate:
                body = f"{{'answer': '{answer}'}}"                       # 单引号：不合法的 JSON
            else:
                body = json.dumps({"answer": answer})
        elif terse:
            body = answer                                                 # 听从"只答一个词"的指令
        else:
            body = f"The answer is {answer}."
        polite = "only" not in instructions and _hash_float("chatty", self.name, req.key()) < self.chatty
        return ("Sure! " + body) if polite else body

    def complete(self, req):
        """
        处理一个请求：生成回答 → 应用停止序列 → 应用 max_tokens 截断 → 记用量和模拟耗时。

        【边界】
            - 消息不合法 → BadRequestError；max_tokens < 1 → BadRequestError
            - 停止序列：在第一次出现的位置之前截断，停止序列本身不包含在输出里（和多数 API 一致）
            - max_tokens 截断发生在 token 边界上；如果切在一个汉字的字节中间，不完整的字节被丢掉
        """
        validate_messages(req.messages)
        if req.max_tokens < 1:
            raise BadRequestError("max_tokens 必须至少是 1")
        text = self.generate_text(req)
        stop_reason = "end_turn"
        cut = [text.find(s) for s in req.stop if s and s in text]
        if cut:
            text, stop_reason = text[:min(cut)], "stop_sequence"
        ids = tokenizer().encode(text)
        if len(ids) > req.max_tokens:
            text = tokenizer().decode(ids[:req.max_tokens], errors="ignore")
            ids, stop_reason = ids[:req.max_tokens], "max_tokens"
        n_in = count_request_tokens(req)
        base, per_in, out_rate = self.latency
        return LLMResponse(text=text, stop_reason=stop_reason, usage=Usage(n_in, len(ids)), model=self.name,
                           latency_s=base + per_in * n_in + len(ids) / out_rate)


class TinyGPTLLM(LLM):
    """
    第 1 讲的迷你 GPT（base 模型），套一个 chat 模板，让它也能接受"消息列表"。

    【思路】base 模型不会"对话"，只会续写文档。这里把最后一条用户消息改写成它熟悉的问答页格式
        "<doc> faq : q : <问题> a :"，再让它续写，遇到 "." 停止。这就是最原始的"把 chat 接口套在
        base 模型上"的做法——chat 模板是**人为约定的文档格式**。
    【边界】只能回答第 1 讲世界里那几种问题；问题里有它词表外的词时会变成 <unk>；
        system 提示被忽略（这个 base 模型从没见过"系统提示"这种东西）。
        用量按第 1 讲的词级分词器计数。需要 PyTorch，第一次使用时如果没有缓存会先训练约 50 秒。
    """

    def __init__(self, name="tiny-gpt-base"):
        self.name = name
        from tiny_gpt import get_model                       # 延迟导入：没装 PyTorch 也能用本模块的其他部分
        import torch
        torch.set_num_threads(1)
        self.model, self.tok, _ = get_model("base", verbose=False)

    def prompt_for(self, req):
        """把消息列表渲染成迷你 GPT 熟悉的问答页格式（只看最后一条用户消息）。"""
        q = req.messages[-1]["content"].lower()
        q = re.sub(r"([?.,!])", r" \1", q)
        q = re.sub(r"\s+", " ", q).strip()
        if not q.endswith("?"):
            q += " ?"
        return f"<doc> faq : q : {q} a :"

    def complete(self, req):
        from tiny_gpt import generate
        validate_messages(req.messages)
        prompt = self.prompt_for(req)
        out = generate(self.model, self.tok, prompt, max_new_tokens=min(req.max_tokens, 20),
                       temperature=req.temperature, stop=(".", "<doc>"), seed=req.seed or 0)
        n_out = len(out.split())
        stop_reason = "end_turn" if out.endswith((".", "<doc>")) else "max_tokens"
        text = out.rstrip(" .").replace("<doc>", "").strip()
        return LLMResponse(text=text, stop_reason=stop_reason, usage=Usage(len(prompt.split()), n_out),
                           model=self.name)


class ClaudeLLM(LLM):
    """
    真实的 Claude（可选）。用官方 anthropic SDK 把 LLMRequest 翻译成 Messages API 的请求。

    【需要】pip install anthropic；凭证：环境变量 ANTHROPIC_API_KEY，或用 `ant auth login` 登录
        （SDK 会自动读取它保存的凭证）。**不要把 key 写进代码**。会花钱。
    【和我们的接口的几处差别（都是真实 API 的特点）】
        1. temperature 和 seed：默认模型 claude-opus-5-5 **不接受**采样参数（传了会返回 400），
           也没有 seed。所以这里不发送它们，第一次遇到时打印一条提醒——实验无法像 SimLLM 那样精确复现。
        2. 思考：这个模型总是先在内部"思考"（adaptive thinking），用 effort 控制深度
           （low / medium / high / xhigh / max；它的默认是 medium）。简单任务用 low 更快更省。
           思考的内容默认不返回，我们只取 text 块；但思考的 token **计入输出用量**。
        3. 拒答：安全分类器可能拒绝请求（stop_reason == "refusal"）。这里默认打开服务端回退
           （fallbacks="default"）：被拒时 API 自动换一个推荐的模型重试；整条链都拒绝时，
           stop_reason 仍是 "refusal"，由调用方处理。
        4. 重试：SDK 自己默认会重试 2 次（429、5xx、连接错误）。如果外面再套一层 RetryingLLM，
           请把 sdk_max_retries 设成 0，否则重试次数会相乘。
    """

    def __init__(self, model="claude-opus-5-5", effort="low", sdk_max_retries=2):
        try:
            import anthropic
        except ImportError as e:
            raise LLMError("没有安装 anthropic SDK：运行 pip install anthropic") from e
        self._anthropic = anthropic
        self.client = anthropic.Anthropic(max_retries=sdk_max_retries)
        self.name, self.effort = model, effort
        self._warned = False

    def complete(self, req):
        anthropic = self._anthropic
        validate_messages(req.messages)
        if (req.temperature or req.seed is not None) and not self._warned:
            print(f"  [ClaudeLLM] 提醒：{self.name} 不接受 temperature / seed，已忽略（结果每次可能不同）")
            self._warned = True
        kwargs = {"system": req.system} if req.system else {}
        if req.stop:
            kwargs["stop_sequences"] = list(req.stop)
        try:
            r = self.client.beta.messages.create(
                model=self.name, max_tokens=req.max_tokens, messages=req.messages,
                output_config={"effort": self.effort},
                betas=["server-side-fallback-2026-07-01"], fallbacks="default", **kwargs)
        except anthropic.RateLimitError as e:
            ra = e.response.headers.get("retry-after")
            raise RateLimitError(str(e), float(ra) if ra else None) from e
        except anthropic.AuthenticationError as e:
            raise AuthError(str(e)) from e
        except anthropic.BadRequestError as e:
            raise BadRequestError(str(e)) from e
        except anthropic.APITimeoutError as e:
            raise LLMTimeout(str(e)) from e
        except anthropic.APIStatusError as e:
            if e.status_code >= 500:
                raise ServerError(str(e)) from e
            raise LLMError(str(e)) from e
        except anthropic.APIConnectionError as e:
            raise ServerError(f"连接失败：{e}") from e
        text = "".join(b.text for b in r.content if b.type == "text").strip()
        return LLMResponse(text=text, stop_reason=r.stop_reason, model=r.model,
                           usage=Usage(r.usage.input_tokens, r.usage.output_tokens))


# ---------------------------------------------------------------------------
# 4. 时钟
# ---------------------------------------------------------------------------

class VirtualClock:
    """
    一个"假"的时钟：sleep(dt) 不真的等待，只把时间往前拨 dt 秒。

    【为什么】重试的等待动辄几秒到几十秒；演示时用假时钟，几毫秒就能模拟完，而且每次结果一样。
        真实系统里换成 time.monotonic（计时）和 time.sleep（等待）。
    """

    def __init__(self):
        self.t = 0.0

    def now(self):
        return self.t

    def sleep(self, dt):
        self.t += max(0.0, dt)


class RealClock:
    """真实的时钟：now() 是单调递增的秒数（time.monotonic），sleep() 真的等待。和 VirtualClock 接口相同。"""

    def now(self):
        import time
        return time.monotonic()

    def sleep(self, dt):
        import time
        time.sleep(max(0.0, dt))


# ---------------------------------------------------------------------------
# 5. 包装器
# ---------------------------------------------------------------------------

class FlakyLLM(LLM):
    """
    在一个后端外面注入失败：按概率抛出 429（附 retry_after）、500，或者超时。

    【参数】p_429 / p_500 / p_timeout：三种失败各自的概率；retry_after：429 时服务端建议等待的秒数；
        timeout_s：超时要等多久才报错（客户端白白等掉的时间，会计入时钟）；seed：随机种子；clock：时钟
    【用途】测试重试逻辑。真实服务的失败率通常远低于这里的设定，但高峰期的 429 并不罕见。
    """

    def __init__(self, inner, p_429=0.2, p_500=0.05, p_timeout=0.0, retry_after=2.0,
                 timeout_s=30.0, seed=0, clock=None):
        self.inner, self.name = inner, inner.name
        self.p_429, self.p_500, self.p_timeout = p_429, p_500, p_timeout
        self.retry_after, self.timeout_s = retry_after, timeout_s
        self.rng, self.clock = random.Random(seed), clock or VirtualClock()
        self.calls = 0

    def complete(self, req):
        self.calls += 1
        r = self.rng.random()
        if r < self.p_429:
            raise RateLimitError("429 Too Many Requests", retry_after=self.retry_after)
        if r < self.p_429 + self.p_500:
            raise ServerError("500 Internal Server Error")
        if r < self.p_429 + self.p_500 + self.p_timeout:
            self.clock.sleep(self.timeout_s)
            raise LLMTimeout(f"{self.timeout_s:.0f} 秒没有回应")
        resp = self.inner.complete(req)
        self.clock.sleep(resp.latency_s)
        return resp


class RetryingLLM(LLM):
    """
    失败时自动重试。只重试 retryable 的错误（429、5xx、超时）；400、401 这类错误立刻抛出。

    【参数】
        max_retries：最多重试几次（总尝试次数 = 1 + max_retries）
        base / cap ：第 k 次重试前的等待 = min(cap, base · 2^k)（指数退避）
        jitter     ："none"（不抖动）或 "full"（在 [0, 上面的值] 里均匀随机取一个，即"全抖动"）
        honor_retry_after：服务端给了 retry-after 时，至少等这么久
        clock      ：时钟（VirtualClock 或真实时钟）；seed：抖动用的随机种子
    【为什么要指数退避 + 抖动】立刻重试只会让过载的服务更过载；指数退避让等待越来越长；
        而"全抖动"让很多同时失败的客户端不会在同一时刻一起重试（04 号脚本演示 3 的"惊群"）。
    【边界】重试用完仍失败 → 抛出最后一个错误；返回的 LLMResponse.attempts 记录了一共试了几次。
    【提示工程联系】一次实验要调用模型成百上千次，没有重试，一次偶然的 429 就会让整组结果缺一块；
        但重试会让"同一个请求被处理了两次"成为可能（例如超时其实是服务端慢了）——对于只读的生成请求无妨，
        对于有副作用的工具调用（发邮件、下订单）就要小心（第 8 讲）。
    """

    def __init__(self, inner, max_retries=4, base=0.5, cap=8.0, jitter="full", honor_retry_after=True,
                 clock=None, seed=0):
        self.inner, self.name = inner, inner.name
        self.max_retries, self.base, self.cap = max_retries, base, cap
        self.jitter, self.honor_retry_after = jitter, honor_retry_after
        self.clock, self.rng = clock or VirtualClock(), random.Random(seed)
        self.retries = 0
        self.waited = 0.0

    def backoff(self, k, err):
        """第 k 次重试（k 从 0 开始）之前要等多少秒。"""
        d = min(self.cap, self.base * 2 ** k)
        if self.jitter == "full":
            d = self.rng.uniform(0, d)
        ra = getattr(err, "retry_after", None)
        if self.honor_retry_after and ra:
            d = max(d, ra)
        return d

    def complete(self, req):
        for k in range(self.max_retries + 1):
            try:
                resp = self.inner.complete(req)
                resp.attempts = k + 1
                return resp
            except LLMError as e:
                if not e.retryable or k == self.max_retries:
                    raise
                d = self.backoff(k, e)
                self.clock.sleep(d)
                self.retries += 1
                self.waited += d


class CachedLLM(LLM):
    """
    精确匹配缓存：请求的指纹（LLMRequest.key()）相同，就直接返回上次的响应，不再调用模型。

    【参数】path：可选的 JSONL 缓存文件（每行一个 {"key", "response"}）——程序重启后缓存还在；
        cache_sampled：温度 > 0 且**没给 seed** 的请求要不要缓存（默认不缓存，理由见下）
    【为什么默认不缓存"没给 seed 的采样请求"】这类请求每次本来就应该得到不同的样本。缓存它们，
        等于把第一次抽到的样本"冻住"：你以为在做 5 次独立采样，其实是同一个答案看了 5 遍。
        给了 seed 的采样请求可以缓存——seed 本身就是指纹的一部分，"同一个 seed 同一个样本"正是你想要的。
    【边界】缓存命中时 LLMResponse.cached = True、不产生新的用量（MeteredLLM 不会给它记费用）。
    """

    def __init__(self, inner, path=None, cache_sampled=False):
        self.inner, self.name = inner, inner.name
        self.path, self.cache_sampled = path, cache_sampled
        self.store = {}
        self.hits = self.misses = self.bypass = 0
        if path and os.path.exists(path):
            with open(path, encoding="utf-8") as f:
                for line in f:
                    rec = json.loads(line)
                    self.store[rec["key"]] = rec["response"]

    def complete(self, req):
        if req.temperature > 0 and req.seed is None and not self.cache_sampled:
            self.bypass += 1
            return self.inner.complete(req)
        k = req.key()
        if k in self.store:
            self.hits += 1
            d = dict(self.store[k])
            d["usage"] = Usage(**d["usage"])
            d.update(cached=True, latency_s=0.0)
            return LLMResponse(**d)
        self.misses += 1
        resp = self.inner.complete(req)
        self.store[k] = resp.to_dict()
        if self.path:
            with open(self.path, "a", encoding="utf-8") as f:
                f.write(json.dumps({"key": k, "response": resp.to_dict()}, ensure_ascii=False) + "\n")
        return resp


# 每百万 token 的价格（美元）：(输入, 输出)。sim-* 是虚构的；Claude 的是编写时（2026-10）官方公布的价格。
PRICES = {
    "sim-small": (0.50, 1.50),
    "sim-large": (3.00, 15.00),
    "tiny-gpt-base": (0.0, 0.0),
    "claude-opus-5-5": (4.00, 20.00),
    "claude-sonnet-5-5": (2.00, 10.00),
    "claude-haiku-5-5": (0.10, 0.50),
}


def cost_usd(model, usage):
    """按价格表算一次调用的费用（美元）。价格表里没有的模型按 0 计，并在 MeteredLLM 里提醒。"""
    pin, pout = PRICES.get(model, (0.0, 0.0))
    return (usage.input_tokens * pin + usage.output_tokens * pout) / 1e6


class MeteredLLM(LLM):
    """
    记账 + 预算：每次调用都按 token 记费用，按 metadata["tag"] 分组汇总；调用之前先估算，超预算就拒绝。

    【参数】budget_usd：预算上限（None 表示不限）
    【估算方法】调用前的估算 = 实际输入 token × 输入价 + **max_tokens** × 输出价（按最坏情况估计输出）。
        这样宁可提前停下，也不会超支。
    【属性】spent（已花费）、calls、by_tag（{标签: [调用数, 输入 token, 输出 token, 费用]}）
    """

    def __init__(self, inner, budget_usd=None):
        self.inner, self.name = inner, inner.name
        self.budget = budget_usd
        self.spent, self.calls = 0.0, 0
        self.by_tag = {}

    def complete(self, req):
        if self.budget is not None:
            worst = cost_usd(req.model, Usage(count_request_tokens(req), req.max_tokens))
            if self.spent + worst > self.budget:
                raise BudgetExceeded(f"已花 {self.spent:.4f} 美元，这次最多再花 {worst:.4f}，会超过预算 {self.budget}")
        resp = self.inner.complete(req)
        c = 0.0 if resp.cached else cost_usd(resp.model, resp.usage)
        self.spent += c
        self.calls += 1
        t = self.by_tag.setdefault(req.metadata.get("tag", "-"), [0, 0, 0, 0.0])
        t[0] += 1
        if not resp.cached:
            t[1] += resp.usage.input_tokens
            t[2] += resp.usage.output_tokens
        t[3] += c
        return resp


class LoggedLLM(LLM):
    """
    把每一次调用（包括失败的）写进一个 JSONL 日志：请求、响应、用量、费用、是否命中缓存、错误。

    【参数】path：日志文件路径（追加写入）；run_id：这次实验的名字；clock：时间戳用的时钟
    【为什么要日志】一个提示实验的结论（"模板 B 比 A 好 4 个百分点"）应该能**只凭日志**重算出来：
        用了哪个模型、哪些参数、每道题模型到底说了什么。没有日志，过一周你自己都说不清楚当时是怎么跑的。
    【边界】日志里有完整的提示和回答——如果里面有敏感信息（用户数据、密钥），要先脱敏再写。
    """

    def __init__(self, inner, path, run_id="run", clock=None):
        self.inner, self.name = inner, inner.name
        self.path, self.run_id, self.clock = path, run_id, clock or VirtualClock()
        self.n = 0

    def complete(self, req):
        self.n += 1
        rec = {"run_id": self.run_id, "call": self.n, "t": round(self.clock.now(), 3),
               "request": req.payload(), "metadata": req.metadata}
        try:
            resp = self.inner.complete(req)
            rec.update(response=resp.to_dict(), cost_usd=0.0 if resp.cached else cost_usd(resp.model, resp.usage))
            return resp
        except LLMError as e:
            rec["error"] = f"{type(e).__name__}: {e}"
            raise
        finally:
            with open(self.path, "a", encoding="utf-8") as f:
                f.write(json.dumps(rec, ensure_ascii=False) + "\n")


def read_log(path):
    """读回 JSONL 日志：返回记录列表（每行一个字典）。"""
    with open(path, encoding="utf-8") as f:
        return [json.loads(line) for line in f if line.strip()]


if __name__ == "__main__":
    import tempfile

    llm = SimLLM()
    c0, c3 = llm.world.countries[0].name, llm.world.countries[3].name
    for q in [f"What is the capital of {c0}?", f"Who leads {c3}? Answer in JSON.", f"Tell me about {c0}", "Hello!"]:
        r = llm.chat(q, max_tokens=40)
        print(f"  {q!r:40s} → {r.text!r}  [{r.stop_reason}, in {r.usage.input_tokens} / out {r.usage.output_tokens}]")
    clock = VirtualClock()
    with tempfile.TemporaryDirectory() as d:
        stack = LoggedLLM(MeteredLLM(CachedLLM(RetryingLLM(FlakyLLM(SimLLM(), p_429=0.5, clock=clock), clock=clock))),
                          os.path.join(d, "log.jsonl"), clock=clock)
        for i in range(6):
            stack.chat("What is the capital of volel?", metadata={"tag": "demo"})
        metered = stack.inner
        print(f"  6 次相同的请求：缓存命中 {metered.inner.hits} 次；重试 {metered.inner.inner.retries} 次；"
              f"花费 {metered.spent:.6f} 美元；日志 {len(read_log(os.path.join(d, 'log.jsonl')))} 行；虚拟时间 {clock.now():.2f} 秒")