from __future__ import annotations

from typing import Any, Literal, TypedDict


ChatRole = Literal["user", "assistant"]
TracePhase = Literal["system", "plan", "tool", "summary", "chat", "error"]


class ChatMessage(TypedDict):
    role: ChatRole
    content: str


class McpToolDescriptor(TypedDict, total=False):
    name: str
    title: str
    description: str
    inputSchema: Any


class McpToolResult(TypedDict, total=False):
    content: list[dict[str, Any]]
    structuredContent: Any


class PlannedToolCall(TypedDict):
    name: str
    arguments: dict[str, Any]


class ToolObservation(TypedDict, total=False):
    name: str
    source: Literal["mcp"]
    arguments: dict[str, Any]
    ok: bool
    result: Any
    error: str
    durationMs: int


class AgentTraceEvent(TypedDict, total=False):
    at: str
    phase: TracePhase
    title: str
    detail: str
    data: Any


class AgentResultTable(TypedDict):
    title: str
    columns: list[str]
    rows: list[dict[str, str]]


class AgentRunResponse(TypedDict):
    content: str
    model: str
    usedMcp: bool
    steps: int
    toolCalls: list[ToolObservation]
    tables: list[AgentResultTable]
    trace: list[AgentTraceEvent]


class AgentState(TypedDict, total=False):
    message: str
    history: list[ChatMessage]
    tools: list[McpToolDescriptor]
    plannedToolCalls: list[PlannedToolCall]
    observations: list[ToolObservation]
    visitedSignatures: list[str]
    trace: list[AgentTraceEvent]
    steps: int
    final: AgentRunResponse