from pydantic import BaseModel, Field, TypeAdapter
from enum import Enum
from typing import Literal, Annotated, Union


__all__ = ['EServerMessage', 'MsgStatus', 'MsgError', 'MsgEventTextChunk', 'MsgEventEnd', 'AgentEventUnion', 'ServerMessage']


class EServerMessage(str, Enum):
    STATUS = "STATUS"
    ERROR = "ERROR"
    GRACEFUL_DISCONNECT = "GRACEFUL_DISCONNECT"
    AGENT_EVENT_TEXT_CHUNK = "AGENT_EVENT_TEXT_CHUNK"
    AGENT_EVENT_END = "AGENT_EVENT_END"


class MsgStatus(BaseModel):
    """
    Отправляется сервером при открытии соединения с клиентом.
    Будет дополнен информацией о готовности агента принимать сообщения.
    """
    type: Literal[EServerMessage.STATUS] = EServerMessage.STATUS


class MsgEventTextChunk(BaseModel):
    """
    Чанк текста ответа агента.
    """
    type: Literal[EServerMessage.AGENT_EVENT_TEXT_CHUNK] = EServerMessage.AGENT_EVENT_TEXT_CHUNK
    text: str


class MsgEventEnd(BaseModel):
    """
    Агент закончил генерацию ответа.
    """
    type: Literal[EServerMessage.AGENT_EVENT_END] = EServerMessage.AGENT_EVENT_END
    tokens_used: int


class MsgError(BaseModel):
    """
    Неопределенная ошибка в работе агента.
    """
    type: Literal[EServerMessage.ERROR] = EServerMessage.ERROR
    code: str
    details: str


class MsgGracefulDisconnect(BaseModel):
    """
    Отправляется перед завершением работы контейнера с агентом. Например, при долгом бездействии.
    Нужно, чтобы отделять обрыв соединения из-за ошибки с необходимостью повторного подключения.
    Приход этого сообщения означает, что агент осознанно завершает работу с клиентом по какой-то причине.
    Для дальнейшего взаимодействия нужно снова обратиться к мастеру.
    """
    type: Literal[EServerMessage.GRACEFUL_DISCONNECT] = EServerMessage.GRACEFUL_DISCONNECT


AgentEventUnion = Union[MsgEventTextChunk, MsgEventEnd]


ServerMessage = TypeAdapter(Annotated[
    Union[MsgStatus, MsgEventTextChunk, MsgEventEnd, MsgError, MsgGracefulDisconnect],
    Field(discriminator="type")
])
"""
Объединяет все типы исходящих сообщений в одно для удобной автоматической десериализации.\n
Pydantic сам определит нужный тип в зависимости от поля ``type``.\n
Использование:\n
msg = ServerMessage.model_validate_json(json)
"""