#!/usr/bin/env python3
"""bolt HTTP 模式端到端测试"""
import http.client, json, subprocess, threading, time

BOLT = "/root/bolt/target/release/bolt"
PORT = 18080
passed = failed = 0
def check(name, cond, detail=""):
    global passed, failed
    if cond:
        passed += 1
        print(f"  PASS {name}")
    else:
        failed += 1
        print(f"  FAIL {name} {detail}")

proc = subprocess.Popen([BOLT, "--http", f"127.0.0.1:{PORT}"], stderr=subprocess.PIPE)
time.sleep(0.5)

def post(obj, auth=None, origin=None, accept="application/json", path="/mcp"):
    c = http.client.HTTPConnection("127.0.0.1", PORT, timeout=30)
    headers = {"Content-Type": "application/json", "Accept": accept}
    if auth: headers["Authorization"] = auth
    if origin: headers["Origin"] = origin
    body = json.dumps(obj)
    c.request("POST", path, body=body, headers=headers)
    r = c.getresponse()
    data = r.read()
    ct = r.getheader("Content-Type", "")
    c.close()
    return r.status, ct, data

# 1. initialize(JSON 响应)
st, ct, data = post({"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {"protocolVersion": "2025-06-18", "capabilities": {}, "clientInfo": {"name": "t", "version": "1"}}})
m = json.loads(data)
check("init json", st == 200 and "application/json" in ct and m["result"]["serverInfo"]["name"] == "bolt", f"{st} {ct} {data[:100]}")

# 2. SSE 响应(accept 含 text/event-stream)
st, ct, data = post({"jsonrpc": "2.0", "id": 2, "method": "ping"}, accept="application/json, text/event-stream")
sseline = data.split(b"data: ")[1].split(b"\r\n")[0].strip()
m = json.loads(sseline)
check("ping sse", st == 200 and "text/event-stream" in ct and m.get("result") == {}, f"{st} {ct} {data[:150]}")

# 3. 通知 → 202
st, ct, data = post({"jsonrpc": "2.0", "method": "notifications/initialized"})
check("notif 202", st == 202 and len(data) == 0, f"{st} {data[:50]}")

# 4. tools/call
st, ct, data = post({"jsonrpc": "2.0", "id": 3, "method": "tools/call", "params": {"name": "run_command", "arguments": {"command": "echo HTTP_OK && exit 5"}}})
m = json.loads(data)
sc = m["result"]["structuredContent"]
check("run over http", sc["exit_code"] == 5 and sc["stdout"].strip() == "HTTP_OK", str(sc))

# 5. GET SSE 流:开流收心跳,并发命令进度通知广播
got_event = threading.Event()
got_heartbeat = threading.Event()
msgs = []
def sse_get():
    c = http.client.HTTPConnection("127.0.0.1", PORT, timeout=35)
    c.request("GET", "/mcp", headers={"Accept": "text/event-stream"})
    r = c.getresponse()
    assert r.status == 200 and "text/event-stream" in r.getheader("Content-Type")
    buf = b""
    end = time.time() + 30
    while time.time() < end:
        chunk = r.read(1)
        if not chunk: break
        buf += chunk
        if buf.endswith(b"\r\n\r\n"):
            msgs.append(buf.decode(errors="replace"))
            buf = b""
            if msgs and (": ping" in msgs[-1]):
                got_heartbeat.set()
    c.close()
t = threading.Thread(target=sse_get)
t.start()
time.sleep(0.7)
# 发一个带进度 token 的命令(约 1.2s)
st, ct, data = post({"jsonrpc": "2.0", "id": 4, "method": "tools/call", "params": {"name": "run_command", "arguments": {"command": "for x in 1 2 3 4; do echo h-$x; sleep 0.3; done"}, "_meta": {"progressToken": "h1"}}})
m = json.loads(data)
check("long cmd done", m["result"]["structuredContent"]["stdout"].count("h-") == 4, str(m)[:100])
saw_progress = any("notifications/progress" in x and "h1" in x for x in msgs)
check("sse got progress broadcast", saw_progress, str(msgs)[:200])
t.join(timeout=25)
saw_hb = got_heartbeat.is_set()
check("sse heartbeat", saw_hb, str(msgs)[:200])

# 6. 取消:连接 A 长命令,连接 B 发 cancelled
st, ct, data = post({"jsonrpc": "2.0", "id": 5, "method": "tools/call", "params": {"name": "run_command", "arguments": {"command": "echo nothing"}}})  # 预热(连接池无关,每次新连接)

def long_cmd():
    return post({"jsonrpc": "2.0", "id": 6, "method": "tools/call", "params": {"name": "run_command", "arguments": {"command": "sleep 30"}}})
res = {}
t2 = threading.Thread(target=lambda: res.update(r=long_cmd()))
t2.start()
time.sleep(0.5)
st, ct, data = post({"jsonrpc": "2.0", "method": "notifications/cancelled", "params": {"requestId": 6}})
check("cancel notif 202", st == 202)
t2.join(timeout=10)
m = json.loads(res["r"][2])
check("http cancel works", m["result"]["structuredContent"]["cancelled"] is True, str(m)[:150])

# 7. 并发两个命令
def run(cmd, rid):
    return post({"jsonrpc": "2.0", "id": rid, "method": "tools/call", "params": {"name": "run_command", "arguments": {"command": cmd}}})
results = {}
ts = [threading.Thread(target=lambda c=cmd, i=i: results.update({i: run(c, i)})) for i, cmd in [(10, "sleep 0.5; echo A"), (11, "sleep 0.1; echo B")]]
for t in ts: t.start()
for t in ts: t.join(timeout=10)
check("http concurrent", json.loads(results[10][2])["result"]["structuredContent"]["stdout"].strip() == "A" and json.loads(results[11][2])["result"]["structuredContent"]["stdout"].strip() == "B")

# 8. 错误路径
st, _, _ = post({"jsonrpc": "2.0", "id": 99, "method": "nope"})
check("unknown method", st == 200)
c = http.client.HTTPConnection("127.0.0.1", PORT, timeout=5)
c.request("POST", "/mcp", body=b"this is {not json", headers={"Content-Type": "application/json", "Accept": "application/json"})
r = c.getresponse(); bd = r.read(); c.close()
check("bad json -> 400", r.status == 400, str(r.status))
st, _, data = post({"jsonrpc": "2.0", "id": 99, "method": "tools/call", "params": {"name": "nope", "arguments": {}}})
check("unknown tool -32602", json.loads(data)["error"]["code"] == -32602)
c = http.client.HTTPConnection("127.0.0.1", PORT, timeout=5)
c.request("GET", "/other", headers={"Accept": "text/event-stream"})
r = c.getresponse(); r.read(); c.close()
check("404 other path", r.status == 404)
c = http.client.HTTPConnection("127.0.0.1", PORT, timeout=5)
c.request("PUT", "/mcp")
r = c.getresponse(); r.read(); c.close()
check("405 method", r.status == 405)
c = http.client.HTTPConnection("127.0.0.1", PORT, timeout=5)
c.request("OPTIONS", "/mcp")
r = c.getresponse(); r.read(); c.close()
check("options cors", r.status == 204 and r.getheader("Access-Control-Allow-Origin") == "*")

# 9. origin 校验
st, ct, data = post({"jsonrpc": "2.0", "id": 50, "method": "ping"}, origin="https://evil.example.com")
check("evil origin 403", st == 403, str(st))
st, ct, data = post({"jsonrpc": "2.0", "id": 51, "method": "ping"}, origin="http://localhost:4173")
check("localhost origin ok", st == 200, str(st))

# 10. GET 无 Accept 头也返回 SSE(宽松)
c = http.client.HTTPConnection("127.0.0.1", PORT, timeout=5)
c.request("GET", "/mcp")
r = c.getresponse()
c.close()
check("get sse default", r.status == 200 and "text/event-stream" in r.getheader("Content-Type"))

proc.kill()
err = proc.stderr.read().decode()
print(f"\n=== HTTP: {passed} passed, {failed} failed ===")
if err.strip():
    print("bolt stderr:", err[:400])
sys_exit = 1 if failed else 0
import sys; sys.exit(sys_exit)