# -*- coding: utf-8 -*-
"""Real QMT WS push test: connect /ws, subscribe whole, receive pushes."""
import socket, struct, base64, os, json, time, urllib.request

def ws_connect(port):
    s = socket.create_connection(("127.0.0.1", port), timeout=5)
    key = base64.b64encode(os.urandom(16)).decode()
    req = ("GET /ws HTTP/1.1\r\nHost: 127.0.0.1:%d\r\nUpgrade: websocket\r\n"
           "Connection: Upgrade\r\nSec-WebSocket-Key: %s\r\n"
           "Sec-WebSocket-Version: 13\r\n\r\n" % (port, key))
    s.sendall(req.encode())
    resp = s.recv(4096)
    assert b"101" in resp, resp[:100]
    return s

def ws_read_frame(s, timeout=15):
    """Read one frame; returns (opcode, payload_str) or None."""
    s.settimeout(timeout)
    hdr = s.recv(2)
    if len(hdr) < 2:
        return None
    opcode = hdr[0] & 0x0F
    b1 = hdr[1]
    ln = b1 & 0x7F
    if ln == 126:
        ln = struct.unpack(">H", s.recv(2))[0]
    elif ln == 127:
        ln = struct.unpack(">Q", s.recv(8))[0]
    payload = b""
    while len(payload) < ln:
        chunk = s.recv(ln - len(payload))
        if not chunk:
            break
        payload += chunk
    return (opcode, payload.decode("utf-8", "ignore"))

# 1. connect
ws = ws_connect(8610)
print("WS connected")

# 2. subscribe whole
req = urllib.request.Request(
    "http://127.0.0.1:8610/data/subscribe",
    data=json.dumps({"type": "whole", "codes": ["SH", "SZ"]}).encode(),
    headers={"Content-Type": "application/json"},
)
with urllib.request.urlopen(req, timeout=10) as r:
    resp = json.loads(r.read().decode())
    print("subscribe:", resp)
    sub_id = resp.get("sub_id")

# 3. receive frames for ~20s, looking for whole pushes
print("waiting for whole pushes (up to 20s)...")
got = 0
deadline = time.time() + 20
while time.time() < deadline and got < 3:
    try:
        frame = ws_read_frame(ws, timeout=5)
    except socket.timeout:
        continue
    if frame is None:
        break
    opcode, payload = frame
    if opcode == 0x1:  # text
        try:
            parsed = json.loads(payload)
            t = parsed.get("type")
            if t == "whole" and parsed.get("data"):
                codes = list(parsed["data"].keys())
                sample = parsed["data"].get(codes[0], {}) if codes else {}
                print("whole push: %d codes, sample=%s" % (len(codes), json.dumps(sample, ensure_ascii=False)[:120]))
                got += 1
            else:
                print("text frame type=%s" % t)
        except Exception as e:
            print("parse err:", e, "payload:", payload[:80])
    elif opcode == 0x9:  # ping -> pong
        ws.sendall(bytes([0x8A, 0x00]))
        print("ping -> pong")
    elif opcode == 0x8:
        print("close frame")
        break

# 4. unsubscribe
req = urllib.request.Request(
    "http://127.0.0.1:8610/data/unsubscribe",
    data=json.dumps({"sub_id": sub_id}).encode(),
    headers={"Content-Type": "application/json"},
)
with urllib.request.urlopen(req, timeout=10) as r:
    print("unsubscribe:", r.read().decode())

ws.close()
print("=== DONE got_whole=%d ===" % got)