"""Agent core types."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Callable, Literal, Protocol, Union

from pi_agents.ai.types import (
    ImageContent,
    Message,
    Model,
    TextContent,
    Tool,
)


ToolExecutionMode = Literal["sequential", "parallel"]

AgentThinkingLevel = Literal["off", "minimal", "low", "medium", "high", "xhigh"]

AgentMessage = Union[Message, dict[str, Any]]  # Message | custom app messages


@dataclass
class AgentToolResult:
    """Result from executing an agent tool."""

    content: list[TextContent | ImageContent] = field(default_factory=list)
    details: Any = None


class AgentTool(Tool):
    """Tool with execution capability."""

    label: str = ""

    async def execute(
        self,
        tool_call_id: str,
        params: dict[str, Any],
        signal: Any = None,
        on_update: Callable[..., None] | None = None,
    ) -> AgentToolResult:
        raise NotImplementedError


@dataclass
class AgentState:
    """Current state of the agent."""

    system_prompt: str = ""
    model: Model | None = None
    thinking_level: str = "off"  # AgentThinkingLevel
    tools: list[AgentTool] = field(default_factory=list)
    messages: list[AgentMessage] = field(default_factory=list)
    is_streaming: bool = False
    stream_message: AgentMessage | None = None
    pending_tool_calls: set[str] = field(default_factory=set)
    error: str | None = None


@dataclass
class BeforeToolCallResult:
    """Result from beforeToolCall hook."""

    block: bool = False
    reason: str | None = None


@dataclass
class AfterToolCallResult:
    """Result from afterToolCall hook."""

    content: list[TextContent | ImageContent] | None = None
    details: Any = None
    is_error: bool | None = None


@dataclass
class BeforeToolCallContext:
    """Context passed to beforeToolCall hook."""

    assistant_message: Any = None
    tool_call: Any = None
    args: dict[str, Any] = field(default_factory=dict)
    context: Any = None


@dataclass
class AfterToolCallContext:
    """Context passed to afterToolCall hook."""

    assistant_message: Any = None
    tool_call: Any = None
    args: dict[str, Any] = field(default_factory=dict)
    result: AgentToolResult | None = None
    is_error: bool = False
    context: Any = None


# Agent event types
AgentEvent = dict[str, Any]
# Types: "agent_start", "agent_end", "turn_start", "turn_end",
#        "message_start", "message_update", "message_end",
#        "tool_execution_start", "tool_execution_update", "tool_execution_end"
