"""
=============================================================================
前置讲 P2 · 代码 03：流式输出 —— 先看到第一个字，和正确地拼起每一个字
=============================================================================

对应讲义：lectures/prereq2_llm_as_a_service.md §4

【这个脚本要回答的问题】
    1. 流式输出（streaming）和一次性返回比，快在哪里？（演示 1：首字时间 vs 总时间）
    2. 流式输出在网络上是什么格式？（演示 2：SSE 事件）
    3. 网络把字节切碎时，客户端怎样才能不把中文变成乱码？（演示 3）

【怎么读输出】
    演示 1 的时间是真实的墙钟时间（⏱ 标记），每次运行会有几毫秒的差别；其余输出每次都一样。

【运行】python3 03_streaming.py     （约 5 秒；会在本机起模拟服务器，不联网）
【依赖】只用 Python 标准库；import 同目录的 mock_server.py、llm_api.py，第 1 讲的 console.py。
"""

import json
import os
import socket
import sys

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "lecture01"))
from console import print_header, print_observations, print_table  # noqa: E402
from llm_api import LLMRequest  # noqa: E402
from mock_server import HttpLLM, MockLLMServer  # noqa: E402

KEY = "sk-test-stream"


def demo1_latency(country):
    """演示 1：首字时间 vs 总时间。"""
    print_header("演示 1：流式输出快在哪里——首字时间（TTFT）vs 总时间")
    delay = 0.02          # 每个 token 0.02 秒，约每秒 50 个 token（真实模型常见的量级）
    with MockLLMServer(api_key=KEY, stream_delay=delay) as srv:
        llm = HttpLLM(srv.url, api_key=KEY)
        req = LLMRequest(messages=[{"role": "user", "content": f"Tell me about {country}"}], max_tokens=200)
        r = llm.complete(req)
        full, ttft, total, _ = llm.stream(req, read_bytes=256)
    rows = [["一次性返回（非流式）", f"⏱ {r.latency_s:.2f} 秒", f"⏱ {r.latency_s:.2f} 秒", r.usage.output_tokens],
            ["流式输出", f"⏱ {ttft:.2f} 秒", f"⏱ {total:.2f} 秒", r.usage.output_tokens]]
    print(f"模拟服务器每生成一个 token 需要 {delay} 秒；问题：'Tell me about {country}'")
    print_table(["方式", "看到第一个字", "收到全部回答", "输出 token"], rows)
    print(f"  两种方式拿到的文字相同？ {full == r.text}")
    print_observations([
        "两种方式的**总时间**差不多（流式还略长一点：每个事件都要单独发送一次）：模型生成每个 token 的速度不会因为流式而变快。",
        "差别在**首字时间**（time to first token, TTFT）：流式输出几乎立刻就能显示第一个字，"
        "非流式要等全部写完才一次性返回。对聊天界面来说，这决定了用户觉得'快'还是'卡'。",
        "长回答还有一个更实际的理由：一次性返回时，HTTP 连接要一直空等到全部生成完，容易撞上超时；"
        "官方 SDK 建议 max_tokens 很大时一律用流式。",
    ])


def demo2_sse(country):
    """演示 2：SSE 的原始格式。"""
    print_header("演示 2：流式输出在网络上是什么——SSE 事件")
    with MockLLMServer(api_key=KEY) as srv:
        host, port = srv.url.replace("http://", "").split(":")
        body = json.dumps({"model": "sim-small", "max_tokens": 30, "stream": True,
                           "messages": [{"role": "user", "content": f"What is the capital of {country}?"}]}).encode()
        req = (f"POST /v1/messages HTTP/1.1\r\nhost: {host}\r\ncontent-type: application/json\r\n"
               f"x-api-key: {KEY}\r\ncontent-length: {len(body)}\r\n\r\n").encode() + body
        with socket.create_connection((host, int(port))) as s:
            s.sendall(req)
            raw = b""
            while True:
                c = s.recv(4096)
                if not c:
                    break
                raw += c
    _, _, stream = raw.decode("utf-8").partition("\r\n\r\n")
    events = [e for e in stream.split("\n\n") if e.strip()]
    shown = events[:4] + [None] + events[-3:]          # 只打印开头 4 个和结尾 3 个事件
    for e in shown:
        if e is None:
            print(f"    | ……（中间省略 {len(events) - 7} 个事件）……\n    |")
            continue
        for line in e.split("\n"):
            print("    | " + line)
        print("    |")
    deltas = [json.loads(e.split("data: ", 1)[1])["delta"]["text"] for e in events if "text_delta" in e]
    print(f"  {len(events)} 个事件，其中 {len(deltas)} 个 text_delta：{deltas}")
    print_observations([
        "SSE（server-sent events）是一种很简单的格式：每个事件是 'event: 名字' 和 'data: JSON' 两行，事件之间空一行。",
        "回答被拆成很多个 text_delta，每个通常是一个或几个 token 的文字；客户端把它们按顺序拼起来。"
        "最后的 message_delta 里才有 stop_reason 和用量——用流式时，也要等流结束再判断'是不是被截断了'。",
        "真实 API 的事件更多（例如思考块、工具调用的增量、ping），但结构是一样的。官方 SDK 会替你解析，"
        "还提供'等流结束后拿到完整消息'的帮助函数。",
    ])


def demo3_bytes():
    """演示 3：字节被切碎时的两种客户端。"""
    print_header("演示 3：网络把字节切碎时——按块解码 vs 先攒齐再解码")
    text = "流式输出也要能处理中文，比如“鲲鹏展翅九万里”和 emoji 😀。"
    rows = []
    with MockLLMServer(api_key=KEY, chunk_bytes=5) as srv:
        llm = HttpLLM(srv.url, api_key=KEY, name="echo")
        req = LLMRequest(messages=[{"role": "user", "content": text}], model="echo", max_tokens=100)
        for mode in ["naive", "buffered"]:
            full, _, _, screen = llm.stream(req, read_bytes=7, decode=mode)
            rows.append([mode, screen, screen == text, screen.count("�")])
    print(f"服务器把每个事件的字节切成 5 字节一块发送，客户端每次读 7 个字节。原文：{text!r}")
    print_table(["客户端的做法", "屏幕上显示的文字", "和原文相同", "'�' 的个数"], rows, aligns=["<", "<", "<", ">"])
    print_observations([
        "naive：每读到一块字节就立刻按 UTF-8 解码显示。一个汉字的 3 个字节、emoji 的 4 个字节被切在两块里，"
        "两边都解不出来，屏幕上就出现一串 '�'——这是流式聊天界面里真实出现过的 bug。",
        "buffered：先按字节攒，攒到一个完整的事件（空行 b'\\n\\n'）再解码。分隔符是 ASCII，不可能出现在多字节字符的中间，"
        "所以切出来的每一段都是完整的 UTF-8（前置讲 P1 §2.4）。",
        "用官方 SDK 时这些都替你处理好了；自己写客户端、或者在中间加一层代理转发流时，一定要注意这一点。",
    ])


def main():
    from llm_api import SimLLM
    country = SimLLM().world.countries[0].name
    demo1_latency(country)
    demo2_sse(country)
    demo3_bytes()


if __name__ == "__main__":
    main()