"""
=============================================================================
CS329Z · 前置讲 P3 共享模块：conc_lab.py —— 慢服务的模拟、并发、限速、缓存、账本、预算
=============================================================================

对应讲义：lectures/prereq3_concurrency_caching.md

【这个文件要回答的问题】
    一次智能体任务要调几十到几百次模型 / 工具。原课 HW1 的"每日摘要"要给一天约 280 封邮件分类——
    一封一封地等，一个任务就要几分钟；同样的调用重复付费，账单还会翻倍。这个文件把解决这些问题
    的那几层从零实现一遍：

    SlowLM          ：一个"慢且会限速、偶尔出错"的假 LLM 服务（线程安全、行为由旋钮决定）
    run_serial / run_threads / run_async：三种跑一批调用的方式
    expected_time() ：一个解析模型：n 个调用、w 个并发、每个 L 秒，理论上要多久
    TokenBucket     ：令牌桶限速器（线程安全；可以用虚拟时钟，便于复现）
    cache_key()     ：按"模型 + 采样设置 + 完整消息"算缓存键
    DiskCache       ：按内容寻址的磁盘缓存（JSON 文件，一个键一个文件）
    Ledger          ：成本账本（JSONL 追加写，按模型 / 按标签汇总）
    Budget          ：预算上限（超了就抛 BudgetExceeded，而不是继续烧钱）
    CachedLM        ：把上面几层叠起来——缓存 → 预算 → 限速 → 重试 → 真正调用 → 记账。
                      这就是原课 HW1 那个 LM 封装的结构，第 2 讲会在它上面接真实的模型。

【关于"确定性"】
    并发和计时天生不确定。本模块的做法是：**让需要复现的东西不依赖真实时间**
    （TokenBucket 可以用虚拟时钟；SlowLM 的回复由内容哈希决定），
    而真实的墙钟测量只在脚本里出现一两处，并明确标注"你的机器上的数字会不同"。

【运行】python3 conc_lab.py   # 自检：跑一小批调用，演示缓存、账本、预算
【依赖】只用 Python 标准库。
"""

import asyncio
import hashlib
import json
import os
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field

# 假装的价格：每百万输入 / 输出 token 多少美元。数量级参考真实的大 / 小模型（大模型贵十几倍），
# 但这些数字是编的，不代表任何具体产品。
PRICES = {
    "big-model":   (3.0, 15.0),
    "small-model": (0.2, 1.0),
}
DEFAULT_PRICE = (3.0, 15.0)


class BudgetExceeded(Exception):
    """预算用完了：在**发出调用之前**抛出，所以不会再花钱。"""


class RateLimitError(Exception):
    """被服务端限速（对应 HTTP 429）。"""


class TransientError(Exception):
    """临时故障（对应 HTTP 5xx）：值得重试。"""


