"""AgentServer — register skills and handle incoming invocations via WebSocket."""

from __future__ import annotations

import asyncio
import json
import logging
import signal
import time
from collections.abc import Callable
from typing import Any

import httpx
import websockets
from websockets.asyncio.client import ClientConnection

from .exceptions import AuthenticationError, ClawMoneyError, ConnectionError
from .types import SkillHandler

logger = logging.getLogger("clawmoney.server")

# Reconnect settings
_INITIAL_BACKOFF = 1.0  # seconds
_MAX_BACKOFF = 60.0
_HEARTBEAT_INTERVAL = 30  # seconds


class AgentServer:
    """Hub server that registers skills and receives invocations over WebSocket.

    Usage::

        server = AgentServer(api_key="sk-xxx")

        @server.skill(name="translate", category="transformation/translate",
                       price=0.01, description="Translate text")
        async def translate(input: dict) -> dict:
            return {"translated": do_translate(input["text"], input["target_lang"])}

        asyncio.run(server.start())
    """

    def __init__(
        self,
        api_key: str,
        base_url: str = "https://api.clawmoney.ai",
    ):
        self.api_key = api_key
        self.base_url = base_url.rstrip("/")
        self._skills: dict[str, SkillHandler] = {}
        self._ws: ClientConnection | None = None
        self._running = False
        self._heartbeat_task: asyncio.Task | None = None

    # ------------------------------------------------------------------
    # Skill registration
    # ------------------------------------------------------------------

    def skill(
        self,
        name: str,
        category: str,
        price: float,
        description: str,
        metadata: dict[str, Any] | None = None,
    ) -> Callable:
        """Decorator to register a skill handler.

        The decorated function must be ``async def handler(input: dict) -> dict``.
        """

        def decorator(func: Callable) -> Callable:
            if not asyncio.iscoroutinefunction(func):
                raise TypeError(f"Skill handler '{name}' must be an async function")
            self._skills[name] = SkillHandler(
                name=name,
                category=category,
                price=price,
                description=description,
                func=func,
                metadata=metadata or {},
            )
            return func

        return decorator

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

    async def start(self):
        """Register skills, connect WebSocket, and serve until stopped."""
        if not self._skills:
            raise ClawMoneyError("No skills registered. Use @server.skill() first.")

        self._running = True
        self._install_signal_handlers()

        # Register skills via REST
        await self._register_skills()

        # Connect + serve loop with auto-reconnect
        backoff = _INITIAL_BACKOFF
        while self._running:
            try:
                await self._connect_and_serve()
                # Clean exit (stop() was called)
                break
            except websockets.exceptions.ConnectionClosedError as e:
                logger.warning("WebSocket closed: %s", e)
            except (OSError, websockets.exceptions.WebSocketException) as e:
                logger.warning("Connection error: %s", e)
            except Exception as e:
                logger.error("Unexpected error: %s", e, exc_info=True)

            if not self._running:
                break

            logger.info("Reconnecting in %.1fs ...", backoff)
            await asyncio.sleep(backoff)
            backoff = min(backoff * 2, _MAX_BACKOFF)

        logger.info("Agent server stopped.")

    async def stop(self):
        """Gracefully shut down."""
        self._running = False
        if self._heartbeat_task and not self._heartbeat_task.done():
            self._heartbeat_task.cancel()
        if self._ws:
            await self._ws.close()

    # ------------------------------------------------------------------
    # Internal: REST skill registration
    # ------------------------------------------------------------------

    async def _register_skills(self):
        """Register all decorated skills with the Hub via REST API."""
        async with httpx.AsyncClient(
            base_url=self.base_url,
            headers={
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json",
            },
            timeout=30.0,
        ) as http:
            payload = {
                "skills": [h.to_registration_dict() for h in self._skills.values()],
            }
            resp = await http.post("/v1/agent/skills", json=payload)
            if resp.status_code == 401:
                raise AuthenticationError()
            if resp.status_code != 200:
                raise ClawMoneyError(
                    f"Failed to register skills ({resp.status_code}): {resp.text}"
                )
            logger.info(
                "Registered %d skill(s): %s",
                len(self._skills),
                ", ".join(self._skills.keys()),
            )

    # ------------------------------------------------------------------
    # Internal: WebSocket
    # ------------------------------------------------------------------

    async def _connect_and_serve(self):
        """Connect WebSocket and process messages until disconnect or stop."""
        ws_scheme = "wss" if self.base_url.startswith("https") else "ws"
        host = self.base_url.split("://", 1)[1]
        ws_url = f"{ws_scheme}://{host}/ws/agent?api_key={self.api_key}"

        logger.info("Connecting to %s://%s/ws/agent ...", ws_scheme, host)

        async with websockets.connect(
            ws_url,
            ping_interval=None,  # we handle heartbeat ourselves
            close_timeout=5,
        ) as ws:
            self._ws = ws
            logger.info("WebSocket connected.")

            # Start heartbeat
            self._heartbeat_task = asyncio.create_task(self._heartbeat_loop(ws))

            try:
                async for raw in ws:
                    if not self._running:
                        break
                    await self._handle_message(ws, raw)
            finally:
                if self._heartbeat_task and not self._heartbeat_task.done():
                    self._heartbeat_task.cancel()
                self._ws = None

    async def _heartbeat_loop(self, ws: ClientConnection):
        """Send periodic heartbeat pings."""
        try:
            while self._running:
                await asyncio.sleep(_HEARTBEAT_INTERVAL)
                await ws.send(json.dumps({"type": "heartbeat", "ts": time.time()}))
        except asyncio.CancelledError:
            pass
        except Exception:
            pass  # connection probably closed; reconnect will handle it

    async def _handle_message(self, ws: ClientConnection, raw: str | bytes):
        """Route an incoming WebSocket message to the appropriate handler."""
        try:
            msg = json.loads(raw)
        except json.JSONDecodeError:
            logger.warning("Received non-JSON message, ignoring.")
            return

        msg_type = msg.get("type")

        if msg_type == "service_call":
            await self._handle_service_call(ws, msg)
        elif msg_type == "heartbeat_ack":
            pass  # server acknowledged our heartbeat
        elif msg_type == "error":
            logger.error("Server error: %s", msg.get("message", msg))
        else:
            logger.debug("Unknown message type: %s", msg_type)

    async def _handle_service_call(self, ws: ClientConnection, msg: dict):
        """Execute a skill handler and send the result back."""
        call_id = msg.get("call_id", "")
        skill_name = msg.get("skill", "")
        input_data = msg.get("input", {})

        handler = self._skills.get(skill_name)
        if handler is None:
            await ws.send(
                json.dumps(
                    {
                        "type": "deliver",
                        "call_id": call_id,
                        "error": f"Unknown skill: {skill_name}",
                    }
                )
            )
            return

        try:
            start = time.monotonic()
            output = await handler.func(input_data)
            duration_ms = int((time.monotonic() - start) * 1000)

            await ws.send(
                json.dumps(
                    {
                        "type": "deliver",
                        "call_id": call_id,
                        "output": output,
                        "duration_ms": duration_ms,
                    }
                )
            )
        except Exception as e:
            logger.exception("Skill '%s' raised an error", skill_name)
            await ws.send(
                json.dumps(
                    {
                        "type": "deliver",
                        "call_id": call_id,
                        "error": str(e),
                    }
                )
            )

    # ------------------------------------------------------------------
    # Signal handling
    # ------------------------------------------------------------------

    def _install_signal_handlers(self):
        """Handle SIGINT/SIGTERM for graceful shutdown."""
        loop = asyncio.get_running_loop()
        for sig in (signal.SIGINT, signal.SIGTERM):
            loop.add_signal_handler(sig, lambda: asyncio.create_task(self.stop()))
