"""
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, timezone
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,
    ):
        if not api_url:
            raise ValueError("api_url is required for WebSocket connection")
        if not api_key:
            raise ValueError("api_key is required for WebSocket connection")

        # Convert HTTP URL to WebSocket URL
        ws_url = api_url.replace("http://", "ws://").replace("https://", "wss://")
        ws_url = ws_url.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,
                        additional_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.now(timezone.utc)
                self._stats["connections_successful"] += 1

                logger.info("WebSocket connected successfully")

                # Post-connect hook (before receive loop so subclass can recv welcome)
                await self._on_connected()

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

                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:
                error_msg = str(e)
                # Provide actionable guidance for common WebSocket errors
                try:
                    import websockets

                    if isinstance(e, websockets.exceptions.InvalidStatusCode):
                        status = getattr(e, "status_code", None)
                        if status == 403:
                            error_msg = (
                                f"WebSocket connection rejected (403 Forbidden). "
                                f"Check: 1) API key is valid, "
                                f"2) Ingress/WAF allows WebSocket upgrades on {target_url}, "
                                f"3) CORS/origin headers are permitted"
                            )
                        elif status == 401:
                            error_msg = (
                                "WebSocket authentication failed (401). "
                                "Verify your API key is correct and not expired."
                            )
                except (ImportError, AttributeError):
                    pass
                logger.error(f"WebSocket connection failed: {error_msg}")
                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_running_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.now(timezone.utc).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()
