#!/usr/bin/env python3
"""replay.py - 把 /log 录下来的真实会话当成测试夹具回放. 

为什么需要它: ligent 的协议适配一直只在自建的 mock 服务端上验证过, 而 mock
是我们自己写的 —— 它接受什么, 完全取决于**我们以为**服务端会接受什么. 真实
服务端的差异是最大的未知数. 

/log 导出的是**原始字节** (请求体、响应头、未经解析的响应原文, 流式就是原始
SSE), 正好够原样伺候回去. 所以这里不重写一个"模拟客户端", 而是当**离线
服务端**: 把录下来的响应按顺序发给一个真的 ligent 实例, 于是 http.c 的分块
读取、llm.c 的 SSE 行切分与 tool_calls 增量重组, 全都在真实字节上跑. 

用法:
    replay.py --summary FILE                 这份录制里有什么
    replay.py --serve PORT FILE [--strict]   当离线服务端
    replay.py --check FILE                   自检: 能不能解析出结构完整的往返

--strict 会把每个进来的请求体与录制时那一份做**结构比对** (顶层键、消息角色
序列、工具名集合), 不一致就回 400 并在 stderr 打印差异. 这是这个工具最有用
的一半: 它把"今天的客户端还在发同样的报文形状"变成一条会失败的断言. 

注意回放时要给 ligent 加 --no-probe: 能力探测会插进额外的请求, 而那几轮
不在录制里. 录制时也建议关掉 (LIGENT_PROBE=0), 两边一致. 
"""
import argparse
import json
import socket
import sys
import threading

# ---------- 解析 /log 导出 ----------


def parse_log(text):
    """把一份 /log 导出切成若干次 HTTP 往返. 

    格式 (见 http.c 的记录点): 
        ===== <stamp> =====
        请求: POST <url>
        ----- 请求体 (N 字节) -----
        <N 字节>
        ----- 响应头 -----
        <header 行...>
        <空行>
        <响应原文>
        ----- 结束: HTTP <status> -----

    日志是**滚动缓冲**, 开头那一段可能是被截断的, 所以要容忍缺标记 —— 跳过
    解析不出来的块, 而不是整个失败. 
    """
    lines = text.split("\n")
    out = []
    i = 0
    while i < len(lines):
        if not (lines[i].startswith("===== ") and lines[i].endswith(" =====")):
            i += 1
            continue
        stamp = lines[i][6:-6]
        i += 1
        ex = {"stamp": stamp, "url": None, "method": None,
              "req": b"", "status": None, "headers": [], "resp": b""}

        # 请求行
        if i < len(lines) and lines[i].startswith("请求: "):
            parts = lines[i][len("请求: "):].split(" ", 1)
            ex["method"] = parts[0]
            ex["url"] = parts[1] if len(parts) > 1 else ""
            i += 1

        # 请求体: 按标记里给出的字节数精确取, 比按行取稳
        if i < len(lines) and lines[i].startswith("----- 请求体 ("):
            n = int(lines[i].split("(")[1].split(" ")[0])
            i += 1
            # 从这一行起累加字节, 直到凑够 n
            blob = ""
            while i < len(lines) and len(blob.encode()) < n:
                blob += (lines[i] + "\n")
                i += 1
            ex["req"] = blob.encode()[:n]
            # 上面多读了一行就把多余的还回去 (blob 里包含了尾部换行)
            consumed = len(blob.encode())
            if consumed > n:
                # 把这一行重新交回处理 (只可能多出换行)
                pass

        # 响应头
        if i < len(lines) and lines[i] == "----- 响应头 -----":
            i += 1
            while i < len(lines) and lines[i] != "":
                ex["headers"].append(lines[i])
                i += 1
            i += 1                     # 跳过空行

        # 响应原文: 直到"结束"行
        start = i
        while i < len(lines) and not lines[i].startswith("----- 结束: "):
            i += 1
        ex["resp"] = ("\n".join(lines[start:i])).encode()

        if i < len(lines) and lines[i].startswith("----- 结束: "):
            tail = lines[i]
            try:
                ex["status"] = int(tail.split("HTTP ")[1].split()[0])
            except (IndexError, ValueError):
                ex["status"] = 200
            i += 1

        # 结构不完整 (多半是滚动缓冲把开头截了)就丢掉
        if ex["url"] and ex["status"] is not None:
            out.append(ex)
    return out


def shape_of(req_bytes):
    """一个请求体的**结构指纹**: 只看形状, 不看内容. 

    内容每轮都不一样 (历史在长), 但形状必须稳定 —— 顶层键、消息角色序列、
    工具名集合. 客户端一旦改了报文形状, 这里就会显示出来. 
    """
    try:
        body = json.loads(req_bytes.decode("utf-8", "replace"))
    except Exception as e:
        return {"error": "不是合法 JSON: %s" % e}
    msgs = body.get("messages") or []
    tools = body.get("tools") or []
    names = []
    for t in tools:
        n = ((t or {}).get("function") or {}).get("name")
        if n:
            names.append(n)
    return {
        "keys": sorted(body.keys()),
        "roles": [m.get("role") for m in msgs],
        "tools": sorted(names),
    }


# ---------- 离线服务端 ----------


