"""Tests for WebSocket hub, framing, and client management."""

from __future__ import annotations

import io
import json
import struct
from unittest.mock import MagicMock

from orxaq.enterprise.events import EventBus
from orxaq.enterprise.state import StateEmitter
from orxaq.experiences.websocket import (
    OP_CLOSE,
    OP_PING,
    OP_TEXT,
    WebSocketClient,
    WebSocketHub,
    _handle_client_message,
    compute_accept_key,
    decode_frame,
    encode_frame,
    handle_websocket_upgrade,
    is_websocket_upgrade,
)


class TestFrameEncoding:
    """Test WebSocket frame encode/decode."""

    def test_encode_small_frame(self):
        data = b"hello"
        frame = encode_frame(data)
        assert frame[0] == 0x81  # FIN + TEXT
        assert frame[1] == 5
        assert frame[2:] == data

    def test_encode_medium_frame(self):
        data = b"x" * 200
        frame = encode_frame(data)
        assert frame[0] == 0x81
        assert frame[1] == 126
        length = struct.unpack("!H", frame[2:4])[0]
        assert length == 200
        assert frame[4:] == data

    def test_encode_large_frame(self):
        data = b"x" * 70000
        frame = encode_frame(data)
        assert frame[0] == 0x81
        assert frame[1] == 127
        length = struct.unpack("!Q", frame[2:10])[0]
        assert length == 70000

    def test_decode_unmasked_frame(self):
        data = b"hello"
        frame = encode_frame(data)
        opcode, payload, consumed = decode_frame(frame)
        assert opcode == OP_TEXT
        assert payload == data
        assert consumed == len(frame)

    def test_decode_masked_frame(self):
        payload = b"test"
        mask = b"\x12\x34\x56\x78"
        masked = bytearray(payload)
        for i in range(len(masked)):
            masked[i] ^= mask[i % 4]

        frame = bytearray()
        frame.append(0x81)
        frame.append(0x80 | len(payload))
        frame.extend(mask)
        frame.extend(masked)

        opcode, decoded, consumed = decode_frame(bytes(frame))
        assert opcode == OP_TEXT
        assert decoded == payload
        assert consumed == len(frame)

    def test_decode_incomplete_frame(self):
        opcode, payload, consumed = decode_frame(b"\x81")
        assert opcode == -1
        assert consumed == 0

    def test_decode_close_frame(self):
        frame = encode_frame(b"", opcode=OP_CLOSE)
        opcode, _, consumed = decode_frame(frame)
        assert opcode == OP_CLOSE

    def test_roundtrip(self):
        for msg in [b"", b"a", b"x" * 125, b"y" * 200, b"z" * 70000]:
            frame = encode_frame(msg)
            opcode, payload, _ = decode_frame(frame)
            assert opcode == OP_TEXT
            assert payload == msg


class TestAcceptKey:
    def test_deterministic(self):
        key = "dGhlIHNhbXBsZSBub25jZQ=="
        result = compute_accept_key(key)
        # Verify deterministic: same key always produces same result.
        assert compute_accept_key(key) == result
        assert len(result) > 0

    def test_different_keys_differ(self):
        assert compute_accept_key("key1==") != compute_accept_key("key2==")


class TestIsWebSocketUpgrade:
    def test_valid_upgrade(self):
        headers = {"Connection": "Upgrade", "Upgrade": "websocket"}
        assert is_websocket_upgrade(headers)

    def test_case_insensitive(self):
        headers = {"Connection": "upgrade", "Upgrade": "WebSocket"}
        assert is_websocket_upgrade(headers)

    def test_missing_upgrade(self):
        headers = {"Connection": "keep-alive"}
        assert not is_websocket_upgrade(headers)

    def test_empty_headers(self):
        assert not is_websocket_upgrade({})


class TestWebSocketClient:
    def test_send_success(self):
        wfile = io.BytesIO()
        client = WebSocketClient(wfile, "test-1")
        assert client.send("hello")
        assert client.connected

    def test_send_json(self):
        wfile = io.BytesIO()
        client = WebSocketClient(wfile, "test-2")
        assert client.send_json({"type": "test", "value": 42})
        wfile.seek(0)
        frame = wfile.read()
        _, payload, _ = decode_frame(frame)
        data = json.loads(payload.decode("utf-8"))
        assert data["type"] == "test"
        assert data["value"] == 42

    def test_send_after_disconnect(self):
        wfile = io.BytesIO()
        client = WebSocketClient(wfile, "test-3")
        client.connected = False
        assert not client.send("hello")

    def test_send_broken_pipe(self):
        wfile = MagicMock()
        wfile.write.side_effect = BrokenPipeError
        client = WebSocketClient(wfile, "test-4")
        assert not client.send("hello")
        assert not client.connected

    def test_close(self):
        wfile = io.BytesIO()
        client = WebSocketClient(wfile, "test-5")
        client.close()
        assert not client.connected

    def test_close_already_disconnected(self):
        wfile = io.BytesIO()
        client = WebSocketClient(wfile, "test-6")
        client.connected = False
        client.close()  # No-op, should not raise.


