from argentic.core.protocol.message import BaseMessage
from typing import Literal, Optional, Dict, Any
from pydantic import BaseModel, Field


class RegisterToolMessage(BaseMessage[None]):
    type: Literal["REGISTER_TOOL"] = "REGISTER_TOOL"
    tool_name: str
    tool_manual: str
    tool_api: str


class ToolRegisteredMessage(BaseMessage[None]):
    type: Literal["TOOL_REGISTERED"] = "TOOL_REGISTERED"
    tool_id: str
    tool_name: str


class ToolUnregisteredMessage(BaseMessage[None]):
    type: Literal["TOOL_UNREGISTERED"] = "TOOL_UNREGISTERED"
    tool_id: str


class ToolRegistrationErrorMessage(BaseMessage[None]):
    type: Literal["TOOL_REGISTRATION_ERROR"] = "TOOL_REGISTRATION_ERROR"
    error: str
    traceback: Optional[str] = None
    tool_id: Optional[str] = None
    tool_name: Optional[str] = None


class UnregisterToolMessage(BaseMessage[None]):
    type: Literal["UNREGISTER_TOOL"] = "UNREGISTER_TOOL"
    tool_id: str


class ToolCallRequest(BaseModel):
    """
    Represents a single tool call request, typically generated by an LLM.
    The 'tool_id' should correspond to the 'id' field of a registered tool.
    """

    tool_id: str
    arguments: Dict[str, Any] = Field(default_factory=dict)
    # If your LLM/Agent assigns a unique ID to each call within a multi-call response:
    # call_id: Optional[str] = None