class ReplayServer:
    def __init__(self, exchanges, strict, loop):
        self.ex = exchanges
        self.strict = strict
        self.loop = loop
        self.pos = 0
        self.mismatch = 0
        self.lock = threading.Lock()

    def next_exchange(self):
        with self.lock:
            if self.pos >= len(self.ex):
                if self.loop:
                    self.pos = 0
                else:
                    return None
            e = self.ex[self.pos]
            self.pos += 1
            return e

    def handle(self, method, path, req_bytes):
        e = self.next_exchange()
        if e is None:
            sys.stderr.write(
                "replay: 录制里没有更多响应了 (已经回了 %d 个). \n"
                "  客户端比录制时多发了请求 —— 多半是多了能力探测 (回放时请加 "
                "--no-probe)或者客户端行为变了. \n" % len(self.ex))
            return 500, {}, b'{"error":{"message":"replay: out of recorded responses"}}'

        if self.strict:
            want = shape_of(e["req"])
            got = shape_of(req_bytes)
            if want != got:
                self.mismatch += 1
                sys.stderr.write(
                    "replay: 第 %d 个请求的**结构**和录制时不一致: \n"
                    "  录制时: %s\n"
                    "  现在  : %s\n"
                    % (self.pos, json.dumps(want, ensure_ascii=False),
                       json.dumps(got, ensure_ascii=False)))
                return (400, {"Content-Type": "application/json"},
                        json.dumps({"error": {
                            "message": "replay strict: request shape changed"}}
                        ).encode())

        ctype = "application/json"
        for h in e["headers"]:
            if h.lower().startswith("content-type:"):
                ctype = h.split(":", 1)[1].strip()
        hdrs = {"Content-Type": ctype}
        return e["status"], hdrs, e["resp"]


def serve(exchanges, port, strict, loop):
    srv = ReplayServer(exchanges, strict, loop)
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    s.bind(("127.0.0.1", port))
    s.listen(8)
    print("replay server on %d (%d 个录制往返%s)"
          % (port, len(exchanges), ", strict" if strict else ""), flush=True)

    while True:
        conn, _ = s.accept()
        threading.Thread(target=handle_conn, args=(conn, srv),
                         daemon=True).start()


def handle_conn(conn, srv):
    try:
        f = conn.makefile("rb")
        req_line = f.readline()
        if not req_line:
            return
        parts = req_line.decode("latin-1").split()
        method = parts[0] if parts else "GET"
        path = parts[1] if len(parts) > 1 else "/"

        clen = 0
        while True:
            line = f.readline()
            if not line or line in (b"\r\n", b"\n"):
                break
            if line.lower().startswith(b"content-length:"):
                clen = int(line.split(b":", 1)[1].strip())
        body = f.read(clen) if clen > 0 else b""

        status, hdrs, resp = srv.handle(method, path, body)
        head = "HTTP/1.1 %d %s\r\n" % (status, "OK" if status == 200 else "Error")
        for k, v in hdrs.items():
            head += "%s: %s\r\n" % (k, v)
        head += "Content-Length: %d\r\nConnection: close\r\n\r\n" % len(resp)
        conn.sendall(head.encode() + resp)
    except Exception as e:                       # noqa: BLE001 - 服务端不该因一个连接死掉
        sys.stderr.write("replay: 连接处理出错: %s\n" % e)
    finally:
        try:
            conn.close()
        except OSError:
            pass


# ---------- 入口 ----------


def main():
    ap = argparse.ArgumentParser(description="回放 /log 录下来的真实会话")
    ap.add_argument("file")
    ap.add_argument("--summary", action="store_true", help="只打印这份录制里有什么")
    ap.add_argument("--check", action="store_true", help="自检能否解析成完整往返")
    ap.add_argument("--serve", type=int, metavar="PORT", help="当离线服务端")
    ap.add_argument("--strict", action="store_true",
                    help="比对每个请求的结构, 不一致就回 400")
    ap.add_argument("--loop", action="store_true", help="响应用完后从头再来")
    a = ap.parse_args()

    text = open(a.file, encoding="utf-8", errors="replace").read()
    ex = parse_log(text)

    if a.summary or a.check:
        print("%s: %d 次往返" % (a.file, len(ex)))
        for i, e in enumerate(ex):
            r = e["resp"]
            kind = "SSE" if r.lstrip().startswith(b"data:") else "JSON"
            ntool = r.count(b'"tool_calls"')
            print("  #%d HTTP %s %-28s 请求 %5d 字节 → 响应 %6d 字节 %s%s"
                  % (i + 1, e["status"], (e["url"] or "").split("/")[-1],
                     len(e["req"]), len(r), kind,
                     "  含 tool_calls" if ntool else ""))
        if a.check:
            bad = [i + 1 for i, e in enumerate(ex) if not e["req"] or e["status"] is None]
            if bad:
                print("解析不完整的往返: %s" % bad, file=sys.stderr)
                return 1
            print("解析自检通过")
        return 0

    if a.serve:
        if not ex:
            print("这份录制里没有可回放的往返", file=sys.stderr)
            return 1
        serve(ex, a.serve, a.strict, a.loop)
        return 0

    ap.error("要么 --summary/--check, 要么 --serve PORT")


if __name__ == "__main__":
    sys.exit(main())