"""Router - handle message router (sock).

Router to manage responses from a socket client.

"""

from __future__ import annotations

from wandb.proto import wandb_internal_pb2 as pb
from wandb.proto import wandb_server_pb2 as spb
from wandb.sdk.lib.sock_client import SockClient, SockClientClosedError
from wandb.sdk.mailbox import Mailbox

from .router import MessageRouter, MessageRouterClosedError


class MessageSockRouter(MessageRouter):
    _sock_client: SockClient
    _mailbox: Mailbox

    def __init__(self, sock_client: SockClient, mailbox: Mailbox) -> None:
        self._sock_client = sock_client
        super().__init__(mailbox=mailbox)

    def _read_message(self) -> spb.ServerResponse | None:
        try:
            return self._sock_client.read_server_response(timeout=1)
        except SockClientClosedError as e:
            raise MessageRouterClosedError from e

    def _send_message(self, record: pb.Record) -> None:
        self._sock_client.send_record_communicate(record)