# ============================================================================ 1. 一个慢服务
@dataclass
class SlowLM:
    """
    一个假的 LLM 服务：慢、可能限速、可能临时出错，但**行为完全由旋钮决定**。

    【字段】
        latency       ：每次调用要多少秒（真的 sleep，模拟等网络和模型生成）
        max_concurrent：服务端允许的最大并发；超过就抛 RateLimitError（None 表示不限）
        fail_every    ：每隔多少次调用故意失败一次（0 表示不失败）
        name          ：模型名（进缓存键和账本）
    【线程安全】计数器都用一把锁保护；`in_flight` 记录"此刻有多少个调用正在进行"。
    【为什么真的 sleep】因为这一讲要讲的正是"等待"：I/O 等待期间 Python 会释放 GIL，
        所以线程能真正并行地等（演示 01）。用假的时间就看不到这件事。
    """
    latency: float = 0.05
    max_concurrent: int = None
    fail_every: int = 0
    name: str = "slow-lm-1"
    calls: int = 0
    in_flight: int = 0
    peak_in_flight: int = 0
    rejected: int = 0
    failed: int = 0
    _lock: threading.Lock = field(default_factory=threading.Lock, repr=False)

    def call(self, messages, temperature=0.0, seed=0):
        """
        调用一次"模型"。

        【参数】messages：消息列表（和第 1 讲的 LM 约定一样）；temperature、seed：进缓存键
        【返回】一段文字（由内容哈希决定，所以同样的输入永远得到同样的输出）
        【边界】
            - 并发超过 max_concurrent → RateLimitError（计入 rejected），**立刻**返回不等待
            - 每 fail_every 次失败一次 → TransientError（计入 failed）
        """
        with self._lock:
            self.calls += 1
            n = self.calls
            if self.max_concurrent is not None and self.in_flight >= self.max_concurrent:
                self.rejected += 1
                raise RateLimitError(f"too many concurrent requests (limit {self.max_concurrent})")
            self.in_flight += 1
            self.peak_in_flight = max(self.peak_in_flight, self.in_flight)
        try:
            if self.fail_every and n % self.fail_every == 0:
                with self._lock:
                    self.failed += 1
                raise TransientError("upstream temporarily unavailable")
            time.sleep(self.latency)                 # 等待：这期间 GIL 被释放
            return self._reply(messages)
        finally:
            with self._lock:
                self.in_flight -= 1

    async def acall(self, messages, temperature=0.0, seed=0):
        """
        asyncio 版本：用 await asyncio.sleep 代替 time.sleep。

        【为什么要有两个版本】asyncio 的并发靠"遇到 await 就让出控制权"，而 time.sleep 不会让出，
            它会把整个事件循环卡住（演示 01 会看到这个经典错误）。
        """
        with self._lock:
            self.calls += 1
            n = self.calls
            self.in_flight += 1
            self.peak_in_flight = max(self.peak_in_flight, self.in_flight)
        try:
            if self.fail_every and n % self.fail_every == 0:
                with self._lock:
                    self.failed += 1
                raise TransientError("upstream temporarily unavailable")
            await asyncio.sleep(self.latency)
            return self._reply(messages)
        finally:
            with self._lock:
                self.in_flight -= 1

    def _reply(self, messages):
        """由内容哈希决定的"回复"：同样的输入永远一样（像温度为 0 的贪心解码）。"""
        blob = json.dumps(messages, sort_keys=True, ensure_ascii=False)
        h = hashlib.sha256(blob.encode()).hexdigest()[:6]
        last = messages[-1]["content"] if messages else ""
        return f"[{self.name}] answer to {last[:40]!r} is {h}"

    def stats(self):
        """【返回】dict：调用次数、峰值并发、被限速次数、失败次数"""
        return {"calls": self.calls, "peak_in_flight": self.peak_in_flight,
                "rejected": self.rejected, "failed": self.failed}

    def reset(self):
        """清零所有计数器（每个演示开始前调用，这样数字只反映这一个演示）。"""
        with self._lock:
            self.calls = self.in_flight = self.peak_in_flight = self.rejected = self.failed = 0


# ============================================================================ 2. 三种跑一批调用的方式
def run_serial(lm, items, make_messages):
    """
    一个一个地调（串行）。

    【参数】lm：有 .call 的对象；items：输入列表；make_messages：item → 消息列表
    【返回】(结果列表, 耗时秒)
    """
    t0 = time.monotonic()
    out = [lm.call(make_messages(it)) for it in items]
    return out, time.monotonic() - t0


def run_threads(lm, items, make_messages, workers=8):
    """
    用线程池并发地调。

    【思路】每个调用的时间几乎全花在**等**（等网络、等模型生成）。Python 的 GIL 只在执行字节码时
        被持有，等待 I/O（包括 time.sleep）时会释放，所以多个线程能真正同时等——这就是
        "I/O 密集型任务用线程有效"的原因。
    【参数】workers：同时最多几个线程
    【返回】(按输入顺序排好的结果列表, 耗时秒)
    【边界】executor.map 会保持输入顺序；任何一个调用抛异常，异常会在读结果时抛出。
    【Python 小知识】with ThreadPoolExecutor(...) as ex 会在退出时等所有任务结束。
    """
    t0 = time.monotonic()
    with ThreadPoolExecutor(max_workers=workers) as ex:
        out = list(ex.map(lambda it: lm.call(make_messages(it)), items))
    return out, time.monotonic() - t0


