"""AnycastAgent — async Python client for the Anycast Agents rendezvous server."""

from __future__ import annotations

import asyncio
import json
import logging
import os
import secrets
import base64
from typing import Any, Callable, Dict, List, Optional

import websockets
import websockets.client

from .types import AgentInfo, AgentOptions, ConnectionState
from .exceptions import (
    AnycastError,
    AnycastAuthError,
    AnycastConnectionError,
    AnycastTimeoutError,
)

logger = logging.getLogger("anycast_agent")


class AnycastAgent:
    """Async client for the Anycast Agents rendezvous server.

    Example::

        agent = AnycastAgent(
            rendezvous_url="wss://agents.anycast.com/rendezvous",
            token="agt_...",
        )

        @agent.on_connect
        async def on_connect(info):
            print(f"Connected as {info.agent_id}")

        @agent.on_peer
        async def on_peer(agent_id, status):
            print(f"Peer {agent_id} is {status}")

        await agent.connect()
        await agent.run_forever()
    """

    def __init__(
        self,
        rendezvous_url: str,
        token: str,
        *,
        machine_id: Optional[str] = None,
        public_key: Optional[str] = None,
        version: Optional[str] = None,
        reconnect: bool = True,
        reconnect_interval: float = 3.0,
        timeout: float = 30.0,
        heartbeat_interval: float = 25.0,
    ):
        self._opts = AgentOptions(
            rendezvous_url=rendezvous_url,
            token=token,
            machine_id=machine_id or secrets.token_hex(16),
            public_key=public_key or base64.b64encode(secrets.token_bytes(32)).decode(),
            version=version,
            reconnect=reconnect,
            reconnect_interval=reconnect_interval,
            timeout=timeout,
            heartbeat_interval=heartbeat_interval,
        )

        self._ws: Optional[websockets.client.WebSocketClientProtocol] = None
        self._state = ConnectionState.DISCONNECTED
        self._heartbeat_task: Optional[asyncio.Task[None]] = None
        self._recv_task: Optional[asyncio.Task[None]] = None
        self._reconnect_attempt = 0
        self._intentional_close = False

        # Public state
        self.agent_id: Optional[str] = None
        self.tenant_id: Optional[str] = None
        self.agent_name: Optional[str] = None
        self.pop_iata: Optional[str] = None

        # Event handlers
        self._handlers: Dict[str, List[Callable[..., Any]]] = {}

    # ------------------------------------------------------------------
    # Properties
    # ------------------------------------------------------------------

    @property
    def connection_state(self) -> ConnectionState:
        return self._state

    @property
    def connected(self) -> bool:
        return self._state == ConnectionState.CONNECTED

    # ------------------------------------------------------------------
    # Event registration (decorator style)
    # ------------------------------------------------------------------

    def on(self, event: str, handler: Callable[..., Any]) -> None:
        """Register an event handler."""
        self._handlers.setdefault(event, []).append(handler)

    def on_connect(self, fn: Callable[..., Any]) -> Callable[..., Any]:
        """Decorator: register a connect handler."""
        self.on("connect", fn)
        return fn

    def on_disconnect(self, fn: Callable[..., Any]) -> Callable[..., Any]:
        """Decorator: register a disconnect handler."""
        self.on("disconnect", fn)
        return fn

    def on_peer(self, fn: Callable[..., Any]) -> Callable[..., Any]:
        """Decorator: register a peer status handler."""
        self.on("peer", fn)
        return fn

    def on_message(self, fn: Callable[..., Any]) -> Callable[..., Any]:
        """Decorator: register a message handler."""
        self.on("message", fn)
        return fn

    def on_error(self, fn: Callable[..., Any]) -> Callable[..., Any]:
        """Decorator: register an error handler."""
        self.on("error", fn)
        return fn

    # ------------------------------------------------------------------
    # Connection
    # ------------------------------------------------------------------

    async def connect(self) -> AgentInfo:
        """Connect to the rendezvous server and register.

        Returns:
            AgentInfo with agent_id, tenant_id, and name.

        Raises:
            AnycastTimeoutError: if registration times out
            AnycastAuthError: if token is invalid
            AnycastConnectionError: if WebSocket connection fails
        """
        if self._state in (ConnectionState.CONNECTED, ConnectionState.CONNECTING):
            if self.agent_id and self.tenant_id and self.agent_name:
                return AgentInfo(self.agent_id, self.tenant_id, self.agent_name)
            raise AnycastConnectionError("Already connecting")

        self._intentional_close = False
        self._set_state(ConnectionState.CONNECTING)

        try:
            self._ws = await asyncio.wait_for(
                websockets.connect(self._opts.rendezvous_url),
                timeout=self._opts.timeout,
            )
        except asyncio.TimeoutError:
            self._set_state(ConnectionState.DISCONNECTED)
            raise AnycastTimeoutError("WebSocket connection timeout")
        except Exception as exc:
            self._set_state(ConnectionState.DISCONNECTED)
            raise AnycastConnectionError(str(exc)) from exc

        # Send register message
        await self._send({
            "type": "register",
            "token": self._opts.token,
            "machineId": self._opts.machine_id,
            "publicKey": self._opts.public_key,
            **({"version": self._opts.version} if self._opts.version else {}),
        })

        # Wait for registered or error
        try:
            raw = await asyncio.wait_for(
                self._ws.recv(),
                timeout=self._opts.timeout,
            )
        except asyncio.TimeoutError:
            await self._cleanup()
            raise AnycastTimeoutError("Registration timeout")

        msg = json.loads(raw)

        if msg.get("type") == "error" and msg.get("fatal"):
            await self._cleanup()
            raise AnycastAuthError(msg.get("message", "Authentication failed"))

        if msg.get("type") != "registered":
            await self._cleanup()
            raise AnycastConnectionError(f"Unexpected message: {msg.get('type')}")

        self.agent_id = msg["agentId"]
        self.tenant_id = msg["tenantId"]
        self.agent_name = msg["name"]
        self.pop_iata = msg.get("popIata")
        self._reconnect_attempt = 0
        self._set_state(ConnectionState.CONNECTED)

        # Start background tasks
        self._heartbeat_task = asyncio.create_task(self._heartbeat_loop())
        self._recv_task = asyncio.create_task(self._recv_loop())

        info = AgentInfo(self.agent_id, self.tenant_id, self.agent_name, self.pop_iata)
        await self._emit("connect", info)
        return info

    async def disconnect(self) -> None:
        """Disconnect from the rendezvous server."""
        self._intentional_close = True
        await self._cleanup()
        self.agent_id = None
        self.tenant_id = None
        self.agent_name = None

    async def run_forever(self) -> None:
        """Block until disconnected. Handles reconnection automatically."""
        while not self._intentional_close:
            if self._recv_task:
                try:
                    await self._recv_task
                except asyncio.CancelledError:
                    break
            if self._intentional_close:
                break
            if self._opts.reconnect:
                await self._reconnect()
            else:
                break

    # ------------------------------------------------------------------
    # Actions
    # ------------------------------------------------------------------

    def request_connection(self, target_agent_id: str) -> str:
        """Request a P2P connection to a peer agent.

        Returns:
            Session ID for this connection.
        """
        self._assert_connected()
        session_id = secrets.token_hex(16)
        asyncio.create_task(self._send({
            "type": "connect_request",
            "targetAgentId": target_agent_id,
            "sessionId": session_id,
        }))
        return session_id

    async def send_signal(
        self,
        session_id: str,
        target_agent_id: str,
        signal: Dict[str, str],
    ) -> None:
        """Send a WebRTC signal to a peer via the rendezvous server."""
        self._assert_connected()
        await self._send({
            "type": "signal",
            "sessionId": session_id,
            "targetAgentId": target_agent_id,
            "signal": signal,
        })

    async def disconnect_session(self, session_id: str) -> None:
        """Close a specific session."""
        self._assert_connected()
        await self._send({"type": "disconnect", "sessionId": session_id})

    # ------------------------------------------------------------------
    # Internal
    # ------------------------------------------------------------------

    async def _send(self, msg: Dict[str, Any]) -> None:
        if self._ws:
            try:
                await self._ws.send(json.dumps(msg))
            except websockets.exceptions.ConnectionClosed:
                pass

    async def _heartbeat_loop(self) -> None:
        try:
            while self._state == ConnectionState.CONNECTED:
                await asyncio.sleep(self._opts.heartbeat_interval)
                await self._send({"type": "heartbeat"})
        except asyncio.CancelledError:
            pass

    async def _recv_loop(self) -> None:
        try:
            assert self._ws is not None
            async for raw in self._ws:
                try:
                    msg = json.loads(raw)
                except json.JSONDecodeError:
                    continue
                await self._handle_message(msg)
        except websockets.exceptions.ConnectionClosed as exc:
            was_connected = self._state == ConnectionState.CONNECTED
            self._set_state(ConnectionState.DISCONNECTED)
            if was_connected:
                await self._emit("disconnect", str(exc))
        except asyncio.CancelledError:
            pass

    async def _handle_message(self, msg: Dict[str, Any]) -> None:
        msg_type = msg.get("type")

        if msg_type == "heartbeat_ack":
            return

        if msg_type == "agent_status":
            await self._emit("peer", msg["agentId"], msg["status"])

        elif msg_type == "connect_incoming":
            await self._emit(
                "connect_incoming",
                msg["sessionId"],
                msg["fromAgentId"],
                msg["fromAgentName"],
            )

        elif msg_type == "connect_accepted":
            await self._emit(
                "connect_accepted",
                msg["sessionId"],
                msg["agentId"],
                msg["publicKey"],
            )

        elif msg_type == "connect_rejected":
            await self._emit("connect_rejected", msg["sessionId"], msg["reason"])

        elif msg_type == "signal_forward":
            await self._emit(
                "signal",
                msg["sessionId"],
                msg["fromAgentId"],
                msg["signal"],
            )

        elif msg_type == "relay_assigned":
            await self._emit(
                "relay_assigned",
                msg["sessionId"],
                msg["relayHost"],
                msg["relayPort"],
                msg["relayToken"],
            )

        elif msg_type == "error":
            err = AnycastError(
                msg.get("message", "Unknown error"),
                msg.get("code", "UNKNOWN"),
                msg.get("fatal", False),
            )
            await self._emit("error", err)
            if err.fatal:
                await self._cleanup()

    async def _emit(self, event: str, *args: Any) -> None:
        for handler in self._handlers.get(event, []):
            try:
                result = handler(*args)
                if asyncio.iscoroutine(result):
                    await result
            except Exception:
                logger.exception("Error in %s handler", event)

    async def _reconnect(self) -> None:
        self._reconnect_attempt += 1
        delay = min(
            self._opts.reconnect_interval * (2 ** (self._reconnect_attempt - 1)),
            30.0,
        )
        self._set_state(ConnectionState.RECONNECTING)
        await self._emit("reconnecting", self._reconnect_attempt)
        logger.info("Reconnecting in %.1fs (attempt %d)", delay, self._reconnect_attempt)
        await asyncio.sleep(delay)
        try:
            await self.connect()
        except AnycastError:
            pass  # connect() failure will trigger another reconnect via run_forever

    async def _cleanup(self) -> None:
        if self._heartbeat_task and not self._heartbeat_task.done():
            self._heartbeat_task.cancel()
            self._heartbeat_task = None
        if self._recv_task and not self._recv_task.done():
            self._recv_task.cancel()
            self._recv_task = None
        if self._ws:
            try:
                await self._ws.close(1000, "Client disconnect")
            except Exception:
                pass
            self._ws = None
        self._set_state(ConnectionState.DISCONNECTED)

    def _set_state(self, state: ConnectionState) -> None:
        if self._state != state:
            self._state = state
            # Fire state_change synchronously inline (no await needed for sync handlers)
            for handler in self._handlers.get("state_change", []):
                try:
                    handler(state)
                except Exception:
                    logger.exception("Error in state_change handler")

    def _assert_connected(self) -> None:
        if self._state != ConnectionState.CONNECTED or not self._ws:
            raise AnycastError("Not connected", "NOT_CONNECTED")
