#!/usr/bin/env python3
"""验证 pCCU WebSocket 推送的 JSON 快照格式与分组完整性。"""
import socket, base64, os, struct, time, json, sys

HOST, PORT = "127.0.0.1", 18081

class WsClient:
    def __init__(self):
        self.s = socket.create_connection((HOST, PORT), timeout=5)
        key = base64.b64encode(os.urandom(16)).decode()
        req = (f"GET /ws HTTP/1.1\r\nHost: {HOST}:{PORT}\r\n"
               f"Upgrade: websocket\r\nConnection: Upgrade\r\n"
               f"Sec-WebSocket-Key: {key}\r\nSec-WebSocket-Version: 13\r\n\r\n")
        self.s.sendall(req.encode())
        self.buf = b""
        # 读取 HTTP 101 响应头
        while b"\r\n\r\n" not in self.buf:
            chunk = self.s.recv(4096)
            if not chunk:
                raise RuntimeError("连接关闭")
            self.buf += chunk
        head, self.buf = self.buf.split(b"\r\n\r\n", 1)
        if b"101" not in head.split(b"\r\n")[0]:
            raise RuntimeError("WebSocket 升级失败: " + head.decode(errors="replace"))

    def _read(self, n):
        while len(self.buf) < n:
            chunk = self.s.recv(4096)
            if not chunk:
                raise RuntimeError("连接关闭")
            self.buf += chunk
        out, self.buf = self.buf[:n], self.buf[n:]
        return out

    def recv_text(self, timeout=6):
        self.s.settimeout(timeout)
        end = time.time() + timeout
        while time.time() < end:
            if len(self.buf) < 2:
                try:
                    chunk = self.s.recv(4096)
                except socket.timeout:
                    return None
                if not chunk:
                    return None
                self.buf += chunk
                continue
            b1, b2 = self.buf[0], self.buf[1]
            opcode = b1 & 0x0F
            length = b2 & 0x7F
            hdr_len = 2
            self.buf = self.buf[2:]  # 消费帧头前2字节
            if length == 126:
                length = struct.unpack(">H", self._read(2))[0]
                hdr_len += 2
            elif length == 127:
                length = struct.unpack(">Q", self._read(8))[0]
                hdr_len += 8
            masked = (b2 & 0x80) != 0
            mask = self._read(4) if masked else b""
            payload = self._read(length)
            if opcode == 1:  # text
                if masked:
                    payload = bytes(b ^ mask[i % 4] for i, b in enumerate(payload))
                return payload.decode("utf-8", "replace")
            # ping/pong/close 跳过，继续
        return None

def main():
    try:
        ws = WsClient()
    except Exception as e:
        print("FAIL: WS 连接失败:", e)
        return 1
    text = ws.recv_text()
    if not text:
        print("FAIL: 未收到快照")
        return 1
    try:
        data = json.loads(text)
    except Exception as e:
        print("FAIL: 快照非合法JSON:", text[:200])
        return 1
    ok = True
    for key in ("fc", "pmCmd", "links", "fcStatusCount", "pmControlCount", "logs"):
        ok = ok and key in data
    print(f"PASS: 快照 JSON 有效, 分组keys={sorted(data.keys())}")
    print(f"  fc.fc_status={data['fc']['fc_status']}, fc.fc_mode={data['fc']['fc_mode']}, "
          f"heartbeat={data['fc']['heartbeat']}")
    print(f"  links={json.dumps(data['links'], ensure_ascii=False)}")
    print(f"  logs条数={len(data['logs'])}, fcStatusCount={data['fcStatusCount']}")
    ws.s.close()
    return 0 if ok else 1

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