class TestWebSocketHub:
    def setup_method(self):
        self.bus = EventBus()
        self.emitter = StateEmitter(version="0.2.0")
        self.hub = WebSocketHub(event_bus=self.bus, state_emitter=self.emitter)

    def teardown_method(self):
        self.hub.shutdown()

    def test_add_client_sends_snapshot(self):
        wfile = io.BytesIO()
        client = self.hub.add_client(wfile)
        assert client.connected
        assert self.hub.client_count() == 1

        wfile.seek(0)
        frame = wfile.read()
        _, payload, _ = decode_frame(frame)
        data = json.loads(payload.decode("utf-8"))
        assert data["type"] == "state_snapshot"
        assert data["data"]["version"] == "0.2.0"

    def test_remove_client(self):
        wfile = io.BytesIO()
        client = self.hub.add_client(wfile)
        assert self.hub.client_count() == 1
        self.hub.remove_client(client)
        assert self.hub.client_count() == 0

    def test_broadcast(self):
        for _ in range(3):
            self.hub.add_client(io.BytesIO())
        count = self.hub.broadcast({"type": "test", "value": 1})
        assert count == 3

    def test_broadcast_removes_dead_clients(self):
        self.hub.add_client(io.BytesIO())

        bad_wfile = MagicMock()
        bad_wfile.write.side_effect = BrokenPipeError
        bad_wfile.flush.side_effect = BrokenPipeError
        self.hub.add_client(bad_wfile)

        assert self.hub.client_count() == 2
        self.hub.broadcast({"type": "test"})
        assert self.hub.client_count() == 1

    def test_event_bus_forwarding(self):
        wfile = io.BytesIO()
        self.hub.add_client(wfile)

        # Clear the initial snapshot.
        wfile.seek(0)
        wfile.truncate()

        self.bus.publish("dag.update", {"node": "gdp_growth"}, source="test")

        wfile.seek(0)
        frame = wfile.read()
        assert len(frame) > 0
        _, payload, _ = decode_frame(frame)
        data = json.loads(payload.decode("utf-8"))
        assert data["type"] == "event"
        assert data["data"]["topic"] == "dag.update"

    def test_shutdown_closes_all(self):
        clients = [self.hub.add_client(io.BytesIO()) for _ in range(3)]
        self.hub.shutdown()
        assert self.hub.client_count() == 0
        for c in clients:
            assert not c.connected

    def test_properties(self):
        assert self.hub.event_bus is self.bus
        assert self.hub.state_emitter is self.emitter


class TestHandleClientMessage:
    def test_get_state(self):
        bus = EventBus()
        emitter = StateEmitter(version="0.2.0")
        hub = WebSocketHub(event_bus=bus, state_emitter=emitter)
        wfile = io.BytesIO()
        client = WebSocketClient(wfile, "test")

        wfile.seek(0)
        wfile.truncate()

        payload = json.dumps({"type": "get_state"}).encode()
        _handle_client_message(client, hub, payload)

        wfile.seek(0)
        frame = wfile.read()
        _, resp, _ = decode_frame(frame)
        data = json.loads(resp.decode("utf-8"))
        assert data["type"] == "state_snapshot"
        hub.shutdown()

    def test_get_history(self):
        bus = EventBus()
        emitter = StateEmitter()
        hub = WebSocketHub(event_bus=bus, state_emitter=emitter)

        bus.publish("test.event", {"x": 1})
        bus.publish("test.event", {"x": 2})

        wfile = io.BytesIO()
        client = WebSocketClient(wfile, "test")
        wfile.seek(0)
        wfile.truncate()

        payload = json.dumps({"type": "get_history", "topic": "test.event"}).encode()
        _handle_client_message(client, hub, payload)

        wfile.seek(0)
        frame = wfile.read()
        _, resp, _ = decode_frame(frame)
        data = json.loads(resp.decode("utf-8"))
        assert data["type"] == "event_history"
        assert len(data["data"]) == 2
        hub.shutdown()

    def test_invalid_json_ignored(self):
        bus = EventBus()
        emitter = StateEmitter()
        hub = WebSocketHub(event_bus=bus, state_emitter=emitter)
        wfile = io.BytesIO()
        client = WebSocketClient(wfile, "test")
        _handle_client_message(client, hub, b"not json{{{")
        hub.shutdown()


