"""WebSocket client for ZenOTC SDK."""

from __future__ import annotations

import asyncio
import hashlib
import hmac
import json
import logging
import time
from typing import Any, Callable, Coroutine, Dict, List, Optional, Set

import websockets
from websockets.client import WebSocketClientProtocol

from zenotc.config import Config
from zenotc.exceptions import ConnectionError

logger = logging.getLogger(__name__)

MessageHandler = Callable[[Dict[str, Any]], Coroutine[Any, Any, None]]


class WebSocketClient:
    """WebSocket client for real-time data streams."""

    def __init__(self, config: Config) -> None:
        self._config = config
        self._ws: Optional[WebSocketClientProtocol] = None
        self._subscriptions: Set[str] = set()
        self._handlers: Dict[str, List[MessageHandler]] = {}
        self._running = False
        self._reconnect_attempts = 0
        self._max_reconnect_attempts = 10
        self._reconnect_delay = 1.0
        self._receive_task: Optional[asyncio.Task[None]] = None
        self._ping_task: Optional[asyncio.Task[None]] = None

    @property
    def is_connected(self) -> bool:
        """Check if WebSocket is connected."""
        return self._ws is not None and self._ws.open

    def _generate_auth_headers(self) -> Dict[str, str]:
        """Generate authentication headers for WebSocket connection.

        Uses the same signature scheme as HTTP requests:
        HMAC-SHA256(SHA256(secret), message)
        """
        timestamp = str(int(time.time() * 1000))
        # For WebSocket, we sign the connection path
        path = "/ws"
        message = f"{timestamp}GET{path}"

        # Hash the secret first - server stores SHA256(secret) as secretHash
        hmac_key = hashlib.sha256(
            self._config.api_secret.encode("utf-8")
        ).hexdigest()
        signature = hmac.new(
            hmac_key.encode("utf-8"),
            message.encode("utf-8"),
            hashlib.sha256,
        ).hexdigest()

        return {
            "X-API-Key": self._config.api_key,
            "X-API-Timestamp": timestamp,
            "X-API-Signature": signature,
        }

    async def connect(self) -> None:
        """Connect to WebSocket server."""
        if self.is_connected:
            return

        url = self._config.websocket_url
        headers = self._generate_auth_headers()

        try:
            self._ws = await websockets.connect(
                url,
                additional_headers=headers,
                ping_interval=None,  # We handle pings manually
            )
            self._running = True
            self._reconnect_attempts = 0
            logger.info(f"Connected to WebSocket: {url}")

            # Start background tasks
            self._receive_task = asyncio.create_task(self._receive_loop())
            self._ping_task = asyncio.create_task(self._ping_loop())

            # Resubscribe to previous subscriptions
            if self._subscriptions:
                await self._resubscribe()

        except Exception as e:
            raise ConnectionError(f"Failed to connect to WebSocket: {e}") from e

    async def disconnect(self) -> None:
        """Disconnect from WebSocket server."""
        self._running = False

        if self._receive_task:
            self._receive_task.cancel()
            try:
                await self._receive_task
            except asyncio.CancelledError:
                pass
            self._receive_task = None

        if self._ping_task:
            self._ping_task.cancel()
            try:
                await self._ping_task
            except asyncio.CancelledError:
                pass
            self._ping_task = None

        if self._ws:
            await self._ws.close()
            self._ws = None

        logger.info("Disconnected from WebSocket")

    async def subscribe(
        self,
        channel: str,
        symbols: Optional[List[str]] = None,
        handler: Optional[MessageHandler] = None,
    ) -> None:
        """
        Subscribe to a channel.

        Args:
            channel: Channel name (e.g., "prices", "orderbook")
            symbols: Optional list of symbols to subscribe to
            handler: Optional callback for messages on this channel
        """
        if not self.is_connected:
            await self.connect()

        subscription_key = f"{channel}:{','.join(symbols or [])}"
        self._subscriptions.add(subscription_key)

        if handler:
            if channel not in self._handlers:
                self._handlers[channel] = []
            self._handlers[channel].append(handler)

        message = {
            "type": "subscribe",
            "channel": channel,
        }
        if symbols:
            message["symbols"] = symbols

        await self._send(message)
        logger.info(f"Subscribed to {channel}: {symbols}")

    async def unsubscribe(
        self,
        channel: str,
        symbols: Optional[List[str]] = None,
    ) -> None:
        """
        Unsubscribe from a channel.

        Args:
            channel: Channel name
            symbols: Optional list of symbols to unsubscribe from
        """
        if not self.is_connected:
            return

        subscription_key = f"{channel}:{','.join(symbols or [])}"
        self._subscriptions.discard(subscription_key)

        message = {
            "type": "unsubscribe",
            "channel": channel,
        }
        if symbols:
            message["symbols"] = symbols

        await self._send(message)
        logger.info(f"Unsubscribed from {channel}: {symbols}")

    async def _send(self, message: Dict[str, Any]) -> None:
        """Send a message to the WebSocket server."""
        if not self._ws:
            raise ConnectionError("Not connected to WebSocket")

        await self._ws.send(json.dumps(message))

    async def _receive_loop(self) -> None:
        """Background task to receive messages."""
        while self._running and self._ws:
            try:
                message = await self._ws.recv()
                data = json.loads(message)
                await self._handle_message(data)
            except websockets.ConnectionClosed:
                logger.warning("WebSocket connection closed")
                if self._running:
                    await self._reconnect()
                break
            except json.JSONDecodeError as e:
                logger.error(f"Failed to parse WebSocket message: {e}")
            except Exception as e:
                logger.error(f"Error in receive loop: {e}")

    async def _ping_loop(self) -> None:
        """Background task to send ping messages."""
        while self._running and self._ws:
            try:
                await asyncio.sleep(self._config.websocket_ping_interval)
                if self._ws and self._ws.open:
                    await self._ws.ping()
            except Exception as e:
                logger.debug(f"Ping failed: {e}")

    async def _handle_message(self, data: Dict[str, Any]) -> None:
        """Handle incoming WebSocket message."""
        channel = data.get("channel")
        if not channel:
            return

        handlers = self._handlers.get(channel, [])
        for handler in handlers:
            try:
                await handler(data)
            except Exception as e:
                logger.error(f"Error in message handler for {channel}: {e}")

    async def _reconnect(self) -> None:
        """Attempt to reconnect to WebSocket server."""
        self._ws = None

        while self._running and self._reconnect_attempts < self._max_reconnect_attempts:
            self._reconnect_attempts += 1
            delay = self._reconnect_delay * (2 ** (self._reconnect_attempts - 1))
            logger.info(
                f"Reconnecting in {delay}s (attempt {self._reconnect_attempts}/"
                f"{self._max_reconnect_attempts})"
            )

            await asyncio.sleep(delay)

            try:
                await self.connect()
                return
            except Exception as e:
                logger.error(f"Reconnection failed: {e}")

        logger.error("Max reconnection attempts reached")
        self._running = False

    async def _resubscribe(self) -> None:
        """Resubscribe to all previous subscriptions after reconnect."""
        for subscription in self._subscriptions:
            channel, symbols_str = subscription.split(":", 1)
            symbols = symbols_str.split(",") if symbols_str else None

            message = {
                "type": "subscribe",
                "channel": channel,
            }
            if symbols and symbols[0]:
                message["symbols"] = symbols

            await self._send(message)
            logger.info(f"Resubscribed to {channel}")
