"""
nodelink.client
~~~~~~~~~~~~~~~
The nodelink Client.
"""

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

import websockets
from websockets.exceptions import ConnectionClosed

logger = logging.getLogger("nodelink.client")


class Client:
    """
    A nodelink client.

    Args:
        host:            server hostname (default ``"localhost"``)
        port:            server port    (default ``5000``)
        reconnect:       automatically reconnect on disconnect (default ``False``)
        reconnect_delay: seconds to wait between attempts (default ``2.0``)
        max_retries:     max reconnect attempts before giving up;
                         ``None`` means retry forever (default ``None``)

    Example::

        from nodelink import Client

        client = Client("localhost", 5000, reconnect=True, reconnect_delay=3.0)

        @client.on("connect")
        def on_connect():
            client.send("hello", {"msg": "hi!"})

        @client.on("welcome")
        def on_welcome(data):
            print(f"Server said: {data}")

        @client.on("disconnect")
        def on_disconnect():
            print("Lost connection.")

        @client.on("reconnecting")
        def on_reconnecting(data):
            print(f"Reconnecting... attempt {data['attempt']}")

        @client.on_error()
        def on_error(error):
            print(f"Error: {error}")

        client.connect()
    """

    def __init__(
        self,
        host: str = "localhost",
        port: int = 5000,
        reconnect: bool = False,
        reconnect_delay: float = 2.0,
        max_retries: Optional[int] = None,
    ):
        self.host = host
        self.port = port
        self._uri = f"ws://{host}:{port}"

        self.reconnect = reconnect
        self.reconnect_delay = reconnect_delay
        self.max_retries = max_retries

        self._handlers:      Dict[str, List[Callable]] = {}
        self._error_handlers: List[Callable]           = []

        self._ws = None
        self._connected = False
        self._loop: Optional[asyncio.AbstractEventLoop] = None
        self._retry_count: int = 0

        # Messages sent before the connection is ready are queued
        # and flushed automatically once connected
        self._send_queue: List[tuple] = []

    def on(self, event: str) -> Callable:
        """
        Register a handler for a server event.

        Built-in events: ``"connect"``, ``"disconnect"``, ``"reconnecting"``

        Handler signatures:

        - ``"connect"`` / ``"disconnect"``: ``fn()``
        - ``"reconnecting"``:               ``fn(data)``  —  ``data["attempt"]``, ``data["delay"]``
        - all others:                        ``fn(data)``

        Example::

            @client.on("game_start")
            def on_start(data):
                print("Game starting!")
        """
        def decorator(fn: Callable) -> Callable:
            self._handlers.setdefault(event, []).append(fn)
            return fn
        return decorator

    def on_error(self) -> Callable:
        """
        Register an error handler.

        Example::

            @client.on_error()
            def on_error(error):
                print(f"Something went wrong: {error}")
        """
        def decorator(fn: Callable) -> Callable:
            self._error_handlers.append(fn)
            return fn
        return decorator

    def connect(self) -> None:
        """Connect to the server.  Blocks until permanently disconnected."""
        try:
            asyncio.run(self._run())
        except KeyboardInterrupt:
            logger.info("Client disconnected.")

    async def _run(self) -> None:
        self._loop = asyncio.get_running_loop()
        try:
            while True:
                did_connect = False
                logger.info(f"Connecting to {self._uri}")
                try:
                    async with websockets.connect(self._uri) as ws:
                        self._ws = ws
                        self._connected = True
                        did_connect = True
                        self._retry_count = 0  # reset on successful connect

                        # Flush messages that were queued while disconnected
                        for queued_event, queued_data in self._send_queue:
                            await self._send_raw(queued_event, queued_data)
                        self._send_queue.clear()

                        await self._fire_single("connect")

                        async for raw in ws:
                            try:
                                message = json.loads(raw)
                                await self._handle_message(message)
                            except json.JSONDecodeError:
                                logger.warning(f"Malformed message: {raw!r}")
                            except Exception as e:
                                logger.error(
                                    f"Message error: {e}", exc_info=True
                                )
                                await self._fire_error(e)

                except ConnectionClosed:
                    logger.info("Disconnected from server")
                except OSError as e:
                    err = ConnectionError(
                        f"Could not connect to {self._uri}: {e}"
                    )
                    logger.error(str(err))
                    await self._fire_error(err)
                finally:
                    self._connected = False
                    self._ws = None
                    if did_connect:
                        await self._fire_single("disconnect")

                # ── Reconnect logic ────────────────────────────────────
                if not self.reconnect:
                    break

                if (
                    self.max_retries is not None
                    and self._retry_count >= self.max_retries
                ):
                    logger.info(
                        f"Max retries ({self.max_retries}) reached. Giving up."
                    )
                    break

                self._retry_count += 1
                logger.info(
                    f"Reconnecting in {self.reconnect_delay}s "
                    f"(attempt {self._retry_count})"
                    + (
                        f" / {self.max_retries}"
                        if self.max_retries is not None
                        else ""
                    )
                )
                await self._fire(
                    "reconnecting",
                    {
                        "attempt": self._retry_count,
                        "delay":   self.reconnect_delay,
                    },
                )
                await asyncio.sleep(self.reconnect_delay)

        finally:
            self._loop = None

    async def _handle_message(self, message: dict) -> None:
        event = message.get("event")
        data  = message.get("data") or {}
        if event:
            await self._fire(event, data)

    def send(self, event: str, data: Any = None) -> None:
        """
        Send an event + data to the server.

        Safe to call from any thread.  Messages sent before the connection
        is established (or while reconnecting) are queued and delivered
        automatically once connected.

        Example::

            client.send("move", {"x": 10, "y": 20})
            client.send("chat", {"msg": "hello!"})
        """
        if not self._connected or not self._ws:
            self._send_queue.append((event, data))
            return

        coro = self._send_raw(event, data)
        try:
            running_loop = asyncio.get_running_loop()
        except RuntimeError:
            running_loop = None

        if running_loop is self._loop:
            # Called from within the event loop — schedule directly
            asyncio.ensure_future(coro)
        elif self._loop:
            # Called from a background thread — wakes the selector
            asyncio.run_coroutine_threadsafe(coro, self._loop)

    async def _send_raw(self, event: str, data: Any = None) -> None:
        if self._ws and self._connected:
            try:
                await self._ws.send(json.dumps({"event": event, "data": data}))
            except ConnectionClosed:
                self._connected = False
            except Exception as e:
                logger.error(f"Send error: {e}")
                await self._fire_error(e)

    async def _fire_single(self, event: str) -> None:
        for handler in self._handlers.get(event, []):
            try:
                result = handler()
                if asyncio.iscoroutine(result):
                    await result
            except Exception as e:
                logger.error(f"Error in '{event}' handler: {e}", exc_info=True)
                await self._fire_error(e)

    async def _fire(self, event: str, data: Any) -> None:
        for handler in self._handlers.get(event, []):
            try:
                result = handler(data)
                if asyncio.iscoroutine(result):
                    await result
            except Exception as e:
                logger.error(f"Error in '{event}' handler: {e}", exc_info=True)
                await self._fire_error(e)

    async def _fire_error(self, error: Exception) -> None:
        if self._error_handlers:
            for handler in self._error_handlers:
                try:
                    result = handler(error)
                    if asyncio.iscoroutine(result):
                        await result
                except Exception as e:
                    logger.error(f"Error in error handler: {e}", exc_info=True)
        else:
            logger.error(
                f"Unhandled error (register @client.on_error() to catch): {error}"
            )