class TestDecodeFrameEdgeCases:
    """Edge cases for decode_frame: medium/large frames, incomplete buffers."""

    def test_medium_length_incomplete_header(self):
        """126-length frame with insufficient header bytes returns -1."""
        # Length byte = 126 means next 2 bytes encode actual length.
        frame = bytearray([0x81, 126])  # Only 2 bytes, need 4 minimum.
        opcode, payload, consumed = decode_frame(bytes(frame))
        assert opcode == -1
        assert consumed == 0

    def test_large_length_incomplete_header(self):
        """127-length frame with insufficient header bytes returns -1."""
        # Length byte = 127 means next 8 bytes encode actual length.
        frame = bytearray([0x81, 127, 0, 0, 0])  # 5 bytes, need 10.
        opcode, payload, consumed = decode_frame(bytes(frame))
        assert opcode == -1
        assert consumed == 0

    def test_masked_frame_incomplete_payload(self):
        """Masked frame with insufficient payload data returns -1."""
        frame = bytearray([0x81, 0x84])  # masked, length 4
        frame.extend(b"\x00\x00\x00\x00")  # 4-byte mask
        frame.extend(b"\x01\x02")  # Only 2 payload bytes, need 4.
        opcode, payload, consumed = decode_frame(bytes(frame))
        assert opcode == -1
        assert consumed == 0

    def test_unmasked_frame_incomplete_payload(self):
        """Unmasked frame with insufficient payload data returns -1."""
        frame = bytearray([0x81, 10])  # 10 bytes expected
        frame.extend(b"short")  # Only 5 bytes.
        opcode, payload, consumed = decode_frame(bytes(frame))
        assert opcode == -1
        assert consumed == 0

    def test_medium_length_decode_roundtrip(self):
        """Encode/decode a 200-byte frame via 126-length path."""
        data = b"m" * 200
        frame = encode_frame(data)
        opcode, payload, consumed = decode_frame(frame)
        assert opcode == OP_TEXT
        assert payload == data
        assert consumed == len(frame)

    def test_large_length_decode_roundtrip(self):
        """Encode/decode a 70000-byte frame via 127-length path."""
        data = b"L" * 70000
        frame = encode_frame(data)
        opcode, payload, consumed = decode_frame(frame)
        assert opcode == OP_TEXT
        assert payload == data
        assert consumed == len(frame)


class TestClientCloseError:
    """Cover the close() BrokenPipeError path."""

    def test_close_broken_pipe(self):
        wfile = MagicMock()
        wfile.write.side_effect = BrokenPipeError
        client = WebSocketClient(wfile, "err-close")
        client.close()
        assert not client.connected

    def test_close_os_error(self):
        wfile = MagicMock()
        wfile.write.side_effect = OSError("connection reset")
        client = WebSocketClient(wfile, "err-close-2")
        client.close()
        assert not client.connected


