"""Tool-calling LLM output model and agent planning models.

Uses ToolArgValue and JSONValue instead of Any/dict[str, Any].
"""

from typing import Literal

from pydantic import BaseModel, Field

from core.common import JSONValue
from core.pathways.models.tool import ToolArgValue


class ToolCallingLLMOutput(BaseModel):
    """Output from tool-calling LLM step."""

    model_config = {"extra": "forbid"}

    class ToolCallTrace(BaseModel):
        """Trace of a single tool call during LLM execution."""

        model_config = {"extra": "forbid"}

        round: int
        tool_name: str
        tool_call_id: str | None = None
        parameters: dict[str, ToolArgValue] = Field(default_factory=dict)
        ok: bool
        result: JSONValue = None
        error: str | None = None
        latency_ms: float | None = None

    response: str
    model: str | None = None
    rounds: int = 1
    tool_trace: list[ToolCallTrace] = Field(default_factory=list)
    error: str | None = None


# =============================================================================
# AGENT PLANNING MODELS
# =============================================================================


class AgentToolCall(BaseModel):
    """A single tool call requested by the agent planner."""

    model_config = {"extra": "forbid"}

    tool_name: str
    parameters: dict[str, ToolArgValue] = Field(default_factory=dict)
    timeout: float | None = None


class AgentPlanInput(BaseModel):
    """Input to the agent planner step."""

    model_config = {"extra": "forbid"}

    message: str
    retrieved_memory: str = ""
    canvas_context: str = ""
    user_approved_tools: bool = False
    available_tools: str = ""


class AgentPlanOutput(BaseModel):
    """Planner output: either respond directly or call tools first."""

    model_config = {"extra": "forbid"}

    action: str = "respond"  # respond | propose_tool_calls | call_tools
    tool_calls: list[AgentToolCall] = Field(default_factory=list)
    notes: str | None = None


class AgentPlanJsonV1(BaseModel):
    """Strict JSON contract for the planner LLM output."""

    model_config = {"extra": "forbid"}

    action: Literal["respond", "propose_tool_calls", "call_tools"] = "respond"
    tool_calls: list[AgentToolCall] = Field(default_factory=list)
    notes: str | None = None


__all__ = [
    "AgentPlanInput",
    "AgentPlanJsonV1",
    "AgentPlanOutput",
    "AgentToolCall",
    "ToolCallingLLMOutput",
]
