import asyncio
import logging
from typing import Awaitable, Callable, Literal, TypedDict
from websockets.exceptions import ConnectionClosedError
from websockets.asyncio.server import serve, Server, ServerConnection
import json
import traceback as tb

logger = logging.getLogger("websocket")

class RecordingPayload(TypedDict):
  running: bool
class RecordingEvent(TypedDict):
  topic: Literal["recording"]
  payload: RecordingPayload

class RecordingStatusEvent(TypedDict):
  topic: Literal["recording-status"]

# inbound messages, coming from page
class WithPageId(TypedDict):
  pageId: str
class ClickEvent(WithPageId, TypedDict):
  topic: Literal["clickEvent"]
  elementTree: str

class InputEvent(WithPageId, TypedDict):
  topic: Literal["inputEvent"]
  elementTree: str

class KeyPressEvent(WithPageId, TypedDict):
  topic: Literal["keyPressEvent"]

class KeyReleaseEvent(WithPageId, TypedDict):
  topic: Literal["keyReleaseEvent"]

class ContextMenuEvent(WithPageId, TypedDict):
  topic: Literal["contextMenuEvent"]
  elementTree: str

class SnapshotPayload(TypedDict):
  id: str
  ariaTree: str

class Snapshot(WithPageId, TypedDict):
  topic: Literal["snapshot"]
  payload: SnapshotPayload

InboundWebsocketMessage = ClickEvent | KeyPressEvent | KeyReleaseEvent | Snapshot | InputEvent | RecordingEvent | RecordingStatusEvent | ContextMenuEvent

# outbound messages, sent to page
class UpdateUserInteractionPayload(TypedDict):
  disabled: bool
class UpdateUserInteraction(TypedDict):
  topic: Literal["interaction"]
  payload: UpdateUserInteractionPayload

class SnapshotRequestPayload(TypedDict):
  id: str

class SnapshotRequest(TypedDict):
  topic: Literal["snapshot-request"]
  page_uuid: str
  payload: SnapshotRequestPayload

OutboundWebsocketMessage = UpdateUserInteraction | SnapshotRequest | RecordingEvent

class WebsocketService:
  server: Server | None
  port: int
  
  def __init__(self, port: int = 12345):
    self.server = None
    self.port = port
    self.listeners: list[Callable[[InboundWebsocketMessage, ServerConnection], None | Awaitable[None]]] = []
    logger.debug(f"WebsocketService created on port {self.port}")

  async def _handle_connection(self, connection: ServerConnection):
    async def handle_result(result: Awaitable[None]):
      try:
        await result
      except Exception as e:
        logger.error(e)
        logger.error(tb.format_exc())

    try:
      async for message in connection:
        str_message: str
        if type(message) == bytes:
          str_message = message.decode()
        elif type(message) == str:
          str_message = message
        else:
          raise Exception(f"Unknown message type: {type(message)}")
        for listener in self.listeners:
          try:
            result = listener(json.loads(str_message), connection)
            if result:
              asyncio.ensure_future(handle_result(result))
          except Exception as e:
            print(f"Listener error: {e}")
            print(tb.format_exc())
    except ConnectionClosedError as e:
      logger.warning(f"Connection closed error: {e}")

  def addListener(self, listener: Callable[[InboundWebsocketMessage, ServerConnection], None | Awaitable[None]]):
    self.listeners.append(listener)
  def removeListener(self, listener: Callable[[InboundWebsocketMessage, ServerConnection], None | Awaitable[None]]):
    self.listeners.remove(listener)

  async def broadcast(self, message: OutboundWebsocketMessage | str):
    if not self.server or not self.server.is_serving():
      return
    if type(message) == str:
      message_str = message
    else:
      message_str = json.dumps(message)
    for conn in self.server.connections:
      await conn.send(message_str, text=False)


  async def __aenter__(self):
    logger.debug("Starting Websocket server...")
    self.server = await serve(self._handle_connection, "localhost", self.port, max_size=None)

    return self
  async def __aexit__(self, exc_type, exc, tb):
    if self.server:
      self.server.close()
      await self.server.wait_closed()

