"""Pydantic models for the Cube Agent API."""

from pydantic import BaseModel, Field
from typing import Optional, List, Literal
from enum import Enum


class EventType(str, Enum):
    """Types of streaming events."""
    THINKING = "thinking"
    TOOL_CALL = "tool_call"
    TOOL_RESULT = "tool_result"
    TEXT = "text"
    DONE = "done"
    ERROR = "error"


class ChatRequest(BaseModel):
    """Request to send a message to the agent."""
    message: str = Field(..., description="The user's message")
    session_id: Optional[str] = Field(None, description="Session ID for conversation continuity")


class StreamEvent(BaseModel):
    """A streaming event from the agent."""
    event: EventType
    data: dict = Field(default_factory=dict)


class ToolCall(BaseModel):
    """A tool call made by the agent."""
    name: str
    input: dict
    id: str


class ToolResult(BaseModel):
    """Result of a tool call."""
    tool_use_id: str
    content: str


class SessionInfo(BaseModel):
    """Information about a chat session."""
    session_id: str
    message_count: int
    created_at: str
    last_active: str


class HealthResponse(BaseModel):
    """Health check response."""
    status: str
    tools: dict
    apollo_configured: bool
    model: str
