"""
Base WebSocket Client — shared connection management for all WS clients.

Provides:
- Persistent connection with auto-reconnect
- Exponential backoff capped at 60s
- Request/response correlation via pending futures
- Background receive and ping loops
- Graceful disconnect with task cleanup
- Statistics tracking

Subclasses implement:
- _get_ws_url() → str: WebSocket endpoint URL
- _get_connect_headers() → dict: Auth/metadata headers
- _on_message(data: dict) → None: Handle incoming messages
- (optional) _on_connected() → None: Post-connect setup
"""

import asyncio
import json
import logging
import time
from abc import ABC, abstractmethod
from datetime import datetime
from enum import Enum
from typing import Any, Dict

logger = logging.getLogger("aigie.realtime.base_ws")


class ConnectionState(Enum):
    """WebSocket connection state."""

    DISCONNECTED = "disconnected"
    CONNECTING = "connecting"
    CONNECTED = "connected"
    RECONNECTING = "reconnecting"
    FAILED = "failed"


class BaseWebSocketClient(ABC):
    """
    Base class for WebSocket clients with auto-reconnect and correlation.

    Subclass this to create domain-specific WebSocket clients.
    """

    MAX_BACKOFF_SEC = 60.0  # Hard cap on reconnect backoff

    def __init__(
        self,
        api_url: str,
        api_key: str,
        *,
        ws_path: str = "/ws",
        auto_reconnect: bool = True,
        reconnect_interval: float = 2.0,
        max_reconnect_attempts: int = 10,
        ping_interval: float = 30.0,
        connect_timeout: float = 10.0,
    ):
        # Convert HTTP URL to WebSocket URL
        ws_url = api_url.replace("http://", "ws://").replace("https://", "wss://")
        ws_url = ws_url.rstrip("/api").rstrip("/")
        self._ws_base_url = f"{ws_url}{ws_path}"

        self._api_key = api_key
        self._auto_reconnect = auto_reconnect
        self._reconnect_interval = reconnect_interval
        self._max_reconnect_attempts = max_reconnect_attempts
        self._ping_interval = ping_interval
        self._connect_timeout = connect_timeout

        # Connection state
        self._state = ConnectionState.DISCONNECTED
        self._websocket = None
        self._reconnect_attempts = 0
        self._last_connect_time: datetime | None = None

        # Request/response correlation
        self._pending_requests: Dict[str, asyncio.Future] = {}

        # Background tasks
        self._receive_task: asyncio.Task | None = None
        self._ping_task: asyncio.Task | None = None
        self._reconnect_task: asyncio.Task | None = None

        # Lock for connect/disconnect
        self._lock = asyncio.Lock()

        # Stats
        self._stats = {
            "connections_attempted": 0,
            "connections_successful": 0,
            "connections_failed": 0,
            "reconnections": 0,
            "messages_sent": 0,
            "messages_received": 0,
            "timeouts": 0,
            "errors": 0,
        }

    # ─── Properties ──────────────────────────────────────────────────

    @property
    def is_connected(self) -> bool:
        return self._state == ConnectionState.CONNECTED and self._websocket is not None

    @property
    def state(self) -> ConnectionState:
        return self._state

    # ─── Abstract methods for subclasses ─────────────────────────────

    def _get_ws_url(self) -> str:
        """Return the WebSocket URL to connect to. Override for custom URL building."""
        return self._ws_base_url

    def _get_connect_headers(self) -> Dict[str, str]:
        """Return headers/query params for WebSocket handshake."""
        return {
            "X-API-Key": self._api_key,
            "X-Client-Type": "sdk",
            "X-Client-Version": "2.0.0",
        }

    @abstractmethod
    async def _on_message(self, data: Dict[str, Any]) -> None:
        """Handle an incoming WebSocket message. Subclass must implement."""
        ...

    async def _on_connected(self) -> None:
        """Called after successful connection. Override for subscriptions etc."""
        pass

    # ─── Connection lifecycle ────────────────────────────────────────

    async def connect(self) -> bool:
        """Establish WebSocket connection."""
        if self._state == ConnectionState.CONNECTED:
            return True

        async with self._lock:
            if self._state == ConnectionState.CONNECTED:
                return True

            self._state = ConnectionState.CONNECTING
            self._stats["connections_attempted"] += 1

            try:
                import websockets
            except ImportError:
                logger.warning("websockets package not installed")
                self._state = ConnectionState.FAILED
                self._stats["connections_failed"] += 1
                return False

            target_url = self._get_ws_url()
            logger.info(f"Connecting to WebSocket: {target_url}")

            try:
                self._websocket = await asyncio.wait_for(
                    websockets.connect(
                        target_url,
                        extra_headers=self._get_connect_headers(),
                        ping_interval=None,  # We handle pings ourselves
                        ping_timeout=10,
                        close_timeout=5,
                    ),
                    timeout=self._connect_timeout,
                )

                self._state = ConnectionState.CONNECTED
                self._reconnect_attempts = 0
                self._last_connect_time = datetime.utcnow()
                self._stats["connections_successful"] += 1

                logger.info("WebSocket connected successfully")

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

                # Post-connect hook
                await self._on_connected()

                return True

            except asyncio.TimeoutError:
                logger.error(f"WebSocket connection timeout after {self._connect_timeout}s")
                self._state = ConnectionState.FAILED
                self._stats["connections_failed"] += 1
                self._schedule_reconnect()
                return False

            except Exception as e:
                logger.error(f"WebSocket connection failed: {e}")
                self._state = ConnectionState.FAILED
                self._stats["connections_failed"] += 1
                self._schedule_reconnect()
                return False

    async def disconnect(self) -> None:
        """Close connection and cancel background tasks."""
        async with self._lock:
            self._state = ConnectionState.DISCONNECTED

            # Cancel background tasks
            for task in [self._receive_task, self._ping_task, self._reconnect_task]:
                if task and not task.done():
                    task.cancel()
                    try:
                        await task
                    except asyncio.CancelledError:
                        pass

            self._receive_task = None
            self._ping_task = None
            self._reconnect_task = None

            # Close WebSocket
            if self._websocket:
                try:
                    await self._websocket.close()
                except Exception:
                    pass
                self._websocket = None

            # Cancel pending requests
            for future in self._pending_requests.values():
                if not future.done():
                    future.cancel()
            self._pending_requests.clear()

            logger.info("WebSocket disconnected")

    # ─── Send helpers ────────────────────────────────────────────────

    async def _send(self, message: Dict[str, Any]) -> bool:
        """Send a JSON message. Returns True if sent."""
        if not self._websocket or not self.is_connected:
            return False
        try:
            await self._websocket.send(json.dumps(message))
            self._stats["messages_sent"] += 1
            return True
        except Exception as e:
            logger.warning(f"Send failed: {e}")
            return False

    async def _send_and_wait(
        self,
        request_id: str,
        message: Dict[str, Any],
        timeout: float,
    ) -> Dict[str, Any] | None:
        """Send a message and wait for a correlated response."""
        if not self._websocket or not self.is_connected:
            return None

        self._stats["messages_sent"] += 1
        start_time = time.perf_counter()

        try:
            future = asyncio.get_event_loop().create_future()
            self._pending_requests[request_id] = future

            await self._websocket.send(json.dumps(message))
            response = await asyncio.wait_for(future, timeout=timeout)
            return response

        except asyncio.TimeoutError:
            self._stats["timeouts"] += 1
            logger.debug(f"Request {request_id} timed out after {timeout}s")
            return None

        except Exception as e:
            self._stats["errors"] += 1
            logger.error(f"Request {request_id} failed: {e}")
            return None

        finally:
            self._pending_requests.pop(request_id, None)

    def _resolve_request(self, request_id: str, data: Dict[str, Any]) -> None:
        """Resolve a pending request future with response data."""
        if request_id in self._pending_requests:
            future = self._pending_requests[request_id]
            if not future.done():
                future.set_result(data)

    # ─── Background loops ────────────────────────────────────────────

    async def _receive_loop(self) -> None:
        """Background task to receive messages."""
        while self._state == ConnectionState.CONNECTED and self._websocket:
            try:
                message = await self._websocket.recv()
                self._stats["messages_received"] += 1
                data = json.loads(message)
                await self._on_message(data)

            except asyncio.CancelledError:
                break
            except Exception as e:
                if self._state == ConnectionState.CONNECTED:
                    if "close frame" not in str(e).lower():
                        logger.error(f"Receive error: {e}")
                    await self._handle_disconnect()
                break

    async def _ping_loop(self) -> None:
        """Background keepalive pings."""
        while self._state == ConnectionState.CONNECTED and self._websocket:
            try:
                await asyncio.sleep(self._ping_interval)
                if self._websocket and self._state == ConnectionState.CONNECTED:
                    await self._websocket.send(
                        json.dumps(
                            {
                                "type": "ping",
                                "timestamp": datetime.utcnow().isoformat(),
                            }
                        )
                    )
                    self._stats["messages_sent"] += 1
            except asyncio.CancelledError:
                break
            except Exception:
                break

    # ─── Reconnection ────────────────────────────────────────────────

    async def _handle_disconnect(self) -> None:
        """Handle unexpected disconnection."""
        self._state = ConnectionState.DISCONNECTED
        self._schedule_reconnect()

    def _schedule_reconnect(self) -> None:
        """Schedule a reconnect attempt."""
        if not self._auto_reconnect:
            return
        if self._reconnect_attempts >= self._max_reconnect_attempts:
            return
        if self._reconnect_task is None or self._reconnect_task.done():
            self._reconnect_task = asyncio.create_task(self._reconnect())

    async def _reconnect(self) -> None:
        """Reconnect with capped exponential backoff."""
        while self._reconnect_attempts < self._max_reconnect_attempts:
            self._state = ConnectionState.RECONNECTING
            self._reconnect_attempts += 1
            self._stats["reconnections"] += 1

            wait_time = min(
                self._reconnect_interval * (2 ** (self._reconnect_attempts - 1)),
                self.MAX_BACKOFF_SEC,
            )
            logger.info(
                f"Reconnecting in {wait_time:.1f}s "
                f"(attempt {self._reconnect_attempts}/{self._max_reconnect_attempts})"
            )

            await asyncio.sleep(wait_time)

            if await self.connect():
                logger.info("Reconnection successful")
                return

        logger.error(f"Max reconnection attempts ({self._max_reconnect_attempts}) reached")
        self._state = ConnectionState.FAILED

    # ─── Stats ───────────────────────────────────────────────────────

    def get_stats(self) -> Dict[str, Any]:
        """Get connection statistics."""
        return {
            **self._stats,
            "state": self._state.value,
            "pending_requests": len(self._pending_requests),
            "reconnect_attempts": self._reconnect_attempts,
            "last_connect_time": self._last_connect_time.isoformat()
            if self._last_connect_time
            else None,
        }

    # ─── Context manager ─────────────────────────────────────────────

    async def __aenter__(self):
        await self.connect()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.disconnect()
