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


__all__ = [
    'EServerMessage', 'MsgStatus', 'MsgError', 'MsgGracefulDisconnect',
    'MsgEventTextChunk', 'MsgEventToolCallChunk', 'MsgEventToolResult',
    'MsgEventCustomUpdate', '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_TOOL_CALL_CHUNK = "AGENT_EVENT_TOOL_CALL_CHUNK"  # Новое
    AGENT_EVENT_TOOL_RESULT = "AGENT_EVENT_TOOL_RESULT"         # Новоеы
    AGENT_EVENT_CUSTOM_UPDATE = "AGENT_EVENT_CUSTOM_UPDATE"     # Новое
    AGENT_EVENT_END = "AGENT_EVENT_END"


class MsgStatus(BaseModel):
    """Отправляется сервером при открытии соединения с клиентом."""
    type: Literal[EServerMessage.STATUS] = EServerMessage.STATUS


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


class MsgGracefulDisconnect(BaseModel):
    """Отправляется перед завершением работы контейнера с агентом."""
    type: Literal[EServerMessage.GRACEFUL_DISCONNECT] = EServerMessage.GRACEFUL_DISCONNECT


# ------------------------------------------------------------------
# AGENT EVENTS (События генерации)
# ------------------------------------------------------------------

class MsgEventTextChunk(BaseModel):
    """Чанк текста ответа агента."""
    type: Literal[EServerMessage.AGENT_EVENT_TEXT_CHUNK] = EServerMessage.AGENT_EVENT_TEXT_CHUNK
    text: str
    # Новое: "main" (главный агент) или "tools:..." (субагент, если будем использовать)
    source: str = "main"  # пока везде будет main


class MsgEventToolCallChunk(BaseModel):
    """Агент решил использовать инструмент и генерирует аргументы."""
    type: Literal[EServerMessage.AGENT_EVENT_TOOL_CALL_CHUNK] = EServerMessage.AGENT_EVENT_TOOL_CALL_CHUNK
    tool_name: Optional[str] = Field(
        None, description="Имя инструмента (приходит обычно в первом чанке)")
    args_chunk: Optional[str] = Field(
        None, description="Кусок JSON-аргументов")
    source: str = "main"


class MsgEventToolResult(BaseModel):
    """Инструмент отработал и вернул результат."""
    type: Literal[EServerMessage.AGENT_EVENT_TOOL_RESULT] = EServerMessage.AGENT_EVENT_TOOL_RESULT
    tool_name: str
    result: Any  # Может быть строкой, словарем или списком
    source: str = "main"


class MsgEventCustomUpdate(BaseModel):
    """Кастомный прогресс (например, скачивание файла) изнутри инструмента."""
    type: Literal[EServerMessage.AGENT_EVENT_CUSTOM_UPDATE] = EServerMessage.AGENT_EVENT_CUSTOM_UPDATE
    payload: Dict[str, Any] = Field(
        ..., description="Любые данные о прогрессе (status, progress и т.д.)")
    source: str = "main"


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


# ------------------------------------------------------------------
# UNIONS & ADAPTERS
# ------------------------------------------------------------------

# Обновлено: добавили новые модели в Union
AgentEventUnion = Union[
    MsgEventTextChunk,
    MsgEventToolCallChunk,
    MsgEventToolResult,
    MsgEventCustomUpdate,
    MsgEventEnd
]

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