#!/usr/bin/env python3
"""全面回归测试: 数据中心→out 多信号添加+立即推送验证"""
import socket, base64, os, struct, json, time

IP = '198.120.0.100'; PORT = 8000

def ws_handshake(sock):
    key = base64.b64encode(os.urandom(16)).decode()
    req = (f"GET /ws HTTP/1.1\r\nHost: {IP}:{PORT}\r\nUpgrade: websocket\r\n"
           f"Connection: Upgrade\r\nSec-WebSocket-Key: {key}\r\n"
           f"Sec-WebSocket-Version: 13\r\n\r\n")
    sock.sendall(req.encode())
    resp = b""
    while b"\r\n\r\n" not in resp:
        c = sock.recv(4096)
        if not c: break
        resp += c
    if b"101" not in resp.split(b"\r\n")[0]: raise Exception("握手失败")
    return resp

def ws_send(sock, msg):
    if isinstance(msg, dict): msg = json.dumps(msg, ensure_ascii=False)
    p = msg.encode(); L = len(p); mask = os.urandom(4)
    m = bytes(p[i] ^ mask[i%4] for i in range(L))
    f = bytearray([0x81])
    if L<126: f.append(0x80|L)
    elif L<65536: f.append(0x80|126); f.extend(struct.pack('>H',L))
    else: f.append(0x80|127); f.extend(struct.pack('>Q',L))
    f.extend(mask); f.extend(m); sock.sendall(bytes(f))

def ws_recv_all(sock, t=2.0):
    sock.settimeout(t); data=b""
    try:
        while True:
            c=sock.recv(65536)
            if not c: break
            data+=c
    except socket.timeout: pass
    msgs=[]; offset=0
    while offset+2<=len(data):
        b0,b1=data[offset],data[offset+1]; op=b0&0xF; L=b1&0x7F; offset+=2
        if L==126: L=struct.unpack('>H',data[offset:offset+2])[0]; offset+=2
        elif L==127: L=struct.unpack('>Q',data[offset:offset+8])[0]; offset+=8
        if op in (1,2) and offset+L<=len(data):
            p=data[offset:offset+L]
            try: msgs.append(json.loads(p.decode()))
            except: msgs.append(p.decode(errors='replace'))
        offset+=L
    return msgs

def find_out_push(msgs):
    for msg in msgs:
        if isinstance(msg, dict) and isinstance(msg.get('out'), list):
            return msg['out']
    return None

print("="*70)
print("  全面回归测试 — 数据中心→out 添加信号")
print("="*70)

sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.settimeout(10); sock.connect((IP,PORT))
ws_handshake(sock); print("[✓] 连接成功")

# 1. 注册 dc 信号
ws_send(sock, {"type":"cmd","data":"datacenter out"})
time.sleep(1); ws_recv_all(sock,1)

# 2. 选择 3 个不同前缀的信号测试
targets = [
    "sys.ch.tcp_s0_if",
    "sys.ch.uart0_if",
    "iec.run_cnt"
]
expected = set(targets)

# 3. 逐个添加，验证每个都能收到立即推送
print(f"\n--- 添加 {len(targets)} 个信号 ---")
all_ok = True
for i, t in enumerate(targets):
    ws_send(sock, {"saddr":t,"signal_type":"out","curd":"add","setting_zone":"0","signal_data":""})
    time.sleep(0.4)
    msgs = ws_recv_all(sock, 2)
    out_data = find_out_push(msgs)
    if out_data:
        saddr_set = {o['saddr'] for o in out_data}
        ok = t in saddr_set
        status = "✅" if ok else "❌"
        print(f"  [{i+1}] {t} add → 推送 {len(out_data)}个信号 {status}")
        if not ok:
            all_ok = False
            print(f"       推送内容: {saddr_set}")
    else:
        print(f"  [{i+1}] {t} add → ❌ 无推送")
        all_ok = False

# 4. 最终验证 — 检查所有 3 个是否都在列表中
print(f"\n--- 最终验证 ---")
ws_send(sock, {"type":"cmd","data":"datacenter out"})
time.sleep(1); msgs = ws_recv_all(sock, 2)
# 忽略 dc_data，看有没有 out 推送
out_data = None
for msg in msgs:
    if isinstance(msg, dict) and isinstance(msg.get('out'), list):
        out_data = msg['out']
if not out_data:
    # 再等 ws_task
    time.sleep(3)
    msgs = ws_recv_all(sock, 3)
    out_data = find_out_push(msgs)

if out_data:
    saddr_set = {o['saddr'] for o in out_data}
    all_found = expected.issubset(saddr_set)
    print(f"  out已注册: {len(out_data)}个 → 目标全部找到: {'✅' if all_found else '❌'}")
    if not all_found:
        print(f"  缺少: {expected - saddr_set}")
else:
    print("  ❌ 未收到 out 推送")

# 5. 清理
print(f"\n--- 清理 ---")
for t in targets:
    ws_send(sock, {"saddr":t,"signal_type":"out","curd":"del","setting_zone":"0","signal_data":""})
    time.sleep(0.3)
    msgs = ws_recv_all(sock, 1)
    out_data = find_out_push(msgs)
    cnt = len(out_data) if out_data else '?'
    print(f"  del {t} → 剩余 {cnt} 个")

print(f"\n{'✅ 全部通过!' if all_ok else '❌ 存在失败'}")
try: sock.sendall(b'\x88\x80'+os.urandom(4))
except: pass
sock.close()