"""Minimal WS client + end-to-end exercise of the 2015 protocol."""
import base64, hashlib, json, os, secrets, socket, struct, sys, threading, time, urllib.request

HOST, PORT = "127.0.0.1", 8080
BASE = "http://%s:%d/api" % (HOST, PORT)


def api(method, path, body=None, token=None):
    data = json.dumps(body).encode() if body is not None else None
    req = urllib.request.Request(BASE + path, data=data, method=method)
    if data:
        req.add_header("Content-Type", "application/json")
    if token:
        req.add_header("Authorization", token)
    try:
        with urllib.request.urlopen(req, timeout=5) as resp:
            raw = resp.read()
            return resp.status, (json.loads(raw) if raw else None)
    except urllib.error.HTTPError as e:
        raw = e.read()
        return e.code, (json.loads(raw) if raw else None)


class WS:
    def __init__(self, host, port, path):
        self.sock = socket.create_connection((host, port), timeout=10)
        key = base64.b64encode(secrets.token_bytes(16)).decode()
        req = (
            "GET %s HTTP/1.1\r\nHost: %s:%d\r\nUpgrade: websocket\r\n"
            "Connection: Upgrade\r\nSec-WebSocket-Key: %s\r\n"
            "Sec-WebSocket-Version: 13\r\n\r\n" % (path, host, port, key)
        )
        self.sock.sendall(req.encode())
        buf = b""
        while b"\r\n\r\n" not in buf:
            chunk = self.sock.recv(1)
            if not chunk:
                raise RuntimeError("handshake failed: %r" % buf)
            buf += chunk
        assert b"101" in buf.split(b"\r\n")[0], buf
        accept = base64.b64encode(
            hashlib.sha1((key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11").encode()).digest()
        ).decode()
        assert accept.encode() in buf, "bad accept key"
        # NB: no makefile() -- a timed-out read on a BufferedReader throws away
        # whatever it had already buffered, which desynchronises the framing.
        # Own the buffer instead so a timeout is always safe.
        self.buf = buf.split(b"\r\n\r\n", 1)[1]

    def send(self, obj):
        payload = json.dumps(obj).encode()
        mask = secrets.token_bytes(4)
        n = len(payload)
        hdr = bytearray([0x81])
        if n < 126:
            hdr.append(0x80 | n)
        elif n < 65536:
            hdr.append(0x80 | 126); hdr += struct.pack(">H", n)
        else:
            hdr.append(0x80 | 127); hdr += struct.pack(">Q", n)
        hdr += mask
        masked = bytes(b ^ mask[i % 4] for i, b in enumerate(payload))
        self.sock.sendall(bytes(hdr) + masked)

    def _need(self, count, deadline):
        while len(self.buf) < count:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                raise TimeoutError("no frame")
            self.sock.settimeout(remaining)
            chunk = self.sock.recv(65536)
            if not chunk:
                raise ConnectionError("closed")
            self.buf += chunk
        out, self.buf = self.buf[:count], self.buf[count:]
        return out

    def recv(self, timeout=5):
        deadline = time.monotonic() + timeout
        hdr = self._need(2, deadline)
        op = hdr[0] & 0x0F
        n = hdr[1] & 0x7F
        if n == 126:
            n = struct.unpack(">H", self._need(2, deadline))[0]
        elif n == 127:
            n = struct.unpack(">Q", self._need(8, deadline))[0]
        data = self._need(n, deadline) if n else b""
        if op == 0x8:
            code = struct.unpack(">H", data[:2])[0] if len(data) >= 2 else None
            return {"__close__": code, "reason": data[2:].decode("utf-8", "replace")}
        return json.loads(data.decode())


def expect(cond, label):
    print(("  PASS  " if cond else "  FAIL  ") + label)
    if not cond:
        globals()["FAILED"] = True


FAILED = False
suffix = secrets.token_hex(3)

print("== accounts ==")
st, alice = api("POST", "/auth/register", {"username": "alice", "email": "alice%s@x.dev" % suffix, "password": "hunter22"})
expect(st == 200 and alice.get("token"), "register alice -> token")
st, bob = api("POST", "/auth/register", {"username": "bob", "email": "bob%s@x.dev" % suffix, "password": "hunter22"})
expect(st == 200 and bob.get("token"), "register bob -> token")
st, dup = api("POST", "/auth/register", {"username": "carol", "email": "alice%s@x.dev" % suffix, "password": "hunter22"})
expect(st == 400 and "email" in dup, "duplicate email rejected with field errors")
st, li = api("POST", "/auth/login", {"email": "alice%s@x.dev" % suffix, "password": "hunter22"})
expect(st == 200 and li.get("token"), "login with correct password")
A, B = alice["token"], bob["token"]

print("== gateway discovery + identify ==")
st, gw = api("GET", "/gateway")
expect(st == 200 and gw["url"].startswith("ws://"), "GET /gateway -> %s" % gw)

ws = WS(HOST, PORT, "/gateway")
ws.send({"op": 2, "d": {"token": A, "properties": {"$os": "Linux", "$browser": "Chrome"}, "v": 3}})
pkt = ws.recv()
expect(pkt["op"] == 0 and pkt["t"] == "READY", "op 2 -> READY dispatch")
d = pkt["d"]
expect(d.get("heartbeat_interval") == 41250, "heartbeat_interval inside READY payload")
expect(d["user"]["username"] == "alice" and d["user"].get("email"), "READY carries private user")
for key in ("session_id", "read_state", "guilds", "private_channels"):
    expect(key in d, "READY has %s" % key)
expect(pkt["s"] == 1, "dispatch carries seq")
session_id = d["session_id"]

bad = WS(HOST, PORT, "/gateway")
bad.send({"op": 2, "d": {"token": "garbage", "properties": {}, "v": 3}})
pkt = bad.recv()
expect(pkt.get("__close__") == 4004, "bad token -> close 4004 (client drops token)")

print("== guild + channel creation ==")
st, guild = api("POST", "/guilds", {"name": "Al Hang", "region": "us-west"}, A)
expect(st == 201 and guild["name"] == "Al Hang", "POST /guilds")
for key in ("channels", "members", "presences", "voice_states", "roles", "joined_at"):
    expect(key in guild, "guild payload has %s (client iterates it unguarded)" % key)
expect(any(r["id"] == guild["id"] for r in guild["roles"]), "@everyone role id == guild id")
pkt = ws.recv()
expect(pkt["t"] == "GUILD_CREATE" and pkt["d"]["id"] == guild["id"], "GUILD_CREATE dispatched to owner")
text = [c for c in guild["channels"] if c["type"] == "text"][0]
voice = [c for c in guild["channels"] if c["type"] == "voice"][0]

print("== messages ==")
nonce = "1234567890"
st, msg = api("POST", "/channels/%s/messages" % text["id"],
              {"content": "hello 2015", "mentions": [], "nonce": nonce, "tts": False}, A)
expect(st == 200 and msg["content"] == "hello 2015", "POST message")
expect(msg.get("nonce") == nonce, "nonce echoed (else the optimistic copy duplicates)")
for key in ("id", "channel_id", "author", "timestamp", "edited_timestamp", "embeds",
            "mentions", "mention_everyone", "attachments", "tts"):
    expect(key in msg, "message has %s" % key)
pkt = ws.recv()
expect(pkt["t"] == "MESSAGE_CREATE" and pkt["d"]["nonce"] == nonce, "MESSAGE_CREATE dispatched with nonce")

st, msgs = api("GET", "/channels/%s/messages?limit=50" % text["id"], None, A)
expect(st == 200 and len(msgs) == 1, "GET messages")

st, edited = api("PATCH", "/channels/%s/messages/%s" % (text["id"], msg["id"]),
                 {"content": "hello, 2015", "mentions": []}, A)
expect(st == 200 and edited["edited_timestamp"], "PATCH message sets edited_timestamp")
expect(ws.recv()["t"] == "MESSAGE_UPDATE", "MESSAGE_UPDATE dispatched")

st, _ = api("POST", "/channels/%s/typing" % text["id"], {}, A)
expect(st == 204, "POST typing")
expect(ws.recv()["t"] == "TYPING_START", "TYPING_START dispatched")

st, _ = api("POST", "/channels/%s/messages/%s/ack" % (text["id"], msg["id"]), {}, A)
expect(st == 200, "POST ack")
expect(ws.recv()["t"] == "MESSAGE_ACK", "MESSAGE_ACK dispatched")

print("== invites + a second member ==")
st, inv = api("POST", "/channels/%s/invites" % text["id"], {"xkcdpass": True}, A)
expect(st == 200 and inv["code"], "create invite -> %s" % inv.get("code"))
st, resolved = api("GET", "/invite/%s" % inv["code"])
expect(st == 200 and resolved["guild"]["name"] == "Al Hang", "resolve invite unauthenticated")

ws_b = WS(HOST, PORT, "/gateway")
ws_b.send({"op": 2, "d": {"token": B, "properties": {}, "v": 3}})
expect(ws_b.recv()["t"] == "READY", "bob READY")

st, accepted = api("POST", "/invite/%s" % inv["code"], {}, B)
expect(st == 200, "bob accepts invite")
pkt = ws_b.recv()
expect(pkt["t"] in ("GUILD_MEMBER_ADD", "GUILD_CREATE"), "bob gets guild events")
seen = set()
for _ in range(3):
    try:
        seen.add(ws.recv(2)["t"])
    except Exception:
        break
expect("GUILD_MEMBER_ADD" in seen, "alice sees GUILD_MEMBER_ADD (saw %s)" % sorted(seen))

st, bobguilds = api("GET", "/gateway")
ws_b2 = WS(HOST, PORT, "/gateway")
ws_b2.send({"op": 2, "d": {"token": B, "properties": {}, "v": 3}})
ready_b = ws_b2.recv()["d"]
expect(len(ready_b["guilds"]) == 1, "bob's READY now contains the guild")
members = ready_b["guilds"][0]["members"]
expect({m["user"]["username"] for m in members} == {"alice", "bob"},
       "members carry user objects (UserStore is seeded from these)")

print("== bob talks, alice hears ==")
st, m2 = api("POST", "/channels/%s/messages" % text["id"],
             {"content": "hey @alice", "mentions": [alice and ready_b["guilds"][0]["members"][0]["user"]["id"]], "nonce": "9", "tts": False}, B)
expect(st == 200, "bob posts")
got = None
for _ in range(4):
    try:
        pkt = ws.recv(2)
    except Exception:
        break
    if pkt["t"] == "MESSAGE_CREATE":
        got = pkt
        break
expect(got is not None and got["d"]["author"]["username"] == "bob", "alice receives bob's MESSAGE_CREATE")

print("== DMs ==")
alice_id = ready_b["guilds"][0]["members"][0]["user"]["id"]
bob_id = [m["user"]["id"] for m in members if m["user"]["username"] == "bob"][0]
st, dm = api("POST", "/users/%s/channels" % alice_id, {"recipient_id": bob_id}, A)
expect(st == 200 and dm["is_private"] and dm["recipient"]["username"] == "bob",
       "open DM -> is_private with singular recipient")

print("== presence + status ==")
ws.send({"op": 3, "d": {"idle_since": int(time.time() * 1000), "game_id": None}})
found = False
for _ in range(4):
    try:
        pkt = ws.recv(2)
    except Exception:
        break
    if pkt["t"] == "PRESENCE_UPDATE" and pkt["d"]["status"] == "idle":
        found = True
        break
expect(found, "op 3 -> PRESENCE_UPDATE idle")

print("== voice stub ==")
ws.send({"op": 4, "d": {"guild_id": guild["id"], "channel_id": voice["id"], "self_mute": False, "self_deaf": False}})
vs = None
for _ in range(4):
    try:
        pkt = ws.recv(2)
    except Exception:
        break
    if pkt["t"] == "VOICE_STATE_UPDATE":
        vs = pkt
        break
expect(vs is not None and vs["d"]["channel_id"] == voice["id"], "op 4 -> VOICE_STATE_UPDATE echoed")
never = True
try:
    for _ in range(2):
        if ws.recv(2)["t"] == "VOICE_SERVER_UPDATE":
            never = False
except Exception:
    pass
expect(never, "no VOICE_SERVER_UPDATE -> client parks, never opens a voice socket")

print("== heartbeat + settings + misc ==")
ws.send({"op": 1, "d": int(time.time() * 1000)})
st, s = api("GET", "/users/@me/settings", None, A)
expect(st == 200 and "muted_channels" in s and "enable_tts_command" in s, "settings keys match client map")
st, s2 = api("PATCH", "/users/@me/settings", {"theme": "light"}, A)
expect(st == 200 and s2["theme"] == "light", "PATCH settings")
st, tut = api("GET", "/tutorial", None, A)
expect(st == 200 and "indicators_suppressed" in tut, "GET tutorial")
st, regions = api("GET", "/voice/regions", None, A)
expect(st == 200 and any(r["id"] == "us-west" for r in regions), "GET voice/regions")
st, ice = api("GET", "/voice/ice", None, A)
expect(st == 200 and "servers" in ice, "GET voice/ice")
st, _ = api("POST", "/track", {"event": "App Opened", "properties": {}}, A)
expect(st == 204, "POST track")
st, me = api("PATCH", "/users/@me", {"username": "alice2", "email": "alice%s@x.dev" % suffix, "password": "hunter22"}, A)
expect(st == 200 and me.get("token") and me["username"] == "alice2", "PATCH /users/@me returns a fresh token")
st, _ = api("POST", "/auth/logout", {}, A)
expect(st == 204, "logout")

print("== unauth ==")
st, _ = api("GET", "/channels/%s/messages" % text["id"], None, None)
expect(st == 401, "no token -> 401")

print("\nFAILURES PRESENT" if FAILED else "\nALL CHECKS PASSED")
sys.exit(1 if FAILED else 0)