"""End-to-end exercise of the v6 protocol against server2018.py."""
import base64, hashlib, json, os, secrets, socket, struct, sys, time, urllib.request, zlib

HOST, PORT = "127.0.0.1", 8093
BASE = "http://%s:%d/api/v6" % (HOST, PORT)
results = []


def E(cond, label):
    results.append(bool(cond))
    print(("  PASS  " if cond else "  FAIL  ") + label, flush=True)


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=10) as r:
            raw = r.read()
            return r.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 GW:
    """v6 gateway client, zlib-stream aware."""

    def __init__(self, compress=True):
        self.compress = compress
        self.sock = socket.create_connection((HOST, PORT), timeout=10)
        key = base64.b64encode(secrets.token_bytes(16)).decode()
        q = "/?encoding=json&v=6" + ("&compress=zlib-stream" if compress else "")
        self.sock.sendall((
            "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" % (q, HOST, PORT, key)).encode())
        buf = b""
        while b"\r\n\r\n" not in buf:
            c = self.sock.recv(4096)
            if not c:
                raise RuntimeError("handshake died")
            buf += c
        head, self.buf = buf.split(b"\r\n\r\n", 1)
        self.status = head.split(b"\r\n")[0]
        accept = base64.b64encode(hashlib.sha1(
            (key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11").encode()).digest()).decode()
        assert accept.encode() in head, "bad accept key"
        self.inflator = zlib.decompressobj() if compress else None
        self.zbuf = b""

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

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

    def recv(self, timeout=10):
        deadline = time.monotonic() + timeout
        while True:
            hdr = self._need(2, deadline)
            op, n = hdr[0] & 0x0F, 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:
                return {"__close__": struct.unpack(">H", data[:2])[0] if len(data) >= 2 else None}
            if self.compress:
                # zlib-stream: accumulate until the 00 00 FF FF sync marker,
                # exactly as the client's pako Inflate does.
                self.zbuf += data
                if len(self.zbuf) < 4 or self.zbuf[-4:] != b"\x00\x00\xff\xff":
                    continue
                raw = self.inflator.decompress(self.zbuf)
                self.zbuf = b""
                return json.loads(raw.decode())
            return json.loads(data.decode())


tag = secrets.token_hex(3)

print("== accounts (v6) ==")
st, u = api("POST", "/auth/register",
            {"username": "alice", "email": "a%s@x.dev" % tag, "password": "hunter22"})
E(st == 201 and u.get("token"), "register -> 201 + token")
A = u["token"]
st, li = api("POST", "/auth/login", {"email": "a%s@x.dev" % tag, "password": "hunter22"})
E(st == 200 and li.get("token"), "login")
st, bad = api("POST", "/auth/login", {"email": "a%s@x.dev" % tag, "password": "nope"})
E(st == 400 and "email" in bad, "bad login -> field error")
st, me = api("GET", "/users/@me", None, A)
E(st == 200 and me["username"] == "alice" and "mfa_enabled" in me, "GET /users/@me (v6 shape)")

print("== gateway v6 handshake ==")
st, gw = api("GET", "/gateway")
E(st == 200 and gw["url"].startswith("ws"), "GET /gateway -> %s" % gw["url"])

ws = GW(compress=True)
E(b"101" in ws.status, "upgrade with compress=zlib-stream")
hello = ws.recv()
E(hello["op"] == 10, "server leads with op 10 HELLO (2015 had no hello frame)")
E(hello["d"]["heartbeat_interval"] == 41250, "HELLO carries heartbeat_interval")

ws.send({"op": 1, "d": None})
ack = ws.recv()
E(ack["op"] == 11, "op 1 HEARTBEAT -> op 11 HEARTBEAT_ACK (else client sees a zombie)")

ws.send({"op": 2, "d": {"token": A, "properties": {"$os": "Linux"}, "compress": False}})
ready = ws.recv()
E(ready["op"] == 0 and ready["t"] == "READY", "IDENTIFY -> READY")
d = ready["d"]
E(d["v"] == 6, "READY v=6")
for k in ("user", "session_id", "user_settings", "user_guild_settings", "read_state",
          "relationships", "private_channels", "guilds", "analytics_token",
          "experiments", "guild_experiments", "connected_accounts",
          "friend_suggestion_count", "notes", "sessions", "consents", "_trace"):
    E(k in d, "READY has %s" % k)
E(d["user"]["username"] == "alice", "READY user correct")

print("== bad token + resume ==")
bad_ws = GW()
E(bad_ws.recv()["op"] == 10, "HELLO before identify on second socket")
bad_ws.send({"op": 2, "d": {"token": "garbage", "properties": {}}})
r = bad_ws.recv()
E(r.get("op") == 9, "bad token -> op 9 INVALID_SESSION")

res = GW()
res.recv()
res.send({"op": 6, "d": {"token": A, "session_id": "nonexistent", "seq": 5}})
r = res.recv()
E(r.get("op") == 9 and r.get("d") is False, "unknown resume -> INVALID_SESSION d=false")

print("== guilds + channels (numeric types) ==")
st, g = api("POST", "/guilds", {"name": "Al Hang 2018", "region": "us-west"}, A)
E(st == 201, "create guild")
E(ws.recv()["t"] == "GUILD_CREATE", "GUILD_CREATE dispatched")
types = sorted(c["type"] for c in g["channels"])
E(types == [0, 2, 4], "channels are numeric types text/voice/category -> %s" % types)
E(any(r["id"] == g["id"] for r in g["roles"]), "@everyone role id == guild id")
E(g["member_count"] == 1 and g["large"] is False, "guild has member_count/large")
text = [c for c in g["channels"] if c["type"] == 0][0]

print("== messages ==")
st, m = api("POST", "/channels/%s/messages" % text["id"],
            {"content": "hello 2018", "nonce": "77", "tts": False}, A)
E(st == 200 and m["content"] == "hello 2018", "post message")
E(m.get("nonce") == "77", "nonce echoed")
for k in ("type", "pinned", "mention_roles", "reactions", "embeds", "attachments",
          "edited_timestamp", "guild_id"):
    E(k in m, "message has %s" % k)
ev = ws.recv()
E(ev["t"] == "MESSAGE_CREATE" and ev["d"]["nonce"] == "77", "MESSAGE_CREATE over zlib-stream")

st, msgs = api("GET", "/channels/%s/messages?limit=50" % text["id"], None, A)
E(st == 200 and len(msgs) == 1, "GET messages")
st, ed = api("PATCH", "/channels/%s/messages/%s" % (text["id"], m["id"]),
             {"content": "hello, 2018"}, A)
E(st == 200 and ed["edited_timestamp"], "edit message")
E(ws.recv()["t"] == "MESSAGE_UPDATE", "MESSAGE_UPDATE")
st, _ = api("POST", "/channels/%s/typing" % text["id"], {}, A)
E(st == 204 and ws.recv()["t"] == "TYPING_START", "typing -> TYPING_START")
st, _ = api("POST", "/channels/%s/messages/%s/ack" % (text["id"], m["id"]), {}, A)
E(st == 200 and ws.recv()["t"] == "MESSAGE_ACK", "ack -> MESSAGE_ACK")

print("== second user, invite, DM ==")
st, u2 = api("POST", "/auth/register",
             {"username": "bob", "email": "b%s@x.dev" % tag, "password": "hunter22"})
B = u2["token"]
st, inv = api("POST", "/channels/%s/invites" % text["id"], {}, A)
E(st == 200 and inv["code"], "create invite -> %s" % inv.get("code"))
E(inv["guild"]["name"] == "Al Hang 2018", "invite embeds guild")
ws_b = GW()
ws_b.recv()
ws_b.send({"op": 2, "d": {"token": B, "properties": {}}})
E(ws_b.recv()["t"] == "READY", "bob READY")
st, _ = api("POST", "/invite/%s" % inv["code"], {}, B)
E(st == 200, "bob accepts invite")
seen = set()
for _ in range(3):
    try:
        seen.add(ws.recv(3)["t"])
    except Exception:
        break
E("GUILD_MEMBER_ADD" in seen, "alice sees GUILD_MEMBER_ADD (%s)" % sorted(seen))

bob_id = u2["user_id"]
st, dm = api("POST", "/users/@me/channels", {"recipient_id": bob_id}, A)
E(st == 200 and dm["type"] == 1, "POST /users/@me/channels -> DM type 1")
E(isinstance(dm.get("recipients"), list) and dm["recipients"][0]["username"] == "bob",
  "DM uses recipients array (2015 used singular recipient)")

print("== presence + voice stub ==")
ws.send({"op": 3, "d": {"status": "idle", "game": None, "since": 0, "afk": False}})
found = False
for _ in range(4):
    try:
        p = ws.recv(3)
    except Exception:
        break
    if p["t"] == "PRESENCE_UPDATE" and p["d"]["status"] == "idle":
        found = True
        break
E(found, "op 3 -> PRESENCE_UPDATE idle")

voice = [c for c in g["channels"] if c["type"] == 2][0]
ws.send({"op": 4, "d": {"guild_id": g["id"], "channel_id": voice["id"],
                        "self_mute": False, "self_deaf": False}})
vs = None
for _ in range(4):
    try:
        p = ws.recv(3)
    except Exception:
        break
    if p["t"] == "VOICE_STATE_UPDATE":
        vs = p
        break
E(vs and vs["d"]["channel_id"] == voice["id"], "op 4 -> VOICE_STATE_UPDATE echoed")
none_server = True
try:
    for _ in range(2):
        if ws.recv(2)["t"] == "VOICE_SERVER_UPDATE":
            none_server = False
except Exception:
    pass
E(none_server, "no VOICE_SERVER_UPDATE -> voice stays a stub")

print("== uncompressed fallback + misc ==")
plain = GW(compress=False)
h = plain.recv()
E(h["op"] == 10, "plain (uncompressed) socket also works")
plain.send({"op": 2, "d": {"token": A, "properties": {}}})
E(plain.recv()["t"] == "READY", "READY without compression")

st, s = api("GET", "/users/@me/settings", None, A)
E(st == 200 and s["theme"] in ("dark", "light") and "locale" in s, "GET settings")
st, ex = api("GET", "/experiments", None, A)
E(st == 200 and "assignments" in ex, "GET /experiments")
st, _ = api("POST", "/science", {"events": []}, A)
E(st == 204, "POST /science")
st, rel = api("GET", "/users/@me/relationships", None, A)
E(st == 200, "GET relationships")
st, gl = api("GET", "/users/@me/guilds", None, A)
E(st == 200 and len(gl) == 1, "GET /users/@me/guilds")
st, _ = api("GET", "/channels/%s/messages" % text["id"], None, None)
E(st == 401, "no token -> 401")

print("== CDN paths ==")
import urllib.request as u2r
for path, label in (("/embed/avatars/0.png", "default avatar PNG"),):
    try:
        with u2r.urlopen("http://%s:%d%s" % (HOST, PORT, path), timeout=5) as r:
            blob = r.read()
        E(r.status == 200 and blob[:8] == b"\x89PNG\r\n\x1a\n", label)
    except Exception as e:
        E(False, "%s (%s)" % (label, e))

print()
print("ALL CHECKS PASSED" if all(results) else "%d FAILURES" % results.count(False))
sys.exit(0 if all(results) else 1)