def run_async(lm, items, make_messages, limit=8, blocking=False):
    """
    用 asyncio 并发地调。

    【参数】limit：用一个信号量限制同时进行的调用数；blocking：True 时故意用阻塞的 time.sleep
        （演示"一个阻塞调用卡住整个事件循环"）
    【返回】(结果列表, 耗时秒)
    【Python 小知识】asyncio 是**单线程**的并发：一个事件循环轮流推进很多个协程，
        每个协程在 await 处让出控制权。所以它只适合 I/O 等待，不适合计算。
    """
    async def main():
        sem = asyncio.Semaphore(limit)

        async def one(it):
            async with sem:
                if blocking:
                    lm.call(make_messages(it))       # 故意的错误：阻塞整个事件循环
                    return "(blocking)"
                return await lm.acall(make_messages(it))

        return await asyncio.gather(*(one(it) for it in items))

    t0 = time.monotonic()
    out = asyncio.run(main())
    return out, time.monotonic() - t0


def expected_time(n, workers, latency):
    """
    解析模型：n 个调用、最多 w 个并发、每个 L 秒，理论上要多久。

    【公式】时间 ≈ ceil(n / w) × L（分成 ceil(n/w) "批"，每批并行跑 L 秒）
    【参数】n、workers、latency
    【返回】秒
    【边界】假设每个调用的时间一样、没有额外开销、服务端不限速。真实情况会更慢
        （线程创建、限速、长尾延迟），所以这是一个**下界**。
    """
    import math
    return math.ceil(n / workers) * latency


# ============================================================================ 3. 限速：令牌桶
class TokenBucket:
    """
    令牌桶限速器：桶里最多 capacity 个令牌，每秒补 rate 个；每次调用取一个令牌，没有就等。

    【为什么是这个算法】它允许**短时突发**（桶里攒着的令牌可以一次用掉），长期速率又不超过 rate。
        真实的 API 限速（每分钟多少请求 / 多少 token）基本都是这个模型。
    【参数】rate：每秒补多少个；capacity：桶的容量（None = 等于 rate）；
        clock / sleep：时间函数，默认用真实时间；传虚拟时钟就能得到可复现的结果
    【线程安全】用一把锁保护令牌数；等待在锁外面进行（否则会把别的线程一起卡住）。
    """

    EPS = 1e-9      # 比较令牌数时的容差（见 acquire 里的说明）

    def __init__(self, rate, capacity=None, clock=None, sleep=None):
        self.rate = float(rate)
        self.capacity = float(capacity if capacity is not None else rate)
        self.clock = clock or time.monotonic
        self.sleep = sleep or time.sleep
        self.tokens = self.capacity
        self.last = self.clock()
        self.waits = []                              # 每次等待了多久（便于演示）
        self._lock = threading.Lock()

    def acquire(self, n=1):
        """
        取 n 个令牌，不够就等到够。

        【返回】等了多少秒（0 表示没等）
        【边界】n > capacity 时会永远等不到——所以先检查并抛 ValueError。
        """
        if n > self.capacity:
            raise ValueError(f"cannot acquire {n} tokens from a bucket of capacity {self.capacity}")
        while True:
            with self._lock:
                now = self.clock()
                self.tokens = min(self.capacity, self.tokens + (now - self.last) * self.rate)
                self.last = now
                # EPS 很重要：浮点运算会让补到的令牌差那么一点点（例如 0.9999999999999998），
                # 于是还要再等 4e-17 秒——真实时钟会直接跨过去，而虚拟时钟加不上这么小的数，
                # 就会死循环。详见代码 README 的"关于实验结果"。
                if self.tokens >= n - self.EPS:
                    self.tokens = max(0.0, self.tokens - n)
                    return 0.0
                need = (n - self.tokens) / self.rate
            self.waits.append(need)
            self.sleep(need)                         # 在锁外面等


class VirtualClock:
    """
    一个虚拟时钟：时间只在"睡眠"时前进。

    【为什么】用它替换 TokenBucket 的 clock / sleep，限速的行为就完全可复现（而且脚本跑得飞快）。
    【用法】vc = VirtualClock(); b = TokenBucket(5, clock=vc.now, sleep=vc.sleep)
    【边界】只适合单线程的演示：多个线程共用一个虚拟时钟没有意义。
    """

    def __init__(self):
        self.t = 0.0

    def now(self):
        return self.t

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


