# import libs
from typing import (
    Optional,
    List,
    Literal,
    Dict
)
import time
from pydantic import BaseModel, Field


class UserMessage(BaseModel):
    """
    Model for user messages in the chat.
    """
    role: str = Field("user", description="Role of the message sender")
    content: str = Field(..., description="Content of the user message")


class AssistantMessage(BaseModel):
    """
    Model for assistant messages in the chat.
    """
    role: str = Field("assistant", description="Role of the message sender")
    content: str = Field(..., description="Content of the assistant message")


class ChatMessage(BaseModel):
    """
    Model for chat messages, which can be either user or assistant messages.
    """
    role: str = Field(..., description="Role of the message sender")
    content: str = Field(..., description="Content of the chat message")
    messages: Optional[list] = Field(
        default_factory=list, description="Additional content for the message"
    )
    thread_id: Optional[str] = Field(
        None, description="Identifier for the chat thread"
    )
    response_time: Optional[float] = Field(
        None, description="Time taken (in seconds) for the agent to respond"
    )
    timestamp: float = Field(
        default_factory=time.time,
        description="Timestamp when the message was created"
    )
    input_tokens: Optional[int] = Field(
        -1, description="Number of input tokens used in the message"
    )
    output_tokens: Optional[int] = Field(
        -1, description="Number of output tokens generated by the message"
    )
    agent_selection: Optional[str] = Field(
        None, description="Selected agent for the chat"
    )


class AgentMessage(BaseModel):
    """
    Model for agent messages in the chat.
    """
    type: Literal["user", "assistant", "system", "tool", "unknown"] = Field(
        ...,
        description="Type of the message, either 'user', 'assistant', 'system', or 'tool'"
    )
    content: str | List[str | Dict] = Field(
        ...,  description="messages in the chat thread"
    )
    name: Optional[str] = Field(
        None, description="Name of the tool or agent"
    )
    tool_call_id: Optional[str] = Field(
        None, description="Identifier for the tool call"
    )
    tool_calls: Optional[list] = Field(
        default_factory=list, description="List of tool calls associated with the message"
    )
    input_tokens: Optional[int] = Field(
        0, description="Number of input tokens used in the message"
    )
    output_tokens: Optional[int] = Field(
        0, description="Number of output tokens generated by the message"
    )


class TokenMetadata(BaseModel):
    """
    Model for metadata related to message tokens.
    """
    input_tokens: int = Field(0, description="Number of input tokens")
    output_tokens: int = Field(0, description="Number of output tokens")
