"""
Core message types for agent communication using Pydantic models.

This module defines the structured message types that agents use to communicate
with each other and with LLMs, following the OpenAI API format.

Message hierarchy:
- BaseMessage: Base class for all message types
  - SystemMessage: Instructions/role definition
  - UserMessage: User input or question
  - AssistantMessage: Model response (text, tool calls, or structured)
  - ToolMessage: Tool execution result

Additional types:
- ToolCallRequest: Structured tool call representation
- MultiModalMessage: Messages supporting multiple content types
- Usage: Token usage statistics
"""

from datetime import datetime
from typing import Any, Dict, List, Literal, Optional, Union

from pydantic import BaseModel, Field, model_validator


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

    All messages in agentbyte conversations have:
    - content: The actual message text
    - source: Who/what sent this message (for tracking and debugging)
    - timestamp: When the message was created (for ordering and analysis)

    This structure enables comprehensive message tracking, routing, and replay
    capabilities essential for multi-agent systems.
    """

    content: str = Field(..., description="The message content")
    source: str = Field(
        ..., description="Source of the message (agent name, system, user, etc.)"
    )
    timestamp: datetime = Field(
        default_factory=datetime.now, description="When the message was created"
    )

    model_config = {"frozen": True}

    def __str__(self) -> str:
        """Returns a user-friendly string representation."""
        time_str = self.timestamp.strftime("%H:%M:%S")
        return f"[{self.source}] {time_str} | {self.content}"

    def __repr__(self) -> str:
        """Returns an unambiguous, developer-friendly representation."""
        class_name = self.__class__.__name__
        content_preview = self.content[:50] + "..." if len(self.content) > 50 else self.content
        return f"{class_name}(source='{self.source}', content='{content_preview}', timestamp='{self.timestamp}')"


class SystemMessage(BaseMessage):
    """
    System message containing instructions/role definition for the agent.

    System messages are typically sent first in a conversation and define:
    - How the assistant should behave
    - What knowledge/context it should use
    - Any constraints or guidelines
    - Personality or tone to adopt

    System messages are not visible to the user but guide all model responses.

    Example:
        system = SystemMessage(
            content="You are a helpful assistant that answers questions concisely.",
            source="system"
        )
    """

    role: Literal["system"] = Field(default="system", description="Message role")


class UserMessage(BaseMessage):
    """
    User message containing input from human or external system.

    User messages represent:
    - Questions the user is asking
    - Statements or context the user is providing
    - Commands or requests for the assistant
    - Feedback or corrections

    User messages appear in the conversation history and drive the assistant's responses.

    Attributes:
        name: Optional name/identifier of the user (useful for multi-user conversations)

    Example:
        user = UserMessage(
            content="What is the capital of France?",
            source="user_123"
        )
        user_named = UserMessage(
            content="How do I make pasta?",
            source="alice",
            name="Alice"
        )
    """

    role: Literal["user"] = Field(default="user", description="Message role")
    name: Optional[str] = Field(default=None, description="Optional name of the user")


class ToolCallRequest(BaseModel):
    """
    Structured representation of an LLM's tool call request.

    When an LLM decides to invoke a tool or function, it generates this structure
    containing:
    - The tool/function name
    - The arguments to pass to it
    - A unique ID to correlate with results

    This structure matches picoagents ToolCallRequest for compatibility.

    Attributes:
        tool_name: Name of the tool to call
        parameters: Dictionary of arguments for the tool
        call_id: Unique identifier for this specific call (used to match tool results)

    Example:
        call = ToolCallRequest(
            tool_name="get_weather",
            parameters={"location": "Paris", "unit": "celsius"},
            call_id="call_abc123"
        )
    """

    tool_name: str = Field(..., description="Name of the tool to call")
    parameters: Dict[str, Any] = Field(..., description="Arguments for the tool")
    call_id: str = Field(..., description="Unique identifier for this call")

    model_config = {"frozen": True}


class Usage(BaseModel):
    """
    Token usage statistics from an LLM API call.

    Tracks how many tokens were used for input (prompt) and output (completion),
    and optionally includes cost estimation data and execution metrics.

    Attributes:
        duration_ms: Total execution time in milliseconds
        llm_calls: Number of LLM API calls made
        tokens_input: Number of tokens in the prompt
        tokens_output: Number of tokens generated in the completion
        tool_calls: Number of tool executions
        memory_operations: Number of memory read/write operations
        cost_estimate: Estimated cost in USD (optional, provider-specific)

    Example:
        usage = Usage(
            duration_ms=250,
            llm_calls=1,
            tokens_input=150,
            tokens_output=45,
            tool_calls=1,
            cost_estimate=0.0025
        )
        total = usage.total_tokens  # 195
    """
    duration_ms: int = Field(
        default=0,
        description="Total execution time in milliseconds",
        ge=0,
    )
    llm_calls: int = Field(
        default=0,
        description="Number of LLM API calls made",
        ge=0,
    )
    tokens_input: int = Field(
        default=0,
        description="Number of tokens used in the prompt",
        ge=0,
    )
    tokens_output: int = Field(
        default=0,
        description="Number of tokens generated in the completion",
        ge=0,
    )
    tool_calls: int = Field(
        default=0,
        description="Number of tool executions",
        ge=0,
    )
    memory_operations: int = Field(
        default=0,
        description="Number of memory read/write operations",
        ge=0,
    )
    cost_estimate: Optional[float] = Field(
        default=None,
        description="Estimated cost in USD (optional, provider-specific)",
        ge=0.0
    )

    @property
    def total_tokens(self) -> int:
        """Total tokens used (input + output)."""
        return self.tokens_input + self.tokens_output

    def __add__(self, other: "Usage") -> "Usage":
        """Aggregate usage statistics from multiple sources."""
        return Usage(
            duration_ms=max(self.duration_ms, other.duration_ms),
            llm_calls=self.llm_calls + other.llm_calls,
            tokens_input=self.tokens_input + other.tokens_input,
            tokens_output=self.tokens_output + other.tokens_output,
            tool_calls=self.tool_calls + other.tool_calls,
            memory_operations=self.memory_operations + other.memory_operations,
            cost_estimate=(self.cost_estimate or 0) + (other.cost_estimate or 0)
            or None,
        )

    model_config = {"frozen": True}


class AssistantMessage(BaseMessage):
    """
    Assistant message containing response from the agent/LLM.

    Assistant messages are responses generated by the language model. They can contain:
    - Plain text response to the user
    - Tool calls if the model needs to take action
    - Structured output when output_format is specified
    - Token usage for monitoring and cost analysis

    The assistant can vocalize its reasoning in content while also requesting
    tool execution, allowing for transparent multi-step reasoning.

    Attributes:
        tool_calls: Optional list of tool calls made by the assistant
        structured_content: Optional structured data when output_format is used
        usage: Optional token usage statistics for this specific message

    Example:
        # Simple text response
        msg = AssistantMessage(
            content="Paris is the capital of France.",
            source="gpt-4"
        )

        # Response with tool call
        msg = AssistantMessage(
            content="I'll check the weather for you.",
            source="gpt-4",
            tool_calls=[
                ToolCallRequest(
                    tool_name="get_weather",
                    parameters={"location": "Paris"},
                    call_id="call_1"
                )
            ]
        )
    """

    role: Literal["assistant"] = Field(default="assistant", description="Message role")
    tool_calls: Optional[List[ToolCallRequest]] = Field(
        default=None, description="Tool calls made by the assistant"
    )
    structured_content: Optional[BaseModel] = Field(
        default=None, description="Structured data when output_format is used"
    )
    usage: Optional[Usage] = Field(
        default=None, description="Token usage for this LLM call"
    )

    def __str__(self) -> str:
        """Returns a user-friendly string representation."""
        time_str = self.timestamp.strftime("%H:%M:%S")

        if self.tool_calls:
            # Show tool calls information
            tool_info = ", ".join(
                [
                    f"{tc.tool_name}({', '.join(f'{k}={v}' for k, v in tc.parameters.items())})"
                    for tc in self.tool_calls
                ]
            )
            if self.content and self.content.strip():
                return (
                    f"[{self.source}] {time_str} | {self.content} [tools: {tool_info}]"
                )
            else:
                return f"[{self.source}] {time_str} | [calling tools: {tool_info}]"
        else:
            return f"[{self.source}] {time_str} | {self.content}"


class ToolMessage(BaseMessage):
    """
    Tool message containing result from tool execution.

    When the agent executes a tool call, it sends the result back via this message.
    This keeps the conversation flowing and provides the information the model needs
    to continue reasoning and generating responses.

    Attributes:
        tool_call_id: ID of the tool call this is responding to
        tool_name: Name of the tool that was executed
        success: Whether tool execution succeeded
        error: Error message if execution failed

    Example:
        # Success case
        result = ToolMessage(
            content="Weather in Paris: Sunny, 22°C",
            tool_call_id="call_1",
            tool_name="get_weather",
            success=True,
            source="system"
        )

        # Error case
        error = ToolMessage(
            content="Invalid location provided",
            tool_call_id="call_1",
            tool_name="get_weather",
            success=False,
            error="Invalid location: 'NewYork' -> did you mean 'New York'?",
            source="system"
        )
    """

    role: Literal["tool"] = Field(default="tool", description="Message role")
    tool_call_id: str = Field(
        ..., description="ID of the tool call this is responding to"
    )
    tool_name: str = Field(..., description="Name of the tool that was executed")
    success: bool = Field(..., description="Whether tool execution succeeded")
    error: Optional[str] = Field(default=None, description="Error message if failed")


class MultiModalMessage(BaseMessage):
    """
    Message supporting multiple content types (text, images, audio, video, etc.).

    This message type allows agents to handle and pass through rich media content
    to models that support multimodal inputs. Content can be provided either inline
    (as binary data or base64) or via URL reference.

    Attributes:
        mime_type: MIME type of the content (e.g., 'text/plain', 'image/jpeg', 'audio/wav')
        data: Binary data (bytes) or base64 string for the content
        media_url: URL to media content if data is not provided
        metadata: Additional content metadata (resolution, duration, etc.)

    Example:
        # Image as inline base64
        msg = MultiModalMessage(
            content="Here's a screenshot",
            source="user",
            role="user",
            mime_type="image/jpeg",
            data=base64_encoded_image
        )

        # Video via URL
        msg = MultiModalMessage(
            content="Check out this demo",
            source="user",
            role="user",
            mime_type="video/mp4",
            media_url="https://example.com/demo.mp4"
        )
    """

    role: Literal["user", "assistant"] = Field(..., description="Message role")
    mime_type: str = Field(
        ...,
        description="MIME type of the content (e.g., 'text/plain', 'image/jpeg', 'audio/wav', 'video/mp4')",
    )
    data: Optional[Union[bytes, str]] = Field(
        None, description="Binary data (bytes) or base64 string for the content"
    )
    media_url: Optional[str] = Field(
        None, description="URL to media content if data is not provided"
    )
    metadata: Dict[str, Any] = Field(
        default_factory=dict, description="Additional content metadata"
    )

    @model_validator(mode="after")
    def validate_media_data(self):
        """Ensure either data or media_url is provided."""
        if self.data is None and self.media_url is None:
            raise ValueError("Either 'data' or 'media_url' must be provided")

        if self.data is not None and self.media_url is not None:
            raise ValueError("Only one of 'data' or 'media_url' should be provided")

        return self

    def is_text(self) -> bool:
        """Check if this is a text message."""
        return self.mime_type.startswith("text/")

    def is_image(self) -> bool:
        """Check if this is an image message."""
        return self.mime_type.startswith("image/")

    def is_audio(self) -> bool:
        """Check if this is an audio message."""
        return self.mime_type.startswith("audio/")

    def is_video(self) -> bool:
        """Check if this is a video message."""
        return self.mime_type.startswith("video/")

    def to_base64(self) -> Optional[str]:
        """Convert data to base64 string for API usage."""
        if self.data is None:
            return None

        # If data is already a string, assume it's base64
        if isinstance(self.data, str):
            return self.data

        # If data is bytes, encode to base64
        import base64

        return base64.b64encode(self.data).decode("utf-8")


# Union type for all message types
Message = Union[
    SystemMessage, UserMessage, AssistantMessage, ToolMessage, MultiModalMessage
]

__all__ = [
    # Base class
    "BaseMessage",
    # Message types
    "SystemMessage",
    "UserMessage",
    "AssistantMessage",
    "ToolMessage",
    "MultiModalMessage",
    # Tool structure
    "ToolCallRequest",
    # Usage tracking
    "Usage",
    # Union type
    "Message",
]