# ============================================================================ 4. 缓存键
def cache_key(model, messages, temperature=0.0, seed=0, extra=None):
    """
    算一个调用的缓存键。

    【思路】键必须覆盖**所有会影响回复的东西**：模型名、完整的消息列表、采样设置。
        少放一样，就可能把 A 的回复当成 B 的（缓存污染）；多放无关的东西（时间戳、请求 id），
        就永远命中不了。
    【参数】model、messages、temperature、seed；extra：别的会影响结果的设置（dict）
    【返回】64 位十六进制字符串（sha256）
    【Python 小知识】json.dumps(..., sort_keys=True) 保证同样的内容得到同样的字符串
        （字典的键顺序不同也算同样的内容）。
    【工程联系】原课 HW1 的 LM 封装就是"按模型、采样设置和**完整的消息**做键"。
        它的 handout 还提醒：提示词里出现每次都变的东西（例如一个临时工作区的绝对路径）
        会让缓存**永远不命中**。
    """
    payload = {"model": model, "messages": messages, "temperature": temperature,
               "seed": seed, "extra": extra or {}}
    blob = json.dumps(payload, sort_keys=True, ensure_ascii=False)
    return hashlib.sha256(blob.encode("utf-8")).hexdigest()


# ============================================================================ 5. 缓存、账本、预算
class DiskCache:
    """
    一个最简单的按内容寻址的磁盘缓存：一个键一个 JSON 文件。

    【为什么用磁盘】进程重启、换个脚本跑，昨天的调用还能复用——开发智能体时这能省大量的钱和时间。
    【线程安全】用一把锁保护"读-改-写"；写文件时先写临时文件再 rename（原子替换），
        避免两个线程同时写出一个半截的文件。
    """

    def __init__(self, directory):
        self.dir = directory
        os.makedirs(self.dir, exist_ok=True)
        self.hits = 0
        self.misses = 0
        self._lock = threading.Lock()

    def path(self, key):
        """键 → 文件路径。"""
        return os.path.join(self.dir, key + ".json")

    def get(self, key):
        """【返回】缓存的回复，或者 None（未命中）。"""
        p = self.path(key)
        with self._lock:
            if not os.path.exists(p):
                self.misses += 1
                return None
            self.hits += 1
        with open(p, encoding="utf-8") as f:
            return json.load(f)["reply"]

    def put(self, key, reply, meta=None):
        """写入缓存。【边界】用临时文件 + os.replace 做原子替换。"""
        p = self.path(key)
        tmp = p + f".tmp{os.getpid()}-{threading.get_ident()}"
        with open(tmp, "w", encoding="utf-8") as f:
            json.dump({"reply": reply, "meta": meta or {}}, f, ensure_ascii=False)
        os.replace(tmp, p)

    def stats(self):
        """【返回】dict：命中、未命中、命中率"""
        total = self.hits + self.misses
        return {"hits": self.hits, "misses": self.misses,
                "hit_rate": (self.hits / total) if total else 0.0}


def count_tokens_approx(text):
    """
    粗略的 token 计数（和第 1 讲 llm.py 的做法一致：按词和标点切）。

    【返回】int。真实计费要按模型自己的分词器算，这里只用来比较数量级。
    """
    import re
    return len(re.findall(r"\w+|[^\w\s]", text))


def cost_usd(tokens_in, tokens_out, model=None):
    """
    按假设的价格算一次调用的费用（美元）。

    【参数】tokens_in、tokens_out；model：模型名（在 PRICES 里查价，查不到用 DEFAULT_PRICE）
    【返回】float
    【工程联系】真实的价目表还要区分缓存读取（更便宜）、批处理（半价）等档位；
        账本的价值不在于算得多准，而在于**按步骤分摊**，让你看出钱花在哪。
    """
    p_in, p_out = PRICES.get(model, DEFAULT_PRICE)
    return tokens_in / 1e6 * p_in + tokens_out / 1e6 * p_out


