from __future__ import annotations

import json
from dataclasses import dataclass

from websockets.sync.client import ClientConnection, connect


@dataclass
class Message:
    type: str
    payload: object


class WsClient:
    def __init__(self, conn: ClientConnection):
        self._conn = conn

    def read_message(self) -> Message:
        message = self._conn.recv()
        if isinstance(message, bytes):
            message = message.decode("utf-8")

        data = json.loads(message)
        if not isinstance(data, dict) or "type" not in data:
            raise ValueError("invalid message envelope")

        return Message(type=str(data["type"]), payload=data.get("payload"))

    def send(self, msg_type: str, payload: object) -> None:
        envelope = {"type": msg_type, "payload": payload}
        self._conn.send(json.dumps(envelope))

    def close(self) -> None:
        self._conn.close()


def connect_ws(server_url: str, token: str) -> WsClient:
    headers: dict[str, str] = {}
    if token:
        headers["Authorization"] = f"Bearer {token}"

    conn = connect(server_url, additional_headers=headers)
    return WsClient(conn)
