"""Grex Agent SDK — conecta um agente IA ao escritorio virtual Grex."""

import asyncio
import json
import logging
from typing import Callable, Optional

import websockets

logger = logging.getLogger("grex")


class Message:
    """Mensagem recebida de um player no escritorio."""

    def __init__(self, data: dict):
        self.agent_id: str = data.get("agent_id", "")
        self.user_id: str = data.get("user_id", "")
        self.text: str = data.get("message", "")
        self.time: int = data.get("time", 0)

    def __repr__(self) -> str:
        return f"Message(user={self.user_id}, text={self.text!r})"


class Agent:
    """Agente IA que vive no escritorio Grex.

    Exemplo:
        agent = Agent(api_key="sua-api-key")

        @agent.on_message
        async def handle(msg):
            await agent.set_state("thinking")
            response = "Recebi: " + msg.text
            await agent.set_state("talking")
            await agent.reply(response)
            await agent.set_state("idle")

        agent.run()
    """

    def __init__(
        self,
        api_key: str,
        server_url: str = "ws://65.21.146.176:8080/ws/agent",
        heartbeat_interval: float = 25.0,
    ):
        self.api_key = api_key
        self.server_url = server_url
        self.heartbeat_interval = heartbeat_interval

        self._ws: Optional[websockets.WebSocketClientProtocol] = None
        self._message_handler: Optional[Callable] = None
        self._connected_handler: Optional[Callable] = None
        self._disconnected_handler: Optional[Callable] = None

        self.agent_id: str = ""
        self.office_id: str = ""
        self.name: str = ""
        self.is_connected: bool = False

    def on_message(self, func: Callable) -> Callable:
        """Decorator: registra handler pra quando um player fala com o agente."""
        self._message_handler = func
        return func

    def on_connected(self, func: Callable) -> Callable:
        """Decorator: chamado quando o agente conecta ao servidor."""
        self._connected_handler = func
        return func

    def on_disconnected(self, func: Callable) -> Callable:
        """Decorator: chamado quando o agente desconecta."""
        self._disconnected_handler = func
        return func

    async def set_state(self, state: str) -> None:
        """Muda o estado visual do agente no escritorio.

        Estados: idle, typing, thinking, talking, meeting, walking, away
        """
        await self._send({
            "type": "update_state",
            "payload": {"state": state},
        })

    async def reply(self, message: str) -> None:
        """Envia mensagem de chat que aparece no jogo."""
        await self._send({
            "type": "agent_chat",
            "payload": {"message": message},
        })

    async def _send(self, msg: dict) -> None:
        if self._ws and self.is_connected:
            await self._ws.send(json.dumps(msg))

    async def _connect(self) -> None:
        url = f"{self.server_url}?api_key={self.api_key}"

        while True:
            try:
                logger.info("Connecting to %s...", self.server_url)
                async with websockets.connect(url) as ws:
                    self._ws = ws
                    self.is_connected = True
                    logger.info("Connected!")

                    # Processa mensagens
                    await self._listen(ws)

            except (websockets.ConnectionClosed, ConnectionRefusedError, OSError) as e:
                logger.warning("Disconnected: %s. Reconnecting in 5s...", e)
                self.is_connected = False
                self._ws = None
                if self._disconnected_handler:
                    try:
                        await self._disconnected_handler()
                    except Exception:
                        pass
                await asyncio.sleep(5)

    async def _listen(self, ws) -> None:
        # Heartbeat task
        heartbeat_task = asyncio.create_task(self._heartbeat_loop())

        try:
            async for raw in ws:
                try:
                    msg = json.loads(raw)
                except json.JSONDecodeError:
                    continue

                msg_type = msg.get("type", "")
                payload = msg.get("payload", {})

                if msg_type == "agent_connected":
                    self.agent_id = payload.get("agent_id", "")
                    self.office_id = payload.get("office_id", "")
                    self.name = payload.get("name", "")
                    logger.info("Agent '%s' online in office %s", self.name, self.office_id)
                    if self._connected_handler:
                        await self._connected_handler()

                elif msg_type == "player_message":
                    if self._message_handler:
                        message = Message(payload)
                        try:
                            await self._message_handler(message)
                        except Exception as e:
                            logger.error("Message handler error: %s", e)

                elif msg_type == "heartbeat_ack":
                    pass

                elif msg_type == "error":
                    logger.error("Server error: [%s] %s",
                                 payload.get("code", "?"),
                                 payload.get("message", "?"))

        finally:
            heartbeat_task.cancel()

    async def _heartbeat_loop(self) -> None:
        while True:
            await asyncio.sleep(self.heartbeat_interval)
            await self._send({"type": "heartbeat", "payload": {}})

    def run(self) -> None:
        """Inicia o agente (blocking). Reconecta automaticamente."""
        logging.basicConfig(
            level=logging.INFO,
            format="%(asctime)s [%(name)s] %(message)s",
            datefmt="%H:%M:%S",
        )
        asyncio.run(self._connect())