class Ledger:
    """
    成本账本：每次**真正发出**的调用追加一行 JSONL，可以按模型、按标签汇总。

    【为什么要标签】一次智能体任务里的调用来自不同的地方（给邮件分类、写摘要、当裁判）。
        打上标签才能回答"钱花在哪一步了"。原课 HW1 的测试套件就给每次调用打上测试名。
    【线程安全】写文件和累加都在锁里（JSONL 的一行一行追加天生适合并发写，但我们还要维护内存里的汇总）。
    """

    def __init__(self, path):
        self.path = path
        os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
        self.rows = []
        self._lock = threading.Lock()

    def record(self, model, tokens_in, tokens_out, tag="", cached=False):
        """
        记一次调用。

        【参数】cached：True 表示这次是缓存命中（**不计费**，但要记下来，才能算出省了多少）
        【返回】这次的费用（缓存命中时为 0.0）
        """
        usd = 0.0 if cached else cost_usd(tokens_in, tokens_out, model)
        row = {"model": model, "tokens_in": tokens_in, "tokens_out": tokens_out,
               "tag": tag, "cached": cached, "usd": usd}
        with self._lock:
            self.rows.append(row)
            with open(self.path, "a", encoding="utf-8") as f:
                f.write(json.dumps(row, ensure_ascii=False) + "\n")
        return usd

    def total_usd(self):
        """实际花掉的钱（不含缓存命中）。"""
        return sum(r["usd"] for r in self.rows)

    def saved_usd(self):
        """缓存**省下**的钱：命中的那些调用如果真发出去会花多少。"""
        return sum(cost_usd(r["tokens_in"], r["tokens_out"], r["model"]) for r in self.rows if r["cached"])

    def by(self, field_name):
        """
        按某个字段汇总。

        【参数】field_name："model" / "tag"
        【返回】dict：值 → {"calls", "cached", "tokens_in", "tokens_out", "usd"}
        """
        out = {}
        for r in self.rows:
            d = out.setdefault(r[field_name], {"calls": 0, "cached": 0, "tokens_in": 0,
                                               "tokens_out": 0, "usd": 0.0})
            d["calls"] += 1
            d["cached"] += 1 if r["cached"] else 0
            d["tokens_in"] += r["tokens_in"]
            d["tokens_out"] += r["tokens_out"]
            d["usd"] += r["usd"]
        return out


class Budget:
    """
    预算：花掉的钱超过上限时，在**下一次调用之前**抛 BudgetExceeded。

    【为什么要有】一个写错了的循环（模型一直复读、工具一直失败重试）能在几分钟里花掉很多钱。
        预算是最后一道保险。原课 HW1 的 LM 封装有两道：课程发的 key 本身有额度，
        外加一个你自己设的 CS329Z_BUDGET_USD。
    【边界】缓存命中不计费，所以预算用完之后，**缓存里已有的调用仍然可以继续**。
    """

    def __init__(self, limit_usd):
        self.limit = float(limit_usd)
        self.spent = 0.0
        self._lock = threading.Lock()

    def check(self):
        """【边界】超了就抛 BudgetExceeded（调用之前检查，所以不会再花钱）。"""
        with self._lock:
            if self.spent >= self.limit:
                raise BudgetExceeded(f"budget of {self.limit:.4f} USD reached (spent {self.spent:.4f})")

    def add(self, usd):
        """记一笔开销。"""
        with self._lock:
            self.spent += usd

    def remaining(self):
        """还剩多少（不会小于 0）。"""
        return max(0.0, self.limit - self.spent)


