"""One WebSocket must isolate subscriptions and preserve the SSE replay cursor."""

import base64
import hashlib
import http.client
import json
import secrets
import socket
import struct
import time
import unittest
import urllib.parse

from harness import Albedo, Provider, text


class WebSocket:
    """A masked stdlib client for the daemon's JSON event socket."""

    def __init__(self, app):
        address = urllib.parse.urlsplit(app.base)
        self.socket = socket.create_connection((address.hostname, address.port), 20)
        self.reader = self.socket.makefile("rb")
        key = base64.b64encode(secrets.token_bytes(16)).decode()
        self.socket.sendall(
            (
                f"GET /events HTTP/1.1\r\nHost: {address.netloc}\r\n"
                "Connection: Upgrade\r\nUpgrade: websocket\r\n"
                f"Sec-WebSocket-Key: {key}\r\nSec-WebSocket-Version: 13\r\n"
                f"Authorization: Bearer {app.connection['token']}\r\n\r\n"
            ).encode()
        )
        try:
            status = self.reader.readline()
            if b" 101 " not in status:
                raise AssertionError(f"WebSocket upgrade failed: {status!r}")
            headers = {}
            while (line := self.reader.readline()) not in (b"\r\n", b""):
                name, value = line.decode().split(":", 1)
                headers[name.lower()] = value.strip()
            expected = base64.b64encode(
                hashlib.sha1(
                    (key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11").encode()
                ).digest()
            ).decode()
            if headers.get("sec-websocket-accept") != expected:
                raise AssertionError(f"bad WebSocket accept: {headers}")
        except BaseException:
            self.close()
            raise

    def close(self):
        self.reader.close()
        self.socket.close()

    def frame(self, opcode, payload):
        mask = secrets.token_bytes(4)
        size = len(payload)
        header = bytes([0x80 | opcode])
        if size < 126:
            header += bytes([0x80 | size])
        elif size < 65536:
            header += bytes([0x80 | 126]) + struct.pack("!H", size)
        else:
            header += bytes([0x80 | 127]) + struct.pack("!Q", size)
        self.socket.sendall(
            header
            + mask
            + bytes(value ^ mask[i % 4] for i, value in enumerate(payload))
        )

    def send(self, **message):
        self.frame(1, json.dumps(message).encode())

    def read_exact(self, size):
        value = self.reader.read(size)
        if len(value) != size:
            raise AssertionError("WebSocket closed before a complete frame")
        return value

    def receive(self):
        message = bytearray()
        while True:
            flags, length = self.read_exact(2)
            opcode = flags & 15
            if length & 128:
                raise AssertionError("server sent a masked frame")
            if length == 126:
                length = struct.unpack("!H", self.read_exact(2))[0]
            elif length == 127:
                length = struct.unpack("!Q", self.read_exact(8))[0]
            if length > 1048576:
                raise AssertionError("event frame exceeded 1 MiB")
            payload = self.read_exact(length)
            if opcode == 8:
                raise AssertionError(f"WebSocket closed: {payload!r}")
            if opcode == 9:
                self.frame(10, payload)
                continue
            if opcode == 10:
                continue
            if opcode not in (0, 1):
                raise AssertionError(f"unexpected opcode {opcode}")
            message.extend(payload)
            if len(message) > 1048576:
                raise AssertionError("event message exceeded 1 MiB")
            if flags & 128:
                return json.loads(message)


def user_texts(batch):
    return [
        part["text"]
        for event in batch["events"]
        if event["type"] == "message" and event["data"]["entry"]["kind"] == "user"
        for part in event["data"]["entry"]["content"]
        if part["kind"] == "text"
    ]


class EventsTest(unittest.TestCase):
    def setUp(self):
        self.provider = Provider(lambda _: text("socket answer"))
        self.addCleanup(self.provider.close)

    def connect(self, app):
        connection = WebSocket(app)
        self.addCleanup(connection.close)
        return connection

    def until(self, connection, predicate, *, reject=()):
        deadline = time.monotonic() + 20
        while time.monotonic() < deadline:
            reply = connection.receive()
            self.assertNotIn("failure", reply, reply)
            self.assertNotIn(reply.get("key"), reject, reply)
            if predicate(reply):
                return reply
        self.fail("event socket did not deliver the expected reply")

    def test_multiplexes_sessions_and_list_and_ends_only_the_deleted_subscription(self):
        with Albedo(self.provider) as app:
            sessions = {"left": app.session(), "right": app.session()}
            connection = self.connect(app)
            for key, session in sessions.items():
                connection.send(key=key, type="session", session_id=session, tail=0)
            connection.send(key="roots", type="sessions", scope="roots")
            initial = {}
            while len(initial) < 3:
                reply = connection.receive()
                self.assertNotIn("failure", reply, reply)
                initial[reply["key"]] = reply["batch"]
            for key, session in sessions.items():
                self.assertEqual(initial[key]["snapshot"]["id"], session)
            self.assertEqual(
                [event["type"] for event in initial["roots"]["events"]],
                ["ready", "reset"],
            )

            for key, session in sessions.items():
                app.prompt(session, f"question for {key}").close()
            observed = {key: [] for key in sessions}
            completed = set()
            listed = set()

            def both_completed(reply):
                key = reply["key"]
                if key in sessions:
                    batch = reply["batch"]
                    observed[key].extend(user_texts(batch))
                    if any(e["type"] == "turn_completed" for e in batch["events"]):
                        completed.add(key)
                elif key == "roots":
                    for event in reply["batch"]["events"]:
                        if event["type"] == "activity":
                            listed.add(event["data"]["session_id"])
                        elif event["type"] == "invalidate":
                            listed.update(event["data"]["session_ids"])
                return len(completed) == 2 and listed.issuperset(sessions.values())

            self.until(connection, both_completed)
            self.assertEqual(
                observed, {key: [f"question for {key}"] for key in sessions}
            )
            for session in sessions.values():
                app.idle(session)

            connection.send(key="left", type="stop")
            self.until(connection, lambda reply: reply == {"key": "left", "end": True})
            app.prompt(sessions["left"], "must not reach stopped subscription").close()
            app.idle(sessions["left"])
            app.prompt(sessions["right"], "right remains subscribed").close()
            self.until(
                connection,
                lambda reply: (
                    reply.get("key") == "right"
                    and "right remains subscribed" in user_texts(reply["batch"])
                ),
                reject=("left",),
            )
            app.idle(sessions["right"])
            resource = f"/sessions/{sessions['right']}?view=configuration"
            with app.api(resource) as response:
                validator = response.getheader("ETag")
                response.read()
            app.api(resource, method="DELETE", headers={"If-Match": validator}).close()
            self.until(
                connection,
                lambda reply: reply == {"key": "right", "end": True},
                reject=("left",),
            )
            connection.send(type="ping")
            self.until(
                connection,
                lambda reply: reply == {"type": "pong"},
                reject=("left", "right"),
            )

    def test_resume_uses_the_same_cursor_without_reset_or_duplicate_messages(self):
        with Albedo(self.provider) as app:
            session = app.session()
            app.prompt(session, "already consumed").close()
            app.idle(session)
            connection = self.connect(app)
            connection.send(key="before", type="session", session_id=session, tail=0)
            before = self.until(connection, lambda reply: reply.get("key") == "before")[
                "batch"
            ]
            connection.close()

            app.prompt(session, "while disconnected").close()
            app.idle(session)
            resumed = self.connect(app)
            resumed.send(
                key="after",
                type="session",
                session_id=session,
                tail=0,
                after_generation=before["generation"],
                after_seq=before["cursor"],
            )
            consumed = before["cursor"]
            observed = []

            def completed(reply):
                nonlocal consumed
                self.assertEqual(reply["key"], "after")
                batch = reply["batch"]
                self.assertEqual(batch["generation"], before["generation"])
                self.assertNotIn("reset", [event["type"] for event in batch["events"]])
                sequences = [event["sequence"] for event in batch["events"]]
                self.assertEqual(
                    sequences, list(range(consumed + 1, consumed + 1 + len(sequences)))
                )
                self.assertEqual(
                    batch["cursor"], sequences[-1] if sequences else consumed
                )
                consumed = batch["cursor"]
                observed.extend(user_texts(batch))
                return any(
                    event["type"] == "turn_completed" for event in batch["events"]
                )

            self.until(resumed, completed)
            self.assertEqual(observed, ["while disconnected"])

    def test_subscription_limits_and_cursor_failures_leave_other_keys_live(self):
        with Albedo(self.provider) as app:
            session = app.session()
            connection = self.connect(app)
            connection.send(key="keep", type="session", session_id=session, tail=0)
            self.until(connection, lambda reply: reply.get("key") == "keep")
            connection.send(
                key="bad-cursor",
                type="session",
                session_id=session,
                tail=0,
                after_generation="invalid",
                after_seq=0,
            )
            failure = connection.receive()
            self.assertEqual(failure["key"], "bad-cursor")
            self.assertIn("failure", failure)
            self.assertNotIn("batch", failure)
            for payload, key in [
                (b"{", ""),
                (b'{"key":"unknown","type":"unknown"}', "unknown"),
                (b'{"key":"duplicate","type":"ping","type":"ping"}', "duplicate"),
                (b'{"key":"extra","type":"stop","extra":true}', "extra"),
            ]:
                with self.subTest(command=payload):
                    connection.frame(1, payload)
                    failure = connection.receive()
                    self.assertEqual(failure["key"], key)
                    self.assertEqual(failure["failure"]["code"], "invalid_request")
                    self.assertNotIn("batch", failure)

            # Every key follows the same session: removing one must not unwatch another.
            others = {f"copy-{index}" for index in range(63)}
            for key in sorted(others):
                connection.send(key=key, type="session", session_id=session, tail=0)
            ready = set()

            def all_ready(reply):
                self.assertIn(reply["key"], others)
                self.assertEqual(reply["batch"]["snapshot"]["id"], session)
                ready.add(reply["key"])
                return ready == others

            self.until(connection, all_ready)
            connection.send(
                key="over-limit", type="session", session_id=session, tail=0
            )
            failure = connection.receive()
            self.assertEqual(failure["key"], "over-limit")
            self.assertEqual(failure["failure"]["code"], "follow_limit")
            self.assertNotIn("batch", failure)

            for key in sorted(others):
                connection.send(key=key, type="stop")
            stopped = set()

            def all_stopped(reply):
                self.assertEqual(reply, {"key": reply["key"], "end": True})
                self.assertIn(reply["key"], others)
                stopped.add(reply["key"])
                return stopped == others

            self.until(connection, all_stopped)
            connection.send(
                key="replacement", type="session", session_id=session, tail=0
            )
            self.until(
                connection,
                lambda reply: reply.get("key") == "replacement",
                reject=others,
            )
            app.prompt(session, "survives per-key failures and unsubscribe").close()
            observed = set()

            def both_live(reply):
                self.assertIn(reply["key"], {"keep", "replacement"})
                if "survives per-key failures and unsubscribe" in user_texts(
                    reply["batch"]
                ):
                    observed.add(reply["key"])
                return observed == {"keep", "replacement"}

            self.until(connection, both_live, reject=others)
            app.idle(session)

    def test_upgrade_requires_authentication_allowed_origin_and_websocket_headers(self):
        with Albedo(self.provider) as app:
            address = urllib.parse.urlsplit(app.base)
            token = {"Authorization": "Bearer " + app.connection["token"]}
            upgrade = {
                "Connection": "Upgrade",
                "Upgrade": "websocket",
                "Sec-WebSocket-Key": base64.b64encode(secrets.token_bytes(16)).decode(),
                "Sec-WebSocket-Version": "13",
            }
            for name, headers, status in [
                ("missing token", upgrade, 401),
                (
                    "forbidden origin",
                    {**upgrade, **token, "Origin": "https://evil.example"},
                    403,
                ),
                ("missing upgrade", token, 400),
            ]:
                with self.subTest(boundary=name):
                    connection = http.client.HTTPConnection(
                        address.hostname, address.port, timeout=20
                    )
                    try:
                        connection.request("GET", "/events", headers=headers)
                        response = connection.getresponse()
                        self.assertEqual(response.status, status)
                        response.read()
                    finally:
                        connection.close()