class TestHandleWebSocketUpgrade:
    """Cover handle_websocket_upgrade read loop paths."""

    def _make_handler(self, key: str = "dGhlIHNhbXBsZSBub25jZQ=="):
        handler = MagicMock()
        handler.headers = {"Sec-WebSocket-Key": key}
        handler.wfile = io.BytesIO()
        return handler

    def test_missing_key_sends_error(self):
        """Line 272-273: missing Sec-WebSocket-Key."""
        handler = MagicMock()
        handler.headers = {"Sec-WebSocket-Key": ""}
        bus = EventBus()
        emitter = StateEmitter()
        hub = WebSocketHub(event_bus=bus, state_emitter=emitter)
        handle_websocket_upgrade(handler, hub)
        handler.send_error.assert_called_once_with(400, "Missing Sec-WebSocket-Key")
        hub.shutdown()

    def test_read_loop_empty_chunk_breaks(self):
        """Line 295-296: empty read returns b'' — breaks loop."""
        handler = self._make_handler()
        handler.rfile = MagicMock()
        handler.rfile.read.return_value = b""
        handler.rfile.read1 = MagicMock(return_value=b"")

        bus = EventBus()
        emitter = StateEmitter()
        hub = WebSocketHub(event_bus=bus, state_emitter=emitter)
        handle_websocket_upgrade(handler, hub)
        assert hub.client_count() == 0  # Client removed in finally
        hub.shutdown()

    def test_read_loop_close_frame(self):
        """Line 305-307: OP_CLOSE frame closes client."""
        close_frame = encode_frame(b"", OP_CLOSE)
        handler = self._make_handler()
        handler.rfile = MagicMock()
        # First read returns close frame, then b""
        handler.rfile.read1 = MagicMock(side_effect=[close_frame, b""])

        bus = EventBus()
        emitter = StateEmitter()
        hub = WebSocketHub(event_bus=bus, state_emitter=emitter)
        handle_websocket_upgrade(handler, hub)
        assert hub.client_count() == 0
        hub.shutdown()

    def test_read_loop_ping_frame(self):
        """Line 308-312: OP_PING gets pong response."""
        ping_frame = encode_frame(b"ping", OP_PING)
        handler = self._make_handler()
        handler.wfile = MagicMock()
        handler.rfile = MagicMock()
        # First read returns ping, second returns empty to exit.
        handler.rfile.read1 = MagicMock(side_effect=[ping_frame, b""])

        bus = EventBus()
        emitter = StateEmitter()
        hub = WebSocketHub(event_bus=bus, state_emitter=emitter)
        handle_websocket_upgrade(handler, hub)
        handler.wfile.write.assert_called()
        handler.wfile.flush.assert_called()
        hub.shutdown()

    def test_read_loop_ping_broken_pipe(self):
        """Line 313-314: OP_PING pong write fails with BrokenPipeError."""
        ping_frame = encode_frame(b"ping", OP_PING)
        handler = self._make_handler()
        handler.wfile = MagicMock()
        handler.wfile.write.side_effect = BrokenPipeError
        handler.rfile = MagicMock()
        handler.rfile.read1 = MagicMock(side_effect=[ping_frame, b""])

        bus = EventBus()
        emitter = StateEmitter()
        hub = WebSocketHub(event_bus=bus, state_emitter=emitter)
        handle_websocket_upgrade(handler, hub)
        hub.shutdown()

    def test_read_loop_text_frame(self):
        """Line 315-316: OP_TEXT dispatches to _handle_client_message."""
        msg = json.dumps({"type": "get_state"}).encode()
        text_frame = encode_frame(msg, OP_TEXT)
        handler = self._make_handler()
        handler.wfile = MagicMock()
        handler.rfile = MagicMock()
        handler.rfile.read1 = MagicMock(side_effect=[text_frame, b""])

        bus = EventBus()
        emitter = StateEmitter()
        hub = WebSocketHub(event_bus=bus, state_emitter=emitter)
        handle_websocket_upgrade(handler, hub)
        hub.shutdown()

    def test_read_loop_oserror_breaks(self):
        """Line 293-294: OSError during read breaks the loop."""
        handler = self._make_handler()
        handler.rfile = MagicMock()
        handler.rfile.read1 = MagicMock(side_effect=OSError("connection reset"))

        bus = EventBus()
        emitter = StateEmitter()
        hub = WebSocketHub(event_bus=bus, state_emitter=emitter)
        handle_websocket_upgrade(handler, hub)
        assert hub.client_count() == 0
        hub.shutdown()

    def test_read_loop_uses_read_fallback(self):
        """Line 292: uses rfile.read when read1 not available."""
        handler = self._make_handler()
        handler.rfile = MagicMock(spec=["read"])  # No read1 attr.
        handler.rfile.read.return_value = b""

        bus = EventBus()
        emitter = StateEmitter()
        hub = WebSocketHub(event_bus=bus, state_emitter=emitter)
        handle_websocket_upgrade(handler, hub)
        handler.rfile.read.assert_called()
        hub.shutdown()


class TestHeartbeatWithClients:
    """Cover heartbeat loop when clients exist (lines 233-236)."""

    def test_heartbeat_broadcasts_to_clients(self):
        bus = EventBus()
        emitter = StateEmitter()
        hub = WebSocketHub(event_bus=bus, state_emitter=emitter)
        hub._heartbeat_interval = 0.05  # Speed up for test.
        wfile = io.BytesIO()
        hub.add_client(wfile)

        # Let the heartbeat fire at least once.
        import time

        time.sleep(0.15)

        wfile.seek(0)
        frame_data = wfile.read()
        # Should have received at least one state_snapshot from heartbeat.
        assert len(frame_data) > 0
        hub.shutdown()
