class BaseChatMessage(BaseModel):
    """Base class for all chat message types.

    Attributes:
        id: Optional unique identifier for the message.
        content: The message content, either as a string or list of Content objects.
        role: The role of the message sender (system, user, assistant, tool).
        metadata: Additional structured metadata about the message.
    """

    id: str | None = None
    content: str | list[Content]
    role: Literal["system", "user", "assistant", "tool"]
    metadata: dict[str, Any] | None = None

    @property
    def text(self) -> str:
        """Get the text content of the message.

        Returns:
            str: The text content of the message. If content is a list,
                 concatenates all text content elements with newlines.
        """
        if isinstance(self.content, str):
            return self.content
        else:
            all_text = [content.text for content in self.content if content.type == "text"]
            return "\n".join(all_text)

class SystemMessage(BaseChatMessage):
    """System message in a chat conversation.

    Attributes:
        role: Always set to "system".
    """

    role: Literal["system"] = "system"  # type: ignore

class UserMessage(BaseChatMessage):
    """User message in a chat conversation.

    Attributes:
        role: Always set to "user".
        tool_call_id: Optional list of tool call IDs this message is responding to.
    """

    role: Literal["user"] = "user"  # type: ignore
    tool_call_id: list[str] | None = None

class AssistantMessage(BaseChatMessage):
    """Assistant message in a chat conversation.

    Attributes:
        role: Always set to "assistant".
        model: Optional identifier for the model that generated this message.
        tool_calls: Optional list of tool calls made by the assistant.
    """

    role: Literal["assistant"] = "assistant"  # type: ignore
    model: str | None = None
    tool_calls: list[ToolCall] | None = None

class DocentAssistantMessage(AssistantMessage):
    """Assistant message in a chat session with additional chat-specific metadata.

    This extends AssistantMessage with fields that are only relevant in Docent chat contexts

    Attributes:
        citations: Optional list of citations referenced in the message content.
        suggested_messages: Optional list of suggested followup messages.
    """

    citations: list[InlineCitation] | None = None
    suggested_messages: list[str] | None = None

class ToolMessage(BaseChatMessage):
    """Tool message in a chat conversation.

    Attributes:
        role: Always set to "tool".
        tool_call_id: Optional ID of the tool call this message is responding to.
        function: Optional name of the function that was called.
        error: Optional error information if the tool call failed.
    """

    role: Literal["tool"] = "tool"  # type: ignore

    tool_call_id: str | None = None
    function: str | None = None
    error: dict[str, Any] | None = None

def parse_chat_message(message_data: dict[str, Any] | ChatMessage) -> ChatMessage:
    """Parse a message dictionary or object into the appropriate ChatMessage subclass.

    This parses base messages without chat-specific fields. For chat sessions,
    use parse_docent_chat_message instead.

    Args:
        message_data: A dictionary or ChatMessage object representing a chat message.

    Returns:
        ChatMessage: An instance of a ChatMessage subclass based on the role.

    Raises:
        ValueError: If the message role is unknown.
    """
    if isinstance(message_data, (SystemMessage, UserMessage, AssistantMessage, ToolMessage)):
        return message_data

    role = message_data.get("role")
    if role == "system":
        return SystemMessage.model_validate(message_data)
    elif role == "user":
        return UserMessage.model_validate(message_data)
    elif role == "assistant":
        return AssistantMessage.model_validate(message_data)
    elif role == "tool":
        return ToolMessage.model_validate(message_data)
    else:
        raise ValueError(f"Unknown message role: {role}")

def parse_docent_chat_message(
    message_data: dict[str, Any] | DocentChatMessage,
) -> DocentChatMessage:
    """Parse a message dictionary or object into the appropriate DocentChatMessage subclass.

    This handles chat session messages which may include DocentAssistantMessage with
    citations and suggested_messages fields.

    Args:
        message_data: A dictionary or DocentChatMessage object representing a chat session message.

    Returns:
        DocentChatMessage: An instance of a DocentChatMessage subclass based on the role.

    Raises:
        ValueError: If the message role is unknown.
    """
    if isinstance(
        message_data,
        (SystemMessage, UserMessage, DocentAssistantMessage, AssistantMessage, ToolMessage),
    ):
        return message_data

    role = message_data.get("role")
    if role == "system":
        return SystemMessage.model_validate(message_data)
    elif role == "user":
        return UserMessage.model_validate(message_data)
    elif role == "assistant":
        return DocentAssistantMessage.model_validate(message_data)
    elif role == "tool":
        return ToolMessage.model_validate(message_data)
    else:
        raise ValueError(f"Unknown message role: {role}")

