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


class CM:
    """
    Namespace для моделей входящих сообщений (от клиента к серверу).\n
    CM = Client Message
    """

    class Type(str, Enum):
        USER_MESSAGE = "USER_MESSAGE"

    # noinspection PyPep8Naming

    class UserMessage(BaseModel):
        """
        Полное сообщение от пользователя.
        """
        type: Literal[CM.Type.USER_MESSAGE]
        text: str
        """
        Текст сообщения.
        """


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


class SM:
    """
    Namespace для моделей исходящих сообщений (от сервера к клиенту).\n
    SM = Server Message
    """

    class Type(str, Enum):
        STATUS = "STATUS"
        AGENT_EVENT = "AGENT_EVENT"
        ERROR = "ERROR"
        GRACEFUL_DISCONNECT = "GRACEFUL_DISCONNECT"

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

    class AgentEventType(str, Enum):
        TEXT_CHUNK = "TEXT_CHUNK"
        END = "END"

    class AgentEvent(BaseModel):
        """
        Базовый класс для ивентов, которые стримит агент во время генерации ответа.
        Конкретный класс для ивента определяется по ``subtype``.
        """
        type: Literal[SM.Type.AGENT_EVENT]
        subtype: SM.AgentEventType

    class EventTextChunk(AgentEvent):
        """
        Чанк текста ответа агента.
        """
        subtype: Literal[SM.AgentEventType.TEXT_CHUNK]
        text: str

    class EventEnd(AgentEvent):
        """
        Агент закончил генерацию ответа.
        """
        subtype: Literal[SM.AgentEventType.END]
        tokens_used: int

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

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


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