from __future__ import annotations

import json
import socket
import time

from server_core import PROTOCOL_VERSION, FalloutTogetherServer


_RECV_BUFFERS: dict[socket.socket, bytes] = {}


def recv_packet(sock: socket.socket, timeout: float = 2.0) -> dict:
    sock.settimeout(timeout)
    data = _RECV_BUFFERS.get(sock, b"")
    while b"\n" not in data:
        chunk = sock.recv(4096)
        if not chunk:
            raise ConnectionError("socket closed before a complete packet was received")
        data += chunk
    line, remainder = data.split(b"\n", 1)
    _RECV_BUFFERS[sock] = remainder
    return json.loads(line.decode("utf-8"))


def recv_until(sock: socket.socket, predicate, timeout: float = 2.0) -> dict:
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        try:
            packet = recv_packet(sock, timeout=max(0.01, deadline - time.monotonic()))
        except socket.timeout:
            break
        if predicate(packet):
            return packet
    raise AssertionError("did not receive matching packet")


def send_packet(sock: socket.socket, packet: dict) -> None:
    sock.sendall(json.dumps(packet, separators=(",", ":")).encode() + b"\n")


def connect_v2(server: FalloutTogetherServer) -> tuple[socket.socket, int]:
    sock = socket.create_connection(("127.0.0.1", server.port), timeout=2.0)
    welcome = recv_until(sock, lambda packet: packet.get("type") == "welcome")
    send_packet(sock, {"type": "hello", "protocolVersion": PROTOCOL_VERSION})
    ready = recv_until(sock, lambda packet: packet.get("type") == "sessionReady")
    assert ready["playerId"] == welcome["playerId"]
    return sock, int(welcome["playerId"])


def transform(cell: str, x: float, y: float = 0.0) -> dict:
    return {
        "type": "transform",
        "x": x,
        "y": y,
        "z": 0.0,
        "angleZ": 0.0,
        "cellId": cell,
        "worldspaceId": "0000003C",
        "movementType": "normal",
    }


def start_server() -> FalloutTogetherServer:
    server = FalloutTogetherServer(host="127.0.0.1", port=0, max_players=24)
    server.start()
    deadline = time.monotonic() + 2.0
    while not server.is_running() and time.monotonic() < deadline:
        time.sleep(0.01)
    return server


def test_sixteen_clients_do_not_cross_broadcast_distant_cell_transforms():
    server = start_server()
    clients: list[tuple[socket.socket, int]] = []
    try:
        clients = [connect_v2(server) for _ in range(16)]
        assert server.get_stats()["connectedClients"] == 16

        near = clients[:8]
        far = clients[8:]
        for index, (sock, _player_id) in enumerate(near):
            send_packet(sock, transform("00000010", float(index * 32)))
        for index, (sock, _player_id) in enumerate(far):
            send_packet(sock, transform("00000020", 30000.0 + float(index * 32)))

        deadline = time.monotonic() + 2.0
        while time.monotonic() < deadline:
            sessions = [server._find_client_by_player_id(player_id) for _sock, player_id in clients]
            if all(session is not None and session.last_transform is not None for session in sessions):
                break
            time.sleep(0.01)
        assert all(
            session is not None and session.last_transform is not None
            for session in (server._find_client_by_player_id(player_id) for _sock, player_id in clients)
        )

        sender_sock, sender_id = near[0]
        near_peer_sock, _near_peer_id = near[1]
        far_peer_sock, _far_peer_id = far[0]
        marker_x = 128.0
        send_packet(sender_sock, transform("00000010", marker_x))

        relayed = recv_until(
            near_peer_sock,
            lambda packet: packet.get("type") == "transform"
            and packet.get("playerId") == sender_id
            and packet.get("x") == marker_x,
        )
        assert relayed["cellId"] == "00000010"

        try:
            recv_until(
                far_peer_sock,
                lambda packet: packet.get("type") == "transform"
                and packet.get("playerId") == sender_id
                and packet.get("x") == marker_x,
                timeout=0.4,
            )
        except AssertionError:
            pass
        else:
            raise AssertionError("distant cell received a transform that should have been interest-filtered")

        stats = server.get_stats()
        assert stats["connectedClients"] == 16
        assert stats["transformPacketsInterestFiltered"] >= len(far)
    finally:
        for sock, _player_id in clients:
            _RECV_BUFFERS.pop(sock, None)
            try:
                sock.close()
            except OSError:
                pass
        server.stop()