"""
=============================================================================
CS349E · 第 15 讲 · 03_context_and_interface.py —— 工具接口与上下文：每一次观察都要进入上下文
=============================================================================

【这个脚本要回答的问题】
    智能体每调用一次工具，工具的输出就原样追加到上下文里，下一轮 LLM 要把整个上下文再读一遍（prefill）。
    工具接口返回什么、返回多长，直接决定了上下文多长、prefill 多少 token、花多少钱。
      演示 1：真实的工具输出有多长：编译器的完整错误输出 vs 只留错误行；错误元素的完整清单 vs 一行摘要。
      演示 2：一个 30 轮的会话：工具输出详细 / 精简 × 历史全保留 / 只留最近 5 次观察 × 思维链保留 / 只留摘要，
              上下文长度、prefill 的 token 数（有 / 没有前缀缓存）、GPU 时间和费用。
      演示 3：截断历史为什么会让前缀缓存失效，以及一个不破坏前缀的做法。

【关于数字】编译器输出是在你的机器上用系统的 C 编译器真实编译得到的（没有编译器时跳过）；token 数按"英文 4 个字符一个、
    汉字一个字一个"估计；会话里每轮的长度、单价是假设值（单价只用来比较相对大小）；GPU 时间按 Llama-3.1-70B 在 H100 上、
    50% MFU 的计算受限估计（第 1 讲的账本）。
【运行】python3 03_context_and_interface.py      （不到 1 秒）
【依赖】numpy；本目录的 agent_lab；../lecture01 的 console、llm_math。
"""

import os
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
sys.path.insert(0, os.path.join(HERE, "..", "lecture01"))

from agent_lab import (compile_log, est_tokens, make_problem, mismatch_dump, observe, prefill_gpu_seconds,  # noqa: E402
                       session_context, summarize_compile_log, tiled_gemm)
from llm_math import LLAMA31_70B  # noqa: E402
from console import print_header, print_observations, print_table  # noqa: E402

TURNS = 30
PROMPT = 4000                 # 系统提示 + 任务 + 参考实现（假设）
PRICE_IN, PRICE_CACHED, PRICE_OUT = 3.0, 0.3, 15.0     # 每百万 token 的单价（美元，假设；缓存命中按 0.1 倍）


def demo1_outputs():
    print_header("演示 1：真实的工具输出有多长")
    rows = []
    log = compile_log()
    if log is None:
        print("（没有找到 C 编译器，跳过编译器输出的部分）")
    else:
        summ = summarize_compile_log(log)
        rows.append(["编译：完整输出（-Wall -Wextra）", f"{log.count(chr(10))} 行", f"{est_tokens(log):,}"])
        rows.append(["编译：只留 error 行", "1 行", f"{est_tokens(summ):,}"])
        print("完整的编译器输出：\n" + log)
        print("压缩之后：" + summ + "\n")
    for bug, shape in (("store_overflow", "edge"), ("b_transposed", "nice")):
        A, B, ref = make_problem(shape)
        C = tiled_gemm(A, B, bug)
        dump = mismatch_dump(C, ref)
        summ = observe("test_verbose", shape, bug)[1]
        rows.append([f"测试（{bug}·{shape}）：每个错误元素一行", f"{dump.count(chr(10)) + 1:,} 行", f"{est_tokens(dump):,}"])
        rows.append([f"测试（{bug}·{shape}）：模式摘要", "1 行", f"{est_tokens(summ):,}"])
    print_table(["工具输出", "长度", "约多少 token"], rows)
    print_observations([
        "一段 20 行的 C 代码，clang 的完整输出就有几百个 token：每个诊断都回显源码、画出插入符号，还有和错误无关的 warning。"
        "真实的 CUDA 扩展编译（nvcc + ninja + PyTorch 的头文件）出错时，输出常常是这个的几十倍。压缩成\"第几行、什么错\"只要几十个 token。",
        "测试的输出差别更大：b_transposed 在 128×128 上几乎每个元素都错，逐个列出要十几万个 token——比很多模型的整个上下文窗口还长。"
        "一行摘要（错了多少、哪些行列、最大误差）只要几十个 token，而且信息更有用（02 脚本：错误的**模式**才是区分 bug 的关键）。",
        "这正是 SWE-agent 的设计原则：动作要简单、反馈要\"有信息量但简洁\"、要有护栏（例如编辑前先跑 linter）。"
        "他们的消融里，搜索结果给摘要（18.0%）比逐条翻看（12.0%）好得多，后者甚至不如不给搜索工具（15.7%）；"
        "文件查看器一次显示 100 行（18.0%）比 30 行（14.3%）和整个文件（12.7%）都好。",
    ])


