import threading
import logging
from binance import ThreadedWebsocketManager

logger = logging.getLogger(__name__)

class BinanceWSClient:
    def __init__(self, api_key: str, api_secret: str):
        self.api_key = api_key
        self.api_secret = api_secret
        self.twm = None
        self.prices: dict[str, float] = {}
        self._lock = threading.Lock()

    def start(self, symbols: list[str]):
        self.twm = ThreadedWebsocketManager(
            api_key=self.api_key, api_secret=self.api_secret
        )
        self.twm.start()
        streams = [s.lower() + "@miniTicker" for s in symbols]
        self.twm.start_multiplex_socket(
            callback=self._handle_message, streams=streams
        )
        logger.info(f"WebSocket started for {len(symbols)} symbols")

    def _handle_message(self, msg):
        if msg.get("e") == "error":
            logger.error(f"WebSocket error: {msg}")
            return
        data = msg.get("data", msg)
        if "s" in data and "c" in data:
            with self._lock:
                self.prices[data["s"]] = float(data["c"])

    def get_price(self, symbol: str) -> float | None:
        with self._lock:
            return self.prices.get(symbol)

    def get_all_prices(self) -> dict[str, float]:
        with self._lock:
            return dict(self.prices)

    def update_symbols(self, symbols: list[str]):
        if self.twm:
            self.stop()
            self.start(symbols)

    def stop(self):
        if self.twm:
            self.twm.stop()
            logger.info("WebSocket stopped")