# ============================================================================ 6. 把几层叠起来
@dataclass
class CachedLM:
    """
    把几层叠成一个"消息进、字符串出"的 LM：缓存 → 预算 → 限速 → 重试 → 真正调用 → 记账。

    【顺序为什么是这样】
        1. **缓存最先**：命中就什么都不用做了，不花钱、不占限速、不用等。
        2. **预算其次**：确认还有钱再发请求。
        3. **限速第三**：主动按自己的速率排队，比被服务端拒绝再重试更有效率。
        4. **重试最后**（包着真正的调用）：只有真的发出去了才需要重试。
        5. **记账在成功之后**：失败的尝试不计费（真实 API 里失败通常也不计费，但要记次数）。
    【字段】inner：真正的服务（有 .call）；cache / ledger / budget / bucket：可选的各层；
        max_retries、base_delay：重试设置；tag：记进账本的标签；sleep：睡眠函数（测试时可替换）
    【工程联系】这就是原课 HW1 那个 LM 封装的结构（缓存、账本、预算），第 2 讲会在它上面接真实的模型。
    """
    inner: object
    cache: DiskCache = None
    ledger: Ledger = None
    budget: Budget = None
    bucket: TokenBucket = None
    max_retries: int = 2
    base_delay: float = 0.05
    tag: str = ""
    sleep: object = time.sleep
    attempts: int = 0

    def __call__(self, messages, temperature=0.0, seed=0):
        """
        【返回】模型的回复（字符串）
        【边界】
            - 缓存命中：直接返回，记一行 cached=True 的账
            - 预算用完：抛 BudgetExceeded（**在**调用之前）
            - RateLimitError / TransientError：退避重试，次数用完后把异常抛出去
        """
        key = cache_key(getattr(self.inner, "name", "lm"), messages, temperature, seed)
        tokens_in = sum(count_tokens_approx(m["content"]) for m in messages)
        if self.cache is not None:
            hit = self.cache.get(key)
            if hit is not None:
                if self.ledger is not None:
                    self.ledger.record(getattr(self.inner, "name", "lm"), tokens_in,
                                       count_tokens_approx(hit), self.tag, cached=True)
                return hit
        if self.budget is not None:
            self.budget.check()
        last = None
        for k in range(self.max_retries + 1):
            if self.bucket is not None:
                self.bucket.acquire()
            self.attempts += 1
            try:
                reply = self.inner.call(messages, temperature=temperature, seed=seed)
            except (RateLimitError, TransientError) as e:
                last = e
                if k == self.max_retries:
                    raise
                self.sleep(self.base_delay * (2 ** k))
                continue
            if self.cache is not None:
                self.cache.put(key, reply, {"model": getattr(self.inner, "name", "lm"), "tag": self.tag})
            if self.ledger is not None:
                usd = self.ledger.record(getattr(self.inner, "name", "lm"), tokens_in,
                                         count_tokens_approx(reply), self.tag, cached=False)
                if self.budget is not None:
                    self.budget.add(usd)
            return reply
        raise last                                   # 理论上到不了这里


if __name__ == "__main__":
    import tempfile

    lm = SlowLM(latency=0.02)
    items = [f"email {i}" for i in range(8)]
    mk = lambda it: [{"role": "user", "content": f"Classify: {it}"}]   # noqa: E731
    _, t_serial = run_serial(lm, items, mk)
    lm.reset()
    _, t_threads = run_threads(lm, items, mk, workers=8)
    print(f"8 个调用（每个 20 ms）的理论下界：串行 {expected_time(8, 1, 0.02):.2f} s，"
          f"线程池(8) {expected_time(8, 8, 0.02):.2f} s")
    print(f"实测的串行 / 并发耗时比至少 3 倍：{t_serial / t_threads > 3}（具体数字每台机器不同）")

    with tempfile.TemporaryDirectory() as d:
        cache = DiskCache(os.path.join(d, "cache"))
        ledger = Ledger(os.path.join(d, "ledger.jsonl"))
        budget = Budget(0.001)
        c = CachedLM(SlowLM(latency=0.0), cache=cache, ledger=ledger, budget=budget, tag="selftest")
        for _ in range(3):
            c([{"role": "user", "content": "same question"}])
        c([{"role": "user", "content": "another question"}])
        print("缓存：", cache.stats())
        print(f"账本：实际花费 {ledger.total_usd():.8f} 美元，缓存省下 {ledger.saved_usd():.8f} 美元")
        print("按标签：", ledger.by("tag"))
        try:
            Budget(0.0).check()
        except BudgetExceeded as e:
            print("预算用完 →", e)

    vc = VirtualClock()
    b = TokenBucket(rate=5, capacity=5, clock=vc.now, sleep=vc.sleep)
    for _ in range(8):
        b.acquire()
    print(f"令牌桶（5/s，容量 5）取 8 次：虚拟时钟走了 {vc.t:.2f} s，等待序列 {[round(w, 2) for w in b.waits]}")