"""Async TCP client for communicating with Godot's GolemServer."""

import asyncio
import json
import uuid
import os
import logging

logger = logging.getLogger("golem_client")

DEFAULT_HOST = "localhost"
DEFAULT_PORT = 3571
HEARTBEAT_INTERVAL = 10.0
REQUEST_TIMEOUT = 30.0
MAX_RECONNECT_DELAY = 30.0


class GolemClient:
    """Async TCP client that speaks JSON lines to the Godot GolemServer."""

    def __init__(
        self,
        host: str | None = None,
        port: int | None = None,
        request_timeout: float = REQUEST_TIMEOUT,
    ):
        self.host = host or os.getenv("GOLEM_HOST", DEFAULT_HOST)
        self.port = port or int(os.getenv("GOLEM_PORT", str(DEFAULT_PORT)))
        self.request_timeout = request_timeout

        self._reader: asyncio.StreamReader | None = None
        self._writer: asyncio.StreamWriter | None = None
        self._pending: dict[str, asyncio.Future] = {}
        self._read_task: asyncio.Task | None = None
        self._heartbeat_task: asyncio.Task | None = None
        self._connected = False
        self._reconnect_delay = 1.0

    # ------------------------------------------------------------------
    # Connection lifecycle
    # ------------------------------------------------------------------

    async def connect(self) -> None:
        """Open TCP connection and start background tasks."""
        try:
            self._reader, self._writer = await asyncio.wait_for(
                asyncio.open_connection(self.host, self.port), timeout=5.0
            )
            self._connected = True
            self._reconnect_delay = 1.0
            self._read_task = asyncio.create_task(self._read_loop())
            self._heartbeat_task = asyncio.create_task(self._heartbeat_loop())
            logger.info("Connected to Godot at %s:%s", self.host, self.port)
        except (OSError, asyncio.TimeoutError) as exc:
            self._connected = False
            raise ConnectionError(
                f"Cannot reach Godot at {self.host}:{self.port} — is the plugin active?"
            ) from exc

    async def disconnect(self) -> None:
        """Cleanly shut down the connection."""
        self._connected = False
        if self._heartbeat_task:
            self._heartbeat_task.cancel()
        if self._read_task:
            self._read_task.cancel()
        if self._writer:
            try:
                self._writer.close()
                await self._writer.wait_closed()
            except Exception:
                pass
        self._reader = None
        self._writer = None
        # Fail all pending requests
        for fut in self._pending.values():
            if not fut.done():
                fut.set_exception(ConnectionError("Disconnected"))
        self._pending.clear()

    async def _reconnect(self) -> None:
        """Reconnect with exponential backoff."""
        await self.disconnect()
        while True:
            logger.info(
                "Reconnecting in %.1fs to %s:%s…",
                self._reconnect_delay,
                self.host,
                self.port,
            )
            await asyncio.sleep(self._reconnect_delay)
            self._reconnect_delay = min(
                self._reconnect_delay * 2, MAX_RECONNECT_DELAY
            )
            try:
                await self.connect()
                return
            except ConnectionError:
                continue

    # ------------------------------------------------------------------
    # Read loop
    # ------------------------------------------------------------------

    async def _read_loop(self) -> None:
        """Read JSON lines from Godot and resolve pending futures."""
        assert self._reader is not None
        try:
            while self._connected:
                line = await self._reader.readline()
                if not line:
                    break
                try:
                    msg = json.loads(line)
                except json.JSONDecodeError:
                    logger.warning("Bad JSON from Godot: %s", line)
                    continue

                msg_id = msg.get("id")
                if msg_id and msg_id in self._pending:
                    fut = self._pending.pop(msg_id)
                    if not fut.done():
                        fut.set_result(msg)
        except (asyncio.CancelledError, ConnectionError):
            pass
        finally:
            if self._connected:
                self._connected = False
                asyncio.create_task(self._reconnect())

    # ------------------------------------------------------------------
    # Heartbeat
    # ------------------------------------------------------------------

    async def _heartbeat_loop(self) -> None:
        """Send periodic pings to detect dead connections."""
        try:
            while self._connected:
                await asyncio.sleep(HEARTBEAT_INTERVAL)
                try:
                    await self.call("ping", {}, timeout=5.0)
                except Exception:
                    logger.warning("Heartbeat failed")
                    break
        except asyncio.CancelledError:
            pass

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    async def call(
        self,
        tool: str,
        args: dict | None = None,
        timeout: float | None = None,
    ) -> dict:
        """Send a command to Godot and wait for the response.

        Returns the full response dict: {ok, result, type, ...} or raises.
        """
        if not self._connected or not self._writer:
            await self.connect()

        timeout = timeout or self.request_timeout
        request_id = str(uuid.uuid4())
        message = {"id": request_id, "tool": tool, "args": args or {}}

        loop = asyncio.get_running_loop()
        fut: asyncio.Future = loop.create_future()
        self._pending[request_id] = fut

        try:
            data = json.dumps(message) + "\n"
            self._writer.write(data.encode())
            await self._writer.drain()
        except (OSError, ConnectionError) as exc:
            self._pending.pop(request_id, None)
            raise ConnectionError("Lost connection to Godot") from exc

        try:
            result = await asyncio.wait_for(fut, timeout=timeout)
        except asyncio.TimeoutError:
            self._pending.pop(request_id, None)
            raise TimeoutError(f"Godot did not respond to '{tool}' within {timeout}s")

        if not result.get("ok"):
            raise RuntimeError(result.get("error", "Unknown error from Godot"))

        return result

    async def ensure_connected(self) -> None:
        """Connect if not already connected."""
        if not self._connected:
            await self.connect()
