from typing_extensions import TypedDict, Any
from dataclasses import dataclass, field
import orjson

from mexc.core import validator, ValidationError
from .streams_rpc import StreamsRPCSocketClient, Message
from .proto import PushDataV3ApiWrapper

MEXC_SPOT_SOCKET_URL = 'wss://wbs-api.mexc.com/ws'

class Response(TypedDict):
  id: int
  code: int
  msg: str

validate_response = validator(Response)

@dataclass
class StreamsClient(StreamsRPCSocketClient[Any, Any, PushDataV3ApiWrapper]):
  url: str = field(default=MEXC_SPOT_SOCKET_URL, kw_only=True)

  async def send(self, msg):
    await (await self.ws).send(orjson.dumps(msg), text=True)

  async def send_request(self, method: str, params=None):
    msg = {'method': method}
    if params is not None:
      msg['params'] = params
    await self.send(msg)

  async def ping(self):
    await self.send('PING')

  async def request_subscription(self, channel: str, params=None):
    await self.send_request('SUBSCRIPTION', [channel])

  async def request_unsubscription(self, channel: str):
    await self.send_request('UNSUBSCRIPTION', [channel])

  def parse_msg(self, msg: bytes) -> Message[Any, PushDataV3ApiWrapper]:
    try:
      data = validate_response(msg)
      return {'kind': 'response', 'response': data}
    except ValidationError:
      proto = PushDataV3ApiWrapper.parse(msg)
      return {'kind': 'subscription', 'channel': proto.channel, 'data': proto}
    
@dataclass
class StreamsMixin:
  ws: StreamsClient = field(default_factory=StreamsClient, kw_only=True)

  async def __aenter__(self):
    await self.ws.__aenter__()
    return self
  
  async def __aexit__(self, exc_type, exc_value, traceback):
    await self.ws.__aexit__(exc_type, exc_value, traceback)

  def subscribe(self, channel: str):
    return self.ws.subscribe(channel)