import asyncio
import json
import websockets
import requests

class OKXPerpStream:
    def __init__(self, quote_filters, chunk_size=50):
        self.url = "wss://ws.okx.com:8443/ws/v5/public"
        self.quote_filters = quote_filters
        self.chunk_size = chunk_size
        self.symbol_chunks = self.get_symbol_chunks()

    def get_symbol_chunks(self):
        try:
            url = "https://www.okx.com/api/v5/public/instruments"
            params = {"instType": "SWAP"}
            data = requests.get(url, params=params).json()
            symbols = [item["instId"] for item in data["data"] if any(q in item["settleCcy"] for q in self.quote_filters)]
            print("[OKX] Starting stream with", len(symbols), "symbols")
            return [symbols[i:i+self.chunk_size] for i in range(0, len(symbols), self.chunk_size)]
        except Exception as e:
            print("[OKX REST ERROR]", e)
            return []

    def normalize_symbol(self, symbol):
        return symbol.replace("-SWAP", "").replace("-", "/")

    def parse_tick(self, data):
        symbol = data["instId"]
        return {
            "exchange": "okx",
            "symbol": self.normalize_symbol(symbol),
            "bid": float(data["bidPx"]),
            "ask": float(data["askPx"]),
            "bid_size": float(data["bidSz"]),
            "ask_size": float(data["askSz"])
        }

    async def stream_chunk(self, symbols):
        args = [{"channel": "tickers", "instId": s} for s in symbols]
        async with websockets.connect(self.url) as ws:
            await ws.send(json.dumps({"op": "subscribe", "args": args}))
            async for msg in ws:
                try:
                    data = json.loads(msg)
                    if "data" in data:
                        for item in data["data"]:
                            tick = self.parse_tick(item)
                            if tick:
                                yield tick
                except Exception as e:
                    print("[OKX PARSE ERROR]", e)

    async def stream(se