"""Conversation memory for tracking message history."""

from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum
from typing import Any


class MessageRole(str, Enum):
    """Message roles in conversation."""

    SYSTEM = "system"
    USER = "user"
    ASSISTANT = "assistant"
    TOOL_RESULT = "tool_result"


@dataclass
class Message:
    """A single message in the conversation."""

    role: MessageRole
    content: str | list[dict[str, Any]]

    # For tool calls from assistant
    tool_calls: list[dict[str, Any]] | None = None

    # For tool results
    tool_use_id: str | None = None
    tool_name: str | None = None

    def to_anthropic_format(self) -> dict[str, Any]:
        """Convert message to Anthropic API format."""
        if self.role == MessageRole.SYSTEM:
            # System messages are handled separately in Anthropic API
            return {"role": "user", "content": self.content}

        if self.role == MessageRole.TOOL_RESULT:
            return {
                "role": "user",
                "content": [
                    {
                        "type": "tool_result",
                        "tool_use_id": self.tool_use_id,
                        "content": self.content,
                    }
                ],
            }

        if self.role == MessageRole.ASSISTANT and self.tool_calls:
            return {
                "role": "assistant",
                "content": self.tool_calls,
            }

        return {
            "role": self.role.value,
            "content": self.content,
        }


@dataclass
class ConversationMemory:
    """Tracks the full message history for the current session."""

    messages: list[Message] = field(default_factory=list)

    def add_user(self, content: str) -> None:
        """Add a user message."""
        self.messages.append(Message(role=MessageRole.USER, content=content))

    def add_assistant(self, content: str) -> None:
        """Add an assistant text response."""
        self.messages.append(Message(role=MessageRole.ASSISTANT, content=content))

    def add_assistant_tool_calls(self, tool_calls: list[dict[str, Any]]) -> None:
        """Add an assistant message with tool calls."""
        self.messages.append(
            Message(role=MessageRole.ASSISTANT, content="", tool_calls=tool_calls)
        )

    def add_tool_result(self, tool_use_id: str, tool_name: str, result: str) -> None:
        """Add a tool result message."""
        self.messages.append(
            Message(
                role=MessageRole.TOOL_RESULT,
                content=result,
                tool_use_id=tool_use_id,
                tool_name=tool_name,
            )
        )

    def get_messages(self) -> list[Message]:
        """Get all messages in the conversation."""
        return self.messages.copy()

    def get_anthropic_messages(self) -> list[dict[str, Any]]:
        """Get messages in Anthropic API format (excluding system)."""
        return [msg.to_anthropic_format() for msg in self.messages]

    def clear(self) -> None:
        """Clear all messages."""
        self.messages.clear()

    def truncate(self, max_messages: int) -> None:
        """Keep only the most recent messages.

        In a real implementation, you'd want token-based truncation.
        This is a simplified version that keeps N most recent messages.
        """
        if len(self.messages) > max_messages:
            self.messages = self.messages[-max_messages:]

    def __len__(self) -> int:
        """Return number of messages."""
        return len(self.messages)