class BaseContent(BaseModel):
    """Base class for all content types in chat messages.

    Provides the foundation for different content types with a discriminator field.

    Attributes:
        type: The content type identifier, used for discriminating between content types.
    """

    type: Literal["text", "reasoning", "image", "audio", "video"]

class ContentText(BaseContent):
    """Text content for chat messages.

    Represents plain text content in a chat message.

    Attributes:
        type: Fixed as "text" to identify this content type.
        text: The actual text content.
        refusal: Optional flag indicating if this is a refusal message.
    """

    type: Literal["text"] = "text"  # type: ignore
    text: str
    refusal: bool | None = None

class ContentReasoning(BaseContent):
    """Reasoning content for chat messages.

    Represents reasoning or thought process content in a chat message.

    Attributes:
        type: Fixed as "reasoning" to identify this content type.
        reasoning: The actual reasoning text.
        signature: Optional signature associated with the reasoning.
        redacted: Flag indicating if the reasoning has been redacted.
    """

    type: Literal["reasoning"] = "reasoning"  # type: ignore
    reasoning: str
    signature: str | None = None
    redacted: bool = False

@dataclass
class ToolCall:
    """Tool call information.

    Attributes:
        id: Unique identifier for tool call.
        type: Type of tool call. Can only be "function" or None.
        function: Function called.
        arguments: Arguments to function.
        parse_error: Error which occurred parsing tool call.
        view: Custom view of tool call input.
    """

    id: str
    function: str
    arguments: dict[str, Any]
    type: Literal["function"] | None = None
    parse_error: str | None = None
    view: ToolCallContent | None = None

class ToolCallContent(BaseModel):
    """Content to include in tool call view.

    Attributes:
        title: Optional (plain text) title for tool call content.
        format: Format (text or markdown).
        content: Text or markdown content.
    """

    title: str | None = None
    format: Literal["text", "markdown"]
    content: str

class ToolParam(BaseModel):
    """A parameter for a tool function.

    Args:
        name: The name of the parameter.
        description: A description of what the parameter does.
        input_schema: JSON Schema describing the parameter's type and validation rules.
    """

    name: str
    description: str
    input_schema: dict[str, Any]

class ToolParams(BaseModel):
    """Description of tool parameters object in JSON Schema format.

    Args:
        type: The type of the parameters object, always 'object'.
        properties: Dictionary mapping parameter names to their ToolParam definitions.
        required: List of required parameter names.
        additionalProperties: Whether additional properties are allowed beyond those
            specified. Always False.
    """

    type: Literal["object"] = "object"
    properties: dict[str, ToolParam] = Field(default_factory=dict)
    required: list[str] = Field(default_factory=list)
    additionalProperties: bool = False

class ToolInfo(BaseModel):
    """Specification of a tool (JSON Schema compatible).

    If you are implementing a ModelAPI, most LLM libraries can
    be passed this object (dumped to a dict) directly as a function
    specification. For example, in the OpenAI provider:

    ```python
    ChatCompletionToolParam(
        type="function",
        function=tool.model_dump(exclude_none=True),
    )
    ```

    In some cases the field names don't match up exactly. In that case
    call `model_dump()` on the `parameters` field. For example, in the
    Anthropic provider:

    ```python
    ToolParam(
        name=tool.name,
        description=tool.description,
        input_schema=tool.parameters.model_dump(exclude_none=True),
    )
    ```

    Attributes:
        name: Name of tool.
        description: Short description of tool.
        parameters: JSON Schema of tool parameters object.
    """

    name: str
    description: str
    parameters: ToolParams = Field(default_factory=ToolParams)