def demo2_sessions():
    print_header(f"演示 2：一个 {TURNS} 轮的会话——上下文、prefill 与费用")
    configs = [
        ("思维链全留 + 详细输出 + 全部历史", 6000, 4000, "full"),
        ("思维链全留 + 精简输出 + 全部历史", 6000, 120, "full"),
        ("只留摘要 + 详细输出 + 全部历史", 1500, 4000, "full"),
        ("只留摘要 + 精简输出 + 全部历史", 1500, 120, "full"),
        ("只留摘要 + 详细输出 + 最近 5 次观察", 1500, 4000, "last_k"),
    ]
    rows = []
    res = {}
    for name, asst, tool, hist in configs:
        sc = session_context(TURNS, PROMPT, asst, tool, history=hist, keep_last=5)
        ctx = sc["context"]
        over128 = next((t + 1 for t, c in enumerate(ctx) if c > 128_000), None)
        no, ca = sum(sc["prefill_nocache"]), sum(sc["prefill_cache"])
        gpu_no = sum(prefill_gpu_seconds(n, 0, LLAMA31_70B) for n in sc["prefill_nocache"])
        gpu_ca = sum(prefill_gpu_seconds(n, c - n, LLAMA31_70B) for n, c in zip(sc["prefill_cache"], ctx))
        out_tok = TURNS * asst
        cost = (ca * PRICE_IN + (no - ca) * PRICE_CACHED + out_tok * PRICE_OUT) / 1e6
        res[name] = (ctx[-1], no, ca, gpu_no, gpu_ca, cost)
        rows.append([name, f"{ctx[-1]:,}", "—" if over128 is None else f"第 {over128} 轮", f"{no / 1e6:.2f} M", f"{ca / 1e6:.2f} M",
                     f"{gpu_no / 3600:.2f}", f"{gpu_ca / 3600:.3f}", f"{cost:.2f}"])
    print_table(["配置", "最后一轮的上下文", "超过 128k", "prefill（无缓存）", "prefill（前缀缓存）",
                 "GPU·小时（无缓存）", "GPU·小时（缓存）", "费用（美元，缓存）"], rows)
    print(f"假设：系统提示 + 任务 {PROMPT:,} token；每轮模型输出 6,000（带思维链）或 1,500（代码 + 摘要）token；工具输出 4,000（详细）或 120（精简）"
          f" token；单价：输入 {PRICE_IN} / 缓存命中 {PRICE_CACHED} / 输出 {PRICE_OUT} 美元每百万 token；GPU 时间按 Llama-3.1-70B、H100、50% MFU。")
    a = res[configs[0][0]]
    d = res[configs[3][0]]
    print_observations([
        f"全部历史、不用前缀缓存时，prefill 的总量随轮数**平方**增长：最\"啰嗦\"的配置 30 轮要 prefill {a[1] / 1e6:.1f} M 个 token，"
        f"最后一轮的上下文 {a[0]:,}；前缀缓存把它变成线性的（每轮只算新增的部分）：{a[2] / 1e6:.2f} M。",
        f"接口设计和上下文管理一起，把最后一轮的上下文从 {a[0]:,} 压到 {d[0]:,}，费用从 {a[5]:.2f} 美元降到 {d[5]:.2f} 美元。"
        "Kevin-32B 在多轮强化学习里也是这样做的：后面的轮次不带完整的思维链，只带模型自己写的简短摘要，否则上下文会爆炸。",
        "没有前缀缓存时，GPU 时间主要花在一遍遍重读同样的历史上——这是推理系统要为智能体负载做前缀缓存（第 7 讲）的直接原因。"
        "智能体的负载是**输入远多于输出**的：每轮输出几千 token，却要读几万到十几万 token 的上下文。",
    ])
    return res


def demo3_truncation():
    print_header("演示 3：截断历史 vs 前缀缓存")
    rows = []
    for name, hist, k in (("全部历史", "full", 5), ("只留最近 5 次观察（旧的换成占位符）", "last_k", 5),
                          ("只留最近 2 次观察", "last_k", 2)):
        sc = session_context(TURNS, PROMPT, 1500, 4000, history=hist, keep_last=k)
        hit = 1 - sum(sc["prefill_cache"]) / sum(sc["prefill_nocache"])
        rows.append([name, f"{sc['context'][-1]:,}", f"{max(sc['context']):,}", f"{sum(sc['prefill_cache']) / 1e6:.2f} M", f"{hit:.0%}"])
    print_table(["历史的管理", "最后一轮的上下文", "最长的上下文", "prefill（前缀缓存）", "缓存命中的比例"], rows)
    full = session_context(TURNS, PROMPT, 1500, 4000)
    l5 = session_context(TURNS, PROMPT, 1500, 4000, history="last_k", keep_last=5)
    print_observations([
        f"只留最近 5 次观察，上下文从 {full['context'][-1]:,} 降到 {l5['context'][-1]:,}，可要 prefill 的 token 反而从 "
        f"{sum(full['prefill_cache']) / 1e6:.2f} M 变成 {sum(l5['prefill_cache']) / 1e6:.2f} M：每一轮都要把\"刚刚变旧\"的那次观察换成占位符，"
        "从那个位置往后的内容都变了，前缀缓存只能用到它之前。上下文变短了，缓存却失效了。",
        "SWE-agent 的消融里，\"只留最近 5 次观察\"（18.0%）比全部历史（15.0%）解决的题更多——短的上下文让模型更专注。"
        "但从推理系统的角度，它把\"只追加\"的上下文变成了\"每轮都改中间\"的上下文。",
        "两全的做法：不在中间改，而是**分段**压缩——攒到一定长度，一次性把一大段旧历史换成摘要（之后又是只追加），"
        "或者只截断末尾附近的内容。上下文管理的策略要和推理引擎的前缀缓存一起设计。",
    ])


if __name__ == "__main__":
    demo1_outputs()
    demo2_sessions()
    demo3_truncation()