"""验证 keep-alive 连接上「不读 body 的接口」不会污染下一个请求。

背景（真实 bug）：`/api/devices/connect` 的 handler 从不读请求体。
在 HTTP/1.1 keep-alive 连接上，没读走的 `{}` 会留在 socket 缓冲区里，
被 BaseHTTPRequestHandler 当成下一个请求的请求行，于是登录到：

    501 Unsupported method ('{}POST')

而且是**间歇性**的（取决于客户端是否复用连接），非常难查。
修法：do_POST 进入路由前先把 body 读完缓存起来。

用法：先起一个 mock 服务，再跑本脚本。
    .venv\\Scripts\\python.exe tools\\_keepalive_probe.py --port 8855
"""
from __future__ import annotations

import argparse
import json
import socket
import sys


def build_post(host: str, port: int, path: str, body: str) -> bytes:
    payload = body.encode("utf-8")
    head = (
        f"POST {path} HTTP/1.1\r\n"
        f"Host: {host}:{port}\r\n"
        f"Content-Type: application/json\r\n"
        f"Content-Length: {len(payload)}\r\n"
        f"Connection: keep-alive\r\n"
        f"\r\n"
    ).encode("ascii")
    return head + payload


def recv_response(sock: socket.socket) -> str:
    """读一个完整响应（按 Content-Length）。"""
    data = b""
    while b"\r\n\r\n" not in data:
        chunk = sock.recv(4096)
        if not chunk:
            return data.decode("utf-8", "replace")
        data += chunk
    head, _, rest = data.partition(b"\r\n\r\n")
    length = 0
    for line in head.split(b"\r\n"):
        if line.lower().startswith(b"content-length:"):
            length = int(line.split(b":", 1)[1].strip())
            break
    while len(rest) < length:
        chunk = sock.recv(4096)
        if not chunk:
            break
        rest += chunk
    return (head + b"\r\n\r\n" + rest).decode("utf-8", "replace")


def status_of(response: str) -> int:
    first = response.split("\r\n", 1)[0]
    parts = first.split()
    try:
        return int(parts[1])
    except (IndexError, ValueError):
        return -1


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=8855)
    args = parser.parse_args()

    failures: list[str] = []

    def check(name: str, ok: bool, detail: str = "") -> None:
        print(f"  [{'PASS' if ok else 'FAIL'}] {name}" + (f" — {detail}" if detail else ""))
        if not ok:
            failures.append(name)

    with socket.create_connection((args.host, args.port), timeout=15) as sock:
        # 第 1 个请求：body 为 {} 且服务端 handler 不读它
        sock.sendall(build_post(args.host, args.port, "/api/devices/connect", "{}"))
        r1 = recv_response(sock)
        s1 = status_of(r1)
        check("第 1 个请求(connect) 成功", s1 in (200, 409), f"status={s1}")
        print(f"    第 1 个响应前 120 字符：{r1[:120]!r}")

        # 第 2 个请求：如果残留的 {} 没被排空，请求行会被污染 → 501
        body2 = json.dumps({"lat": 39.9087, "lng": 116.3975})
        sock.sendall(build_post(args.host, args.port, "/api/location", body2))
        r2 = recv_response(sock)
        s2 = status_of(r2)
        check("同连接第 2 个请求(location) 未被污染", s2 == 200, f"status={s2}")
        if s2 != 200:
            print(f"    第 2 个响应：{r2[:300]!r}")

        # 第 3 个请求：再压一次，确认连续多请求也稳
        sock.sendall(build_post(args.host, args.port, "/api/location", body2))
        r3 = recv_response(sock)
        s3 = status_of(r3)
        check("同连接第 3 个请求仍然正常", s3 == 200, f"status={s3}")

    print()
    if failures:
        print(f"失败 {len(failures)} 项：{failures}")
        return 1
    print("全部通过：keep-alive 上不会串请求")
    return 0


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