"""UDP request/response link to the claude-co-engineer mod running inside the game."""

import asyncio
import itertools
import json
import socket
from typing import Any

from mcp.server.mcpserver.exceptions import ToolError


class GameError(ToolError):
    """The game couldn't answer a query. A ToolError so the message reaches the model."""


class _Protocol(asyncio.DatagramProtocol):
    def __init__(self, link: "GameLink"):
        self.link = link

    def datagram_received(self, data: bytes, addr: Any) -> None:
        # The game replies from its --enable-lua-udp port; ignore anything else on localhost.
        if addr[:2] == self.link.peer:
            self.link._on_packet(data)


class GameLink:
    """Sends JSON requests to the game's --enable-lua-udp port and reassembles chunked replies.

    The mod replies to whatever port a request came from, so we bind an ephemeral port.
    """

    def __init__(self, game_port: int, host: str = "127.0.0.1", timeout: float = 5.0):
        self.game_addr = (host, game_port)
        self.host = host
        self.timeout = timeout
        self._ids = itertools.count(1)
        self._pending: dict[int, asyncio.Future] = {}
        self._parts: dict[int, dict[int, bytes]] = {}
        self._transport: asyncio.DatagramTransport | None = None

    async def _ensure_open(self) -> asyncio.DatagramTransport:
        if self._transport is None or self._transport.is_closing():
            loop = asyncio.get_running_loop()
            # The address replies come from, resolved so a hostname like "localhost" matches.
            self.peer = (socket.gethostbyname(self.host), self.game_addr[1])
            sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            # Large replies arrive as a burst of ~32KB packets; give the kernel room to queue them.
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 8 * 1024 * 1024)
            sock.bind((self.host, 0))
            self._transport, _ = await loop.create_datagram_endpoint(lambda: _Protocol(self), sock=sock)
        return self._transport

    def _on_packet(self, data: bytes) -> None:
        try:
            raw_id, raw_seq, raw_total, fragment = data.split(b"|", 3)
            request_id, seq, total = int(raw_id), int(raw_seq), int(raw_total)
        except ValueError:
            return
        future = self._pending.get(request_id)
        if future is None or future.done():
            return
        parts = self._parts.setdefault(request_id, {})
        parts[seq] = fragment
        if len(parts) < total:
            return
        del self._parts[request_id]
        try:
            future.set_result(json.loads(b"".join(parts[i] for i in range(1, total + 1))))
        except (KeyError, ValueError) as e:
            future.set_exception(GameError(f"malformed reply from game: {e}"))

    async def request(self, op: str, **args: Any) -> Any:
        transport = await self._ensure_open()
        request_id = next(self._ids)
        future = asyncio.get_running_loop().create_future()
        self._pending[request_id] = future
        payload = {"id": request_id, "op": op, "args": {k: v for k, v in args.items() if v is not None}}
        try:
            transport.sendto(json.dumps(payload).encode(), self.game_addr)
            response = await asyncio.wait_for(future, self.timeout)
        except TimeoutError:
            raise GameError(
                f"No reply from Factorio within {self.timeout:g}s on UDP port {self.game_addr[1]}. "
                "Check that the game was launched with --enable-lua-udp on that port, the claude-co-engineer "
                "mod is enabled, a save is loaded, and the game isn't paused (the Esc menu pauses singleplayer)."
            ) from None
        finally:
            self._pending.pop(request_id, None)
            self._parts.pop(request_id, None)
        if not response.get("ok"):
            raise GameError(response.get("error", "unknown error"))
        return response.get("result")