import asyncio
import json
import random
import time
from dataclasses import dataclass
from typing import Any, Callable, List, Optional, Tuple

import websockets.asyncio.client
import websockets.exceptions
from loguru import logger

from src.configs import WebSocketConfig


# === WEBSOCKET FORMAT CONSTANTS ===

# WebSocket Frontend Orderbook Format Constants
# The frontend orderbook WebSocket uses the following format:
# - Prices: Always in 10^18 format (universal across all markets)
# - Sizes: In the market's size_precision format (varies by market)
WEBSOCKET_PRICE_PRECISION = 1_000_000_000_000_000_000  # 10^18 - prices always in wei-like format


# === DATACLASSES ===


@dataclass
class SubscriptionRequest:
    """Subscription request to server."""

    type: str  # "subscribe"
    channel: str  # "frontendOrderbook"
    market: str  # Market address (lowercase)


@dataclass
class MarketParams:
    """Market configuration parameters."""

    price_precision: int
    size_precision: int
    base_asset: str
    base_asset_decimals: int
    quote_asset: str
    quote_asset_decimals: int
    tick_size: int
    min_size: int
    max_size: int
    taker_fee_bps: int
    maker_fee_bps: int


@dataclass
class VaultParams:
    """Vault liquidity parameters."""

    vault_best_bid: int
    bid_partially_filled_size: int
    vault_best_ask: int
    ask_partially_filled_size: int
    vault_bid_order_size: int
    vault_ask_order_size: int
    spread: int


@dataclass
class SubscriptionResponse:
    """Response to subscription request."""

    type: str  # "subscribed"
    status: str  # "success" | "pending" | "error"
    message: Optional[str] = None
    data: Optional[dict] = None  # FrontendOrderbookData as dict


@dataclass
class FrontendEvent:
    """Individual orderbook event.

    Note: Prices (p) are in 10^18 format. Sizes (s) are in the market's size_precision format.
    Use KuruFrontendOrderbookClient.format_websocket_price() and format_websocket_size() to convert.
    """

    e: str  # Event type (e.g., "Trade", "OrderCreated", etc.)
    ts: int  # Timestamp
    mad: str  # Market address
    th: Optional[str] = None  # Transaction hash
    p: Optional[int] = None  # Price (in 10^18 format - use format_websocket_price())
    s: Optional[int] = None  # Size (in size_precision format - use format_websocket_size())
    ib: Optional[bool] = None  # Is buy
    t: Optional[str] = None  # Taker address
    m: Optional[str] = None  # Maker address


@dataclass
class FrontendOrderbookUpdate:
    """Incremental orderbook update.

    Note: Prices in bids/asks are in 10^18 format. Sizes are in the market's size_precision format.
    Use KuruFrontendOrderbookClient.format_websocket_price() and format_websocket_size() to convert.
    """

    events: List[FrontendEvent]
    b: Optional[List[Tuple[int, int]]] = None  # Bids: [(price_10e18, size_precision), ...]
    a: Optional[List[Tuple[int, int]]] = None  # Asks: [(price_10e18, size_precision), ...]
    v: Optional[VaultParams] = None  # Updated vault params


# === MAIN CLIENT CLASS ===


class KuruFrontendOrderbookClient:
    """
    WebSocket client for Kuru frontend orderbook data.

    Features:
    - Automatic reconnection with exponential backoff
    - Built-in heartbeat monitoring
    - BigInt parsing for large numbers
    - Type-safe dataclass structures
    - Async context manager support
    - Queue-based orderbook updates

    Args:
        ws_url: WebSocket server URL (ws:// or wss://)
        market_address: Market contract address
        update_queue: asyncio.Queue to receive orderbook updates
        websocket_config: WebSocket behavior configuration
        on_error: Optional callback for errors
    """

    def __init__(
        self,
        ws_url: str,
        market_address: str,
        update_queue: asyncio.Queue[FrontendOrderbookUpdate],
        websocket_config: Optional[WebSocketConfig] = None,
        on_error: Optional[Callable[[Exception], None]] = None,
    ) -> None:
        """
        Initialize the KuruFrontendOrderbookClient.

        Args:
            ws_url: WebSocket server URL
            market_address: Market contract address
            update_queue: Queue to receive orderbook updates
            websocket_config: WebSocket behavior configuration.
                            If None, uses default WebSocketConfig()
            on_error: Optional callback for errors

        Raises:
            ValueError: If any validation fails
        """
        # Use default config if not provided
        if websocket_config is None:
            websocket_config = WebSocketConfig()

        # Store config
        self.websocket_config = websocket_config

        # Validation
        if not ws_url:
            raise ValueError("ws_url cannot be empty")
        if not ws_url.startswith(("ws://", "wss://")):
            raise ValueError("ws_url must start with ws:// or wss://")
        if not market_address:
            raise ValueError("market_address cannot be empty")
        if not isinstance(update_queue, asyncio.Queue):
            raise ValueError("update_queue must be an asyncio.Queue")

        # Connection parameters
        self._ws_url = ws_url
        self._market_address = market_address.lower()  # Normalize to lowercase
        self._update_queue = update_queue
        self._on_error = on_error

        # Get reconnection parameters from config
        self._max_reconnect_attempts = websocket_config.max_reconnect_attempts
        self._reconnect_delay = websocket_config.reconnect_delay

        # Connection state
        self._websocket: Optional[websockets.asyncio.client.ClientConnection] = None
        self._connected = False
        self._closing = False
        self._subscribed = False
        self._initial_snapshot_received = False

        # Reconnection state
        self._reconnect_count = 0
        self._current_reconnect_delay = self._reconnect_delay

        # Background tasks
        self._message_loop_task: Optional[asyncio.Task] = None
        self._heartbeat_task: Optional[asyncio.Task] = None

        # Heartbeat configuration from config
        self._heartbeat_interval = websocket_config.heartbeat_interval
        self._heartbeat_timeout = websocket_config.heartbeat_timeout

        # Thread safety
        self._lock = asyncio.Lock()

        logger.info(
            f"Initialized KuruFrontendOrderbookClient for market {self._market_address}"
        )

    @staticmethod
    def format_websocket_price(raw_price: int) -> float:
        """
        Convert WebSocket price to human-readable decimal.

        The frontend orderbook WebSocket always sends prices in 10^18 format,
        regardless of the market's contract precision.

        Args:
            raw_price: Raw price from WebSocket (in 10^18 format)

        Returns:
            Human-readable price as float

        Example:
            241470000000000000000 -> 241.47
        """
        return raw_price / WEBSOCKET_PRICE_PRECISION

    @staticmethod
    def format_websocket_size(raw_size: int, size_precision: int) -> float:
        """
        Convert WebSocket size to human-readable decimal.

        The frontend orderbook WebSocket sends sizes in the market's size_precision format.
        You must provide the size_precision from MarketConfig to convert properly.

        Args:
            raw_size: Raw size from WebSocket (in size_precision format)
            size_precision: Size precision from MarketConfig

        Returns:
            Size as float

        Example:
            format_websocket_size(100000000000, 10000000000) -> 10.0
        """
        if size_precision == 0:
            return float(raw_size)
        return raw_size / size_precision

    async def connect(self) -> None:
        """
        Establish WebSocket connection and start message processing.

        This method:
        1. Connects to the WebSocket server
        2. Starts the message processing loop
        3. Starts the heartbeat monitor
        4. Subscribes to the market orderbook

        Raises:
            RuntimeError: If already connected or closing
            ConnectionError: If connection fails
        """
        async with self._lock:
            if self._connected:
                logger.warning("Already connected")
                return

            if self._closing:
                raise RuntimeError("Cannot connect while closing")

            try:
                logger.info(f"Connecting to {self._ws_url}")

                # Use websockets.asyncio.client.connect with built-in ping/pong
                self._websocket = await websockets.asyncio.client.connect(
                    self._ws_url,
                    ping_interval=self._heartbeat_interval,
                    ping_timeout=self._heartbeat_timeout,
                    open_timeout=10.0,
                    close_timeout=10.0,
                    max_size=10 * 1024 * 1024,  # 10MB max message size
                )

                self._connected = True
                self._reconnect_count = 0  # Reset on successful connection

                logger.info("WebSocket connected successfully")

                # Start background tasks
                self._message_loop_task = asyncio.create_task(self._message_loop())
                self._heartbeat_task = asyncio.create_task(self._heartbeat_monitor())

                # Subscribe to market
                await self.subscribe()

            except websockets.exceptions.InvalidURI as e:
                self._connected = False
                self._websocket = None
                logger.error(f"Invalid WebSocket URI: {e}")
                raise ValueError(f"Invalid WebSocket URI: {self._ws_url}") from e
            except OSError as e:
                self._connected = False
                self._websocket = None
                logger.error(f"Network error: {e}")
                raise ConnectionError(f"Failed to connect to {self._ws_url}") from e
            except Exception as e:
                self._connected = False
                self._websocket = None
                logger.error(f"Connection failed: {e}")
                raise

    async def subscribe(self) -> None:
        """
        Subscribe to the market's frontend orderbook.

        Sends subscription request and waits for confirmation.

        Raises:
            RuntimeError: If not connected
        """
        if not self._connected or self._websocket is None:
            raise RuntimeError("Cannot subscribe - not connected")

        if self._subscribed:
            logger.warning("Already subscribed")
            return

        # Build subscription request
        request_dict = {
            "type": "subscribe",
            "channel": "frontendOrderbook",
            "market": self._market_address,
        }

        logger.info(f"Subscribing to market {self._market_address}")

        try:
            await self._websocket.send(json.dumps(request_dict))
            logger.debug(f"Subscription request sent: {request_dict}")
        except Exception as e:
            logger.error(f"Failed to send subscription request: {e}")
            raise

    async def close(self) -> None:
        """
        Close the WebSocket connection and cleanup resources.

        This method:
        1. Sets closing flag to prevent reconnection
        2. Cancels background tasks
        3. Closes WebSocket connection
        4. Resets state
        """
        async with self._lock:
            if self._closing:
                return

            self._closing = True
            logger.info("Closing WebSocket connection")

            await self._cleanup_connection()

            logger.info("WebSocket closed successfully")

    def is_connected(self) -> bool:
        """
        Check if currently connected to WebSocket.

        Returns:
            True if connected, False otherwise
        """
        return self._connected and self._websocket is not None

    async def _cleanup_connection(self) -> None:
        """Clean up connection resources."""
        # Cancel background tasks
        if self._message_loop_task:
            self._message_loop_task.cancel()
            try:
                await self._message_loop_task
            except asyncio.CancelledError:
                pass
            self._message_loop_task = None

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

        # Close WebSocket
        if self._websocket:
            try:
                await self._websocket.close()
            except Exception as e:
                logger.warning(f"Error closing websocket: {e}")
            self._websocket = None

        self._connected = False
        self._subscribed = False
        self._initial_snapshot_received = False

    async def _reconnect(self) -> None:
        """
        Attempt to reconnect with exponential backoff.

        Backoff formula: delay = base_delay * (2 ^ attempt) + jitter
        Max delay capped at 60 seconds.
        """
        async with self._lock:
            if self._closing:
                logger.info("Not reconnecting - client is closing")
                return

            if self._reconnect_count >= self._max_reconnect_attempts:
                error_msg = f"Max reconnection attempts ({self._max_reconnect_attempts}) reached"
                logger.error(error_msg)
                await self._invoke_error_callback(ConnectionError(error_msg))
                return

            # Calculate exponential backoff with jitter
            backoff_multiplier = 2**self._reconnect_count
            self._current_reconnect_delay = min(
                self._reconnect_delay * backoff_multiplier + random.uniform(0, 1),
                60.0,  # Cap at 60 seconds
            )

            self._reconnect_count += 1

            logger.info(
                f"Reconnection attempt {self._reconnect_count}/{self._max_reconnect_attempts} "
                f"in {self._current_reconnect_delay:.2f}s"
            )

            await asyncio.sleep(self._current_reconnect_delay)

            try:
                await self._cleanup_connection()
                await self.connect()
                logger.info("Reconnection successful")
            except Exception as e:
                logger.error(f"Reconnection failed: {e}")
                # Schedule another reconnection attempt
                asyncio.create_task(self._reconnect())

    async def _message_loop(self) -> None:
        """
        Main message processing loop.

        Continuously receives and processes messages from WebSocket.
        Handles connection errors and triggers reconnection.
        """
        logger.info("Message loop started")

        try:
            while self._connected and not self._closing:
                if self._websocket is None:
                    logger.warning("WebSocket is None in message loop")
                    break

                try:
                    # Receive message with timeout
                    message = await asyncio.wait_for(
                        self._websocket.recv(),
                        timeout=self._heartbeat_interval
                        + self._heartbeat_timeout
                        + 5.0,
                    )

                    # Process message
                    self._handle_message(message)

                except asyncio.TimeoutError:
                    logger.warning("Message receive timeout - connection may be stale")
                    break

                except websockets.exceptions.ConnectionClosed as e:
                    logger.warning(f"WebSocket connection closed: {e}")
                    break

                except Exception as e:
                    logger.error(f"Error receiving message: {e}")
                    await self._invoke_error_callback(e)
                    # Continue loop - don't break on parse errors

        except asyncio.CancelledError:
            logger.debug("Message loop cancelled")

        except Exception as e:
            logger.error(f"Fatal error in message loop: {e}")
            await self._invoke_error_callback(e)

        finally:
            logger.info("Message loop ended")

            # If we exited due to connection issue (not intentional close), reconnect
            if not self._closing:
                asyncio.create_task(self._handle_connection_loss())

    def _handle_message(self, message: str) -> None:
        """
        Parse and route incoming message to appropriate handler.

        Args:
            message: Raw JSON message string
        """
        try:
            data = json.loads(message)

            # Detect message type
            msg_type = data.get("type")

            if msg_type == "subscribed":
                # Subscription response
                self._handle_subscription_response(data)

            elif msg_type == "snapshot" or (
                msg_type is None and "b" in data and "a" in data
            ):
                # Initial snapshot or update with full orderbook
                # Check if this is first snapshot
                if not self._initial_snapshot_received:
                    logger.info("Received initial orderbook snapshot")
                    self._initial_snapshot_received = True
                    asyncio.create_task(self._handle_initial_snapshot(data))
                else:
                    # It's an update with full orderbook state
                    asyncio.create_task(self._handle_orderbook_update(data))

            elif "events" in data:
                # Incremental update with events
                asyncio.create_task(self._handle_orderbook_update(data))

            else:
                logger.warning(
                    f"Unknown message type: {msg_type}, data keys: {list(data.keys())}"
                )

        except json.JSONDecodeError as e:
            logger.error(f"Failed to parse JSON message: {e}")
            asyncio.create_task(
                self._invoke_error_callback(ValueError(f"Invalid JSON: {e}"))
            )

        except Exception as e:
            logger.error(f"Error handling message: {e}")
            asyncio.create_task(self._invoke_error_callback(e))

    def _handle_subscription_response(self, data: dict) -> None:
        """
        Handle subscription response message.

        Args:
            data: Parsed JSON message
        """
        try:
            response = SubscriptionResponse(
                type=data.get("type"),
                status=data.get("status"),
                message=data.get("message"),
                data=data.get("data"),
            )

            logger.info(
                f"Subscription response: status={response.status}, message={response.message}"
            )

            if response.status == "success":
                self._subscribed = True

                # If initial data provided, process it as snapshot
                if response.data:
                    logger.info(
                        "Processing initial orderbook snapshot from subscription response"
                    )
                    asyncio.create_task(self._handle_initial_snapshot(response.data))

            elif response.status == "pending":
                logger.info("Subscription pending")

            else:
                error_msg = f"Subscription failed: {response.message}"
                logger.error(error_msg)
                asyncio.create_task(self._invoke_error_callback(ValueError(error_msg)))

        except Exception as e:
            logger.error(f"Error handling subscription response: {e}")
            asyncio.create_task(self._invoke_error_callback(e))

    async def _handle_initial_snapshot(self, data: dict) -> None:
        """
        Handle initial orderbook snapshot.

        Args:
            data: FrontendOrderbookData dict
        """
        try:
            # Parse bids and asks
            bids = [
                (self._parse_big_int(bid[0]), self._parse_big_int(bid[1]))
                for bid in data.get("b", [])
            ]
            asks = [
                (self._parse_big_int(ask[0]), self._parse_big_int(ask[1]))
                for ask in data.get("a", [])
            ]

            # Parse vault params if present
            vault_params = None
            if "vp" in data and data["vp"]:
                vault_params = self._parse_vault_params(data["vp"])

            # Create update with empty events for initial snapshot
            update = FrontendOrderbookUpdate(
                events=[],
                b=bids,
                a=asks,
                v=vault_params,
            )

            # Put update on queue
            await self._update_queue.put(update)

        except Exception as e:
            logger.error(f"Error handling initial snapshot: {e}")
            asyncio.create_task(self._invoke_error_callback(e))

    async def _handle_orderbook_update(self, data: dict) -> None:
        """
        Handle incremental orderbook update.

        Args:
            data: FrontendOrderbookUpdate dict
        """
        try:
            # Parse events
            events = []
            if "events" in data:
                for event_data in data["events"]:
                    events.append(self._parse_frontend_event(event_data))

            # Parse optional fields
            bids = None
            if "b" in data and data["b"]:
                bids = [
                    (self._parse_big_int(bid[0]), self._parse_big_int(bid[1]))
                    for bid in data["b"]
                ]

            asks = None
            if "a" in data and data["a"]:
                asks = [
                    (self._parse_big_int(ask[0]), self._parse_big_int(ask[1]))
                    for ask in data["a"]
                ]

            vault_params = None
            if "v" in data and data["v"]:
                vault_params = self._parse_vault_params(data["v"])

            # Create update
            update = FrontendOrderbookUpdate(
                events=events,
                b=bids,
                a=asks,
                v=vault_params,
            )

            # Put update on queue
            await self._update_queue.put(update)

        except Exception as e:
            logger.error(f"Error handling orderbook update: {e}")
            asyncio.create_task(self._invoke_error_callback(e))

    def _parse_big_int(self, value: Any) -> int:
        """
        Parse BigInt value from various formats to Python int.

        Handles:
        - int: Return as-is
        - str: Parse as decimal or hexadecimal integer (supports '0x' prefix)
        - None: Return 0
        - Other: Attempt int() conversion

        Args:
            value: Value to parse

        Returns:
            Python int

        Raises:
            ValueError: If value cannot be parsed as integer
        """
        if value is None:
            return 0

        if isinstance(value, int):
            return value

        if isinstance(value, str):
            try:
                # Remove any whitespace
                value_stripped = value.strip()

                # Check if it's a hexadecimal string (starts with '0x')
                if value_stripped.startswith('0x') or value_stripped.startswith('0X'):
                    return int(value_stripped, 16)
                else:
                    # Parse as decimal
                    return int(value_stripped)
            except ValueError as e:
                logger.error(f"Failed to parse string as int: '{value}'")
                raise ValueError(f"Invalid integer string: {value}") from e

        # Attempt generic int() conversion
        try:
            return int(value)
        except (ValueError, TypeError) as e:
            logger.error(f"Failed to convert {type(value)} to int: {value}")
            raise ValueError(f"Cannot convert {type(value).__name__} to int") from e

    def _parse_market_params(self, data: dict) -> MarketParams:
        """
        Convert dict to MarketParams dataclass.

        Args:
            data: Dict with market params

        Returns:
            MarketParams dataclass
        """
        return MarketParams(
            price_precision=self._parse_big_int(data.get("price_precision", 0)),
            size_precision=self._parse_big_int(data.get("size_precision", 0)),
            base_asset=data.get("base_asset", ""),
            base_asset_decimals=self._parse_big_int(data.get("base_asset_decimals", 0)),
            quote_asset=data.get("quote_asset", ""),
            quote_asset_decimals=self._parse_big_int(data.get("quote_asset_decimals", 0)),
            tick_size=self._parse_big_int(data.get("tick_size", 0)),
            min_size=self._parse_big_int(data.get("min_size", 0)),
            max_size=self._parse_big_int(data.get("max_size", 0)),
            taker_fee_bps=self._parse_big_int(data.get("taker_fee_bps", 0)),
            maker_fee_bps=self._parse_big_int(data.get("maker_fee_bps", 0)),
        )

    def _parse_vault_params(self, data: dict) -> VaultParams:
        """
        Convert dict to VaultParams dataclass.

        Args:
            data: Dict with vault params

        Returns:
            VaultParams dataclass
        """
        return VaultParams(
            vault_best_bid=self._parse_big_int(data.get("vault_best_bid", 0)),
            bid_partially_filled_size=self._parse_big_int(
                data.get("bid_partially_filled_size", 0)
            ),
            vault_best_ask=self._parse_big_int(data.get("vault_best_ask", 0)),
            ask_partially_filled_size=self._parse_big_int(
                data.get("ask_partially_filled_size", 0)
            ),
            vault_bid_order_size=self._parse_big_int(
                data.get("vault_bid_order_size", 0)
            ),
            vault_ask_order_size=self._parse_big_int(
                data.get("vault_ask_order_size", 0)
            ),
            spread=self._parse_big_int(data.get("spread", 0)),
        )

    def _parse_frontend_event(self, data: dict) -> FrontendEvent:
        """
        Convert dict to FrontendEvent dataclass.

        Args:
            data: Dict with event data

        Returns:
            FrontendEvent dataclass
        """
        return FrontendEvent(
            e=data["e"],  # Required
            ts=self._parse_big_int(data["ts"]),  # Required - handle hex timestamps
            mad=data["mad"],  # Required
            th=data.get("th"),
            p=(
                self._parse_big_int(data["p"])
                if "p" in data and data["p"] is not None
                else None
            ),
            s=(
                self._parse_big_int(data["s"])
                if "s" in data and data["s"] is not None
                else None
            ),
            ib=data.get("ib"),
            t=data.get("t"),
            m=data.get("m"),
        )

    async def _heartbeat_monitor(self) -> None:
        """
        Monitor connection health using websockets' built-in ping/pong.

        This monitor checks if the connection is alive by relying on
        websockets library's automatic ping/pong mechanism. If the
        connection becomes stale, trigger reconnection.
        """
        while self._connected and not self._closing:
            try:
                await asyncio.sleep(self._heartbeat_interval)

                # Check if websocket is still open
                if self._websocket is None:
                    logger.warning("WebSocket is None in heartbeat monitor")
                    asyncio.create_task(self._handle_connection_loss())
                    break

                # websockets library handles ping/pong automatically
                # If connection is stale, it will raise an exception
                # We just need to catch it in the message loop

            except asyncio.CancelledError:
                logger.debug("Heartbeat monitor cancelled")
                break
            except Exception as e:
                logger.error(f"Heartbeat monitor error: {e}")
                await self._invoke_error_callback(e)
                asyncio.create_task(self._handle_connection_loss())
                break

    async def _handle_connection_loss(self) -> None:
        """
        Handle detected connection loss.

        This method:
        1. Logs the connection loss
        2. Cleans up current connection
        3. Initiates reconnection
        """
        if self._closing:
            return

        logger.warning("Connection loss detected")

        async with self._lock:
            self._connected = False

        await self._invoke_error_callback(ConnectionError("Connection lost"))

        # Initiate reconnection
        asyncio.create_task(self._reconnect())

    async def _invoke_error_callback(self, error: Exception) -> None:
        """
        Invoke the on_error callback if provided.

        Args:
            error: Exception to pass to callback
        """
        if self._on_error is None:
            return

        try:
            if asyncio.iscoroutinefunction(self._on_error):
                await self._on_error(error)
            else:
                loop = asyncio.get_event_loop()
                await loop.run_in_executor(None, self._on_error, error)

        except Exception as e:
            logger.error(f"Error in on_error callback: {e}", exc_info=True)

    async def __aenter__(self) -> "KuruFrontendOrderbookClient":
        """
        Async context manager entry.

        Automatically connects when entering the context.

        Returns:
            Self for use in the context
        """
        await self.connect()
        return self

    async def __aexit__(
        self,
        exc_type: Optional[type],
        exc_val: Optional[Exception],
        exc_tb: Optional[Any],
    ) -> bool:
        """
        Async context manager exit.

        Automatically closes connection when exiting the context.

        Args:
            exc_type: Exception type if an exception occurred
            exc_val: Exception value if an exception occurred
            exc_tb: Exception traceback if an exception occurred

        Returns:
            False to propagate any exceptions
        """
        await self.close()
        return False  # Propagate exceptions
