"""
WebSocket support for real-time updates in Python SDK
"""

import asyncio
import json
from typing import Optional, Callable, Dict, Any
from enum import Enum
import websockets
from websockets.client import WebSocketClientProtocol


class WebSocketEvent(str, Enum):
    CONNECTED = "connected"
    DISCONNECTED = "disconnected"
    ERROR = "error"
    MESSAGE = "message"
    APPROVAL_PENDING = "approval_pending"
    APPROVAL_APPROVED = "approval_approved"
    APPROVAL_DENIED = "approval_denied"
    POLICY_UPDATED = "policy_updated"
    TRACE_CREATED = "trace_created"


class AegisWebSocket:
    """WebSocket client for real-time Aegis updates"""

    def __init__(
        self,
        url: str,
        api_key: Optional[str] = None,
        reconnect: bool = True,
        reconnect_interval: int = 5,
        max_reconnect_attempts: int = 10,
    ):
        self.url = url
        self.api_key = api_key
        self.reconnect = reconnect
        self.reconnect_interval = reconnect_interval
        self.max_reconnect_attempts = max_reconnect_attempts
        
        self.ws: Optional[WebSocketClientProtocol] = None
        self.reconnect_attempts = 0
        self.handlers: Dict[str, list] = {}
        self._running = False

    def on(self, event: str, handler: Callable):
        """Register event handler"""
        if event not in self.handlers:
            self.handlers[event] = []
        self.handlers[event].append(handler)

    def emit(self, event: str, data: Any = None):
        """Emit event to handlers"""
        if event in self.handlers:
            for handler in self.handlers[event]:
                handler(data)

    async def connect(self):
        """Connect to WebSocket server"""
        url = self.url
        if self.api_key:
            url += f"?apiKey={self.api_key}"

        try:
            self.ws = await websockets.connect(url)
            self.reconnect_attempts = 0
            self.emit(WebSocketEvent.CONNECTED)
            self._running = True
            
            # Start message loop
            await self._message_loop()
            
        except Exception as e:
            self.emit(WebSocketEvent.ERROR, e)
            if self.reconnect and self.reconnect_attempts < self.max_reconnect_attempts:
                self.reconnect_attempts += 1
                await asyncio.sleep(self.reconnect_interval)
                await self.connect()

    async def disconnect(self):
        """Disconnect from WebSocket server"""
        self._running = False
        if self.ws:
            await self.ws.close()
            self.ws = None
            self.emit(WebSocketEvent.DISCONNECTED)

    async def send(self, message_type: str, data: Any):
        """Send message to server"""
        if not self.ws:
            raise RuntimeError("WebSocket is not connected")

        message = {
            "type": message_type,
            "data": data,
            "timestamp": asyncio.get_event_loop().time(),
        }
        
        await self.ws.send(json.dumps(message))

    async def subscribe_to_approvals(self):
        """Subscribe to approval updates"""
        await self.send("subscribe", {"channel": "approvals"})

    async def subscribe_to_policy_updates(self):
        """Subscribe to policy updates"""
        await self.send("subscribe", {"channel": "policy"})

    async def subscribe_to_traces(self):
        """Subscribe to trace updates"""
        await self.send("subscribe", {"channel": "traces"})

    async def _message_loop(self):
        """Main message processing loop"""
        try:
            async for message in self.ws:
                try:
                    data = json.loads(message)
                    self._handle_message(data)
                except json.JSONDecodeError as e:
                    self.emit(WebSocketEvent.ERROR, e)
        except websockets.exceptions.ConnectionClosed:
            self.emit(WebSocketEvent.DISCONNECTED)
            if self.reconnect and self.reconnect_attempts < self.max_reconnect_attempts:
                self.reconnect_attempts += 1
                await asyncio.sleep(self.reconnect_interval)
                await self.connect()

    def _handle_message(self, message: Dict[str, Any]):
        """Handle incoming message"""
        self.emit(WebSocketEvent.MESSAGE, message)

        # Emit specific events based on message type
        msg_type = message.get("type")
        if msg_type == "approval_pending":
            self.emit(WebSocketEvent.APPROVAL_PENDING, message.get("data"))
        elif msg_type == "approval_approved":
            self.emit(WebSocketEvent.APPROVAL_APPROVED, message.get("data"))
        elif msg_type == "approval_denied":
            self.emit(WebSocketEvent.APPROVAL_DENIED, message.get("data"))
        elif msg_type == "policy_updated":
            self.emit(WebSocketEvent.POLICY_UPDATED, message.get("data"))
        elif msg_type == "trace_created":
            self.emit(WebSocketEvent.TRACE_CREATED, message.get("data"))

    def is_connected(self) -> bool:
        """Check if connected"""
        return self.ws is not None and not self.ws.closed
