from __future__ import annotations

import asyncio
import json
from dataclasses import dataclass

from engine.network.protocol import DEFAULT_PORT, encode_message


@dataclass
class ConnectedClient:
    entity_id: str
    writer: asyncio.StreamWriter


class GameServer:
    """Minimal authoritative-relay TCP server: broadcasts join/move/leave to every other client.

    Wire format is JSON-lines (see engine/network/protocol.py). This is a foundation, not a
    full sync engine - it relays messages verbatim rather than validating/simulating world
    state server-side.
    """

    def __init__(self, host: str = "0.0.0.0", port: int = DEFAULT_PORT):
        self.host = host
        self.port = port
        self.clients: dict[str, ConnectedClient] = {}
        self._server: asyncio.AbstractServer | None = None

    async def start(self) -> None:
        self._server = await asyncio.start_server(self._handle_client, self.host, self.port)
        self.port = self._server.sockets[0].getsockname()[1]

    async def serve_forever(self) -> None:
        assert self._server is not None
        async with self._server:
            await self._server.serve_forever()

    async def stop(self) -> None:
        if self._server is not None:
            self._server.close()
            await self._server.wait_closed()
            self._server = None

    async def _handle_client(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
        entity_id: str | None = None
        try:
            while True:
                line = await reader.readline()
                if not line:
                    break
                message = json.loads(line)
                if message.get("type") == "hello":
                    entity_id = str(message["entity_id"])
                    self.clients[entity_id] = ConnectedClient(entity_id, writer)
                    await self._broadcast(
                        {"type": "join", "entity_id": entity_id, "name": message.get("name", "")}, exclude=entity_id
                    )
                else:
                    await self._broadcast(message, exclude=entity_id)
        finally:
            if entity_id is not None:
                self.clients.pop(entity_id, None)
                await self._broadcast({"type": "leave", "entity_id": entity_id})
            writer.close()

    async def _broadcast(self, message: dict, exclude: str | None = None) -> None:
        data = encode_message(message)
        for client_id, client in list(self.clients.items()):
            if client_id == exclude:
                continue
            try:
                client.writer.write(data)
                await client.writer.drain()
            except ConnectionError:
                self.clients.pop(client_id, None)