import json
import logging
import websockets
import asyncio
from typing import Optional, Callable, Dict, Any

class AxiomTradeWebSocketClient:    
    def __init__(self, auth_token=None, refresh_token=None, log_level=logging.INFO, auth_manager=None) -> None:
        self.ws_url = "wss://cluster3.axiom.trade/"
        self.ws_url_token_price = "wss://socket8.axiom.trade/"
        self.ws: Optional[websockets.WebSocketClientProtocol] = None
        self.auth_token = auth_token
        self.refresh_token = refresh_token
        self.auth_manager = auth_manager
        
        # Setup logging
        self.logger = logging.getLogger("AxiomTradeWebSocket")
        self.logger.setLevel(log_level)
        
        # Create console handler if none exists
        if not self.logger.handlers:
            handler = logging.StreamHandler()
            handler.setLevel(log_level)
            formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
            handler.setFormatter(formatter)
            self.logger.addHandler(handler)
        
        self._callbacks: Dict[str, Callable] = {}

    async def connect(self, is_token_price: bool = False) -> bool:
        """Connect to the WebSocket server."""
        headers = {
            'Origin': 'https://axiom.trade',
            'Cache-Control': 'no-cache',
            'Accept-Language': 'en-US,en;q=0.9',
            'Pragma': 'no-cache',
            'Connection': 'Upgrade', # Usually handled by the library
            'Upgrade': 'websocket', # Usually handled by the library
            'Sec-WebSocket-Key': '0/XbexFur6y5gruQBvuI/Q==', # Generated by the library
            'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/133.0.0.0 Safari/537.36 OPR/118.0.0.0',
            'Sec-WebSocket-Version': '13', # Handled by the library
            'Sec-WebSocket-Extensions': 'permessage-deflate; client_max_window_bits' # Negotiated by the library
        }
        
        if self.auth_token:
            headers["Cookie"] = f"auth-access-token={self.auth_token}"
            if self.refresh_token:
                headers["Cookie"] += f"; auth-refresh-token={self.refresh_token}"
        
        try:
            if is_token_price:
                self.ws_url = self.ws_url_token_price
            self.ws = await websockets.connect(
                self.ws_url,
                extra_headers=headers
            )
            self.logger.info("Connected to WebSocket server")
            return True
        except Exception as e:
            self.logger.error(f"Failed to connect to WebSocket: {e}")
            return False

    async def subscribe_new_tokens(self, callback: Callable[[Dict[str, Any]], None]):
        """Subscribe to new token updates."""
        if not self.ws:
            if not await self.connect():
                return False

        self._callbacks["new_pairs"] = callback
        
        try:
            await self.ws.send(json.dumps({
                "action": "join",
                "room": "new_pairs"
            }))
            self.logger.info("Subscribed to new token updates")
            return True
        except Exception as e:
            self.logger.error(f"Failed to subscribe to new tokens: {e}")
            return False

    async def subscribe_token_price(self, token: str, callback: Callable[[Dict[str, Any]], None]):
        """Subscribe to token price updates."""
        if not self.ws:
            if not await self.connect(is_token_price=True):
                return False

        self._callbacks[token] = callback
        
        try:
            await self.ws.send(json.dumps({
                "action": "join",
                "room": token
            }))
            self.logger.info(f"Subscribed to price updates for token: {token}")
            return True
        except Exception as e:
            self.logger.error(f"Failed to subscribe to token price: {e}")
            return False
    async def _message_handler(self):
        """Handle incoming WebSocket messages."""
        if not self.ws:
            return

        try:
            while True:
                message = await self.ws.recv()
                data = json.loads(message)
                
                if "room" in data and data["room"] in self._callbacks:
                    callback = self._callbacks[data["room"]]
                    await callback(data["content"])
                
        except websockets.exceptions.ConnectionClosed:
            self.logger.warning("WebSocket connection closed")
        except Exception as e:
            self.logger.error(f"Error handling message: {e}")

    async def start(self):
        """Start the WebSocket client and message handler."""
        if not self.ws:
            if not await self.connect():
                return
        
        await self._message_handler()

    async def close(self):
        """Close the WebSocket connection."""
        if self.ws:
            await self.ws.close()
            self.ws = None
            self.logger.info("WebSocket connection closed")
