import base64
from enum import Enum
from typing import Optional
from .model import TypedModel
from .transcriber import TranscriberConfig
from .agent import AgentConfig
from .synthesizer import SynthesizerConfig


class WebSocketMessageType(str, Enum):
    BASE = "websocket_base"
    START = "websocket_start"
    AUDIO = "websocket_audio"
    READY = "websocket_ready"
    STOP = "websocket_stop"


class WebSocketMessage(TypedModel, type=WebSocketMessageType.BASE):
    pass


class AudioMessage(WebSocketMessage, type=WebSocketMessageType.AUDIO):
    data: str

    @classmethod
    def from_bytes(cls, chunk: bytes):
        return cls(data=base64.b64encode(chunk).decode("utf-8"))

    def get_bytes(self) -> bytes:
        return base64.b64decode(self.data)


class StartMessage(WebSocketMessage, type=WebSocketMessageType.START):
    transcriber_config: TranscriberConfig
    agent_config: AgentConfig
    synthesizer_config: SynthesizerConfig
    conversation_id: Optional[str] = None


class ReadyMessage(WebSocketMessage, type=WebSocketMessageType.READY):
    pass


class StopMessage(WebSocketMessage, type=WebSocketMessageType.STOP):
    pass