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",
    "MsgEventSendFile",
    "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_SEND_FILE = "AGENT_EVENT_SEND_FILE"  # Новое
    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 MsgEventSendFile(BaseModel):
    """Агент отправляет файл пользователю."""

    type: Literal[EServerMessage.AGENT_EVENT_SEND_FILE] = (
        EServerMessage.AGENT_EVENT_SEND_FILE
    )
    path: str = Field(..., description="Путь к файлу относительно /workspace")


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,
    MsgEventSendFile,
    MsgEventEnd,
]

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