"""WebSocket connection manager.

Handles WebSocket connections, channel subscriptions, message routing,
and automatic reconnection with exponential backoff.
"""

from __future__ import annotations

import asyncio
import json
import logging
from typing import Any, Awaitable, Callable, TypeVar

import websockets
from websockets.asyncio.client import ClientConnection

from .channel import ChannelImpl
from .reconnect import ReconnectManager
from .types import (
    ChannelSubscriptionOptions,
    ConnectionEventListener,
    ConnectionState,
    MessageCallback,
    ReconnectConfig,
    WebSocketClientConfig,
)

logger = logging.getLogger(__name__)

T = TypeVar("T")


class ChannelManager:
    """WebSocket connection manager with automatic reconnection.

    Manages the WebSocket connection lifecycle, channel subscriptions,
    and provides a clean API for real-time communication.
    """

    def __init__(
        self,
        config: WebSocketClientConfig,
        get_token: Callable[[], str | None] | Callable[[], Awaitable[str | None]],
        on_connect: Callable[[], None] | None = None,
        on_disconnect: Callable[[], None] | None = None,
        on_error: Callable[[Exception], None] | None = None,
    ) -> None:
        """Initialize the channel manager.

        Args:
            config: WebSocket client configuration
            get_token: Function to get the current auth token (sync or async)
            on_connect: Called when connection is established
            on_disconnect: Called when connection is closed
            on_error: Called when an error occurs
        """
        self._config = config
        self._get_token = get_token
        self._on_connect = on_connect or (lambda: None)
        self._on_disconnect = on_disconnect or (lambda: None)
        self._on_error = on_error or (lambda e: None)

        self._ws: ClientConnection | None = None
        self._channels: dict[str, ChannelImpl] = {}
        self._message_queue: list[dict[str, Any]] = []
        self._reconnect_manager = ReconnectManager(
            ReconnectConfig(
                initial_delay=1.0,
                max_delay=30.0,
                multiplier=2.0,
                max_attempts=None,  # Infinite
            )
        )
        self._connection_state = ConnectionState.DISCONNECTED
        self._connection_listeners: set[ConnectionEventListener] = set()
        self._intentional_disconnect = False
        self._heartbeat_task: asyncio.Task[None] | None = None
        self._receive_task: asyncio.Task[None] | None = None

    async def connect(self) -> None:
        """Connect to the WebSocket server."""
        if self._connection_state in (
            ConnectionState.CONNECTING,
            ConnectionState.CONNECTED,
        ):
            return

        self._intentional_disconnect = False
        self._set_connection_state(ConnectionState.CONNECTING)

        try:
            # Get token
            token = self._get_token()
            if asyncio.iscoroutine(token):
                token = await token

            # Build URL with token
            url = self._config.url
            if token:
                separator = "&" if "?" in url else "?"
                url = f"{url}{separator}token={token}"

            # Connect
            self._ws = await websockets.connect(url)

            # Start message receiving loop
            self._receive_task = asyncio.create_task(self._receive_loop())

            # Handle successful connection
            await self._handle_open()

        except Exception as e:
            self._handle_error(e if isinstance(e, Exception) else Exception(str(e)))
            raise

    async def disconnect(self) -> None:
        """Disconnect from the WebSocket server."""
        self._intentional_disconnect = True
        self._set_connection_state(ConnectionState.DISCONNECTING)
        self._reconnect_manager.cancel()
        self._stop_heartbeat()

        # Cancel receive task
        if self._receive_task and not self._receive_task.done():
            self._receive_task.cancel()
            try:
                await self._receive_task
            except asyncio.CancelledError:
                pass
            self._receive_task = None

        # Close WebSocket
        if self._ws:
            await self._ws.close()
            self._ws = None

        # Clear all channels
        self._channels.clear()
        self._message_queue.clear()
        self._set_connection_state(ConnectionState.DISCONNECTED)

    def subscribe(
        self,
        channel_name: str,
        *,
        params: dict[str, str] | None = None,
        on_subscribed: Callable[[], None] | None = None,
        on_error: Callable[[Exception], None] | None = None,
        on_message: MessageCallback | None = None,
    ) -> ChannelImpl:
        """Subscribe to a WebSocket channel.

        Args:
            channel_name: Name of the channel to subscribe to
            params: Additional parameters to send when subscribing
            on_subscribed: Called when successfully subscribed
            on_error: Called when an error occurs
            on_message: Called when a message is received

        Returns:
            Channel instance for receiving messages and sending events
        """
        # Check if channel already exists
        channel = self._channels.get(channel_name)

        if not channel:
            # Create send function for this channel
            def send_fn(event: str, data: Any) -> None:
                self._send_channel_message(channel_name, event, data)

            options = ChannelSubscriptionOptions(params=params)

            # Create new channel
            channel = ChannelImpl(
                name=channel_name,
                send_fn=send_fn,
                options=options,
                on_unsubscribe=lambda: self.unsubscribe(channel_name),
                on_subscribed=on_subscribed,
                on_error=on_error,
                on_message=on_message,
            )
            self._channels[channel_name] = channel

            # Send subscribe action if connected
            self._send_subscribe(channel_name, params)

        return channel

    def unsubscribe(self, channel_name: str) -> None:
        """Unsubscribe from a WebSocket channel.

        Args:
            channel_name: Name of the channel to unsubscribe from
        """
        channel = self._channels.get(channel_name)
        if channel:
            channel.mark_unsubscribed()
            del self._channels[channel_name]

            # Send unsubscribe action if connected
            self._send_unsubscribe(channel_name)

    def is_connected(self) -> bool:
        """Check if connected to the server."""
        return self._connection_state == ConnectionState.CONNECTED

    def get_connection_state(self) -> ConnectionState:
        """Get the current connection state."""
        return self._connection_state

    def on_connection_change(
        self, callback: ConnectionEventListener
    ) -> Callable[[], None]:
        """Register a listener for connection state changes.

        Args:
            callback: Function to call when connection state changes

        Returns:
            Unsubscribe function
        """
        self._connection_listeners.add(callback)

        def unsubscribe() -> None:
            self._connection_listeners.discard(callback)

        return unsubscribe

    async def _receive_loop(self) -> None:
        """Receive messages from the WebSocket."""
        if not self._ws:
            return

        try:
            async for message in self._ws:
                if isinstance(message, bytes):
                    message = message.decode("utf-8")
                self._handle_message(message)
        except websockets.ConnectionClosed:
            await self._handle_close()
        except Exception as e:
            self._handle_error(e)

    async def _handle_open(self) -> None:
        """Handle WebSocket open event."""
        self._set_connection_state(ConnectionState.CONNECTED)
        self._reconnect_manager.reset()
        self._on_connect()

        # Start heartbeat
        self._start_heartbeat()

        # Resubscribe to all channels
        for channel_name, channel in self._channels.items():
            options = channel.get_options()
            self._send_subscribe(channel_name, options.params)

        # Flush message queue
        await self._flush_message_queue()

    async def _handle_close(self) -> None:
        """Handle WebSocket close event."""
        self._stop_heartbeat()
        self._on_disconnect()

        # Mark all channels as unsubscribed
        for channel in self._channels.values():
            channel.mark_unsubscribed()

        # Attempt to reconnect if not intentional
        if not self._intentional_disconnect and self._config.auto_reconnect:
            await self._attempt_reconnect()
        else:
            self._set_connection_state(ConnectionState.DISCONNECTED)

    def _handle_error(self, error: Exception) -> None:
        """Handle WebSocket error event."""
        self._on_error(error)
        self._set_connection_state(ConnectionState.DISCONNECTED)

    def _handle_message(self, raw_message: str) -> None:
        """Handle incoming WebSocket message.

        Server messages use the format:
        - { channel: "...", event: "subscribed", data: {...} }
        - { channel: "...", event: "data", data: {...} }
        - { channel: "...", event: "error", data: { code: "...", message: "..." } }
        - { event: "pong" }
        """
        try:
            message = json.loads(raw_message)
        except json.JSONDecodeError:
            logger.error(f"Failed to parse WebSocket message: {raw_message}")
            return

        event = message.get("event")

        # Handle pong (no channel)
        if event == "pong":
            return

        # Handle channel messages
        channel_name = message.get("channel")
        data = message.get("data")

        if not channel_name:
            # Global error without channel
            if event == "error" and data:
                error_msg = data.get("message", "Unknown error")
                self._on_error(Exception(error_msg))
            return

        if event == "subscribed":
            self._handle_subscribed(channel_name)
        elif event == "error":
            self._handle_error_message(channel_name, data)
        else:
            # All other events are data messages
            self._handle_data_message(channel_name, event, data)

    def _handle_subscribed(self, channel_name: str) -> None:
        """Handle subscribed confirmation."""
        channel = self._channels.get(channel_name)
        if channel:
            channel.mark_subscribed()

    def _handle_data_message(
        self, channel_name: str, event: str, data: Any
    ) -> None:
        """Handle data message."""
        channel = self._channels.get(channel_name)
        if channel:
            channel.handle_message(event, data)

    def _handle_error_message(self, channel_name: str, data: Any) -> None:
        """Handle error message."""
        error_msg = data.get("message", "Channel error") if data else "Channel error"
        error = Exception(error_msg)

        channel = self._channels.get(channel_name)
        if channel:
            channel.handle_error(error)

    def _send_subscribe(
        self, channel_name: str, params: dict[str, str] | None
    ) -> None:
        """Send subscribe action."""
        message = {
            "action": "subscribe",
            "channel": channel_name,
        }
        if params:
            message["payload"] = params

        self._send(message)

    def _send_unsubscribe(self, channel_name: str) -> None:
        """Send unsubscribe action."""
        message = {
            "action": "unsubscribe",
            "channel": channel_name,
        }
        self._send(message)

    def _send_channel_message(
        self, channel_name: str, event: str, data: Any
    ) -> None:
        """Send a message to a channel."""
        message = {
            "action": "message",
            "channel": channel_name,
            "event": event,
            "payload": data,
        }
        self._send(message)

    def _send(self, message: dict[str, Any]) -> None:
        """Send a message to the server."""
        if self._ws and self._connection_state == ConnectionState.CONNECTED:
            asyncio.create_task(self._ws.send(json.dumps(message)))
        else:
            # Queue message for later
            self._message_queue.append(message)

    async def _flush_message_queue(self) -> None:
        """Flush queued messages."""
        if not self._ws or self._connection_state != ConnectionState.CONNECTED:
            return

        for message in self._message_queue:
            await self._ws.send(json.dumps(message))

        self._message_queue.clear()

    async def _attempt_reconnect(self) -> None:
        """Attempt to reconnect."""
        self._set_connection_state(ConnectionState.RECONNECTING)

        try:
            await self._reconnect_manager.schedule_reconnect(self.connect)
        except Exception as e:
            self._handle_error(e if isinstance(e, Exception) else Exception(str(e)))

    def _set_connection_state(self, state: ConnectionState) -> None:
        """Set connection state and notify listeners."""
        if self._connection_state == state:
            return

        self._connection_state = state

        # Notify listeners
        for listener in self._connection_listeners:
            try:
                listener(state)
            except Exception as e:
                logger.error(f"Error in connection state listener: {e}")

    def _start_heartbeat(self) -> None:
        """Start heartbeat."""
        self._stop_heartbeat()

        async def heartbeat_loop() -> None:
            while self._ws and self._connection_state == ConnectionState.CONNECTED:
                try:
                    await self._ws.send(json.dumps({"action": "ping"}))
                    await asyncio.sleep(self._config.heartbeat_interval)
                except Exception:
                    break

        self._heartbeat_task = asyncio.create_task(heartbeat_loop())

    def _stop_heartbeat(self) -> None:
        """Stop heartbeat."""
        if self._heartbeat_task and not self._heartbeat_task.done():
            self._heartbeat_task.cancel()
            self._heartbeat_task = None

    async def destroy(self) -> None:
        """Destroy the channel manager and clean up resources."""
        await self.disconnect()
        self._reconnect_manager.destroy()
        self._connection_listeners.clear()
