# -*- coding: utf-8 -*-
"""End-to-end test: WS connect -> subscribe -> QMT push -> WS receive."""
import sys, io, os, json, time, socket, struct, base64, urllib.request

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "src"))
if sys.stdout is None:
    sys.stdout = io.StringIO()

import bridge_util, bridge_http_server, bridge_subscription

class FakeCtx:
    def __init__(self):
        self.subs = {}
        self._next = 1
        self.callback = None
    def subscribe_whole_quote(self, codes, callback=None):
        h = self._next
        self._next += 1
        self.subs[h] = list(codes)
        self.callback = callback
        return h
    def unsubscribe_quote(self, h):
        self.subs.pop(h, None)

fake = FakeCtx()
bridge_util.CTX = fake
bridge_http_server.PORT = 18691
bridge_http_server.TOKEN = ""
bridge_http_server.start_server()
time.sleep(0.5)

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_recv(s, timeout=5):
    s.settimeout(timeout)
    hdr = s.recv(2)
    if len(hdr) < 2:
        return None
    b1 = hdr[1]
    ln = b1 & 0x7F
    if ln == 126:
        ln = struct.unpack(">H", s.recv(2))[0]
    payload = b""
    while len(payload) < ln:
        chunk = s.recv(ln - len(payload))
        if not chunk:
            break
        payload += chunk
    return payload.decode("utf-8", "ignore")

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

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

# 3. trigger QMT push
time.sleep(0.3)
fake.callback({"600000.SH": {"lastPrice": 10.5, "volume": 200}})

# 4. WS should receive whole message
msg = ws_recv(ws)
print("ws received:", msg)
assert msg is not None, "ws received nothing!"
parsed = json.loads(msg)
assert parsed.get("type") == "whole", "expected type=whole, got %s" % parsed.get("type")
assert "600000.SH" in parsed.get("data", {}), "data missing code"
print("=== WS PUSH CHAIN OK ===")
ws.close()
bridge_http_server.stop_server()