from typing import Any, Dict, List, Optional
from pydantic import BaseModel, Field
from datetime import datetime
from enum import Enum


class MessageRole(str, Enum):
    SYSTEM = "system"
    USER = "user"
    ASSISTANT = "assistant"
    TOOL = "tool"


class Message(BaseModel):
    role: MessageRole = Field(..., description="Message role")
    content: str = Field(..., description="Message content")
    timestamp: datetime = Field(default_factory=datetime.utcnow)
    metadata: Dict[str, Any] = Field(default_factory=dict)
    tool_calls: Optional[List[Dict[str, Any]]] = None
    tool_results: Optional[List[Dict[str, Any]]] = None

    class Config:
        json_encoders = {datetime: lambda v: v.isoformat()}


class Conversation(BaseModel):
    id: str = Field(..., description="Conversation ID")
    messages: List[Message] = Field(default_factory=list)
    created_at: datetime = Field(default_factory=datetime.utcnow)
    updated_at: datetime = Field(default_factory=datetime.utcnow)
    metadata: Dict[str, Any] = Field(default_factory=dict)

    def add_message(self, message: Message) -> None:
        self.messages.append(message)
        self.updated_at = datetime.utcnow()

    def get_messages(self, role: Optional[MessageRole] = None) -> List[Message]:
        if role is None:
            return self.messages.copy()
        return [msg for msg in self.messages if msg.role == role]

    def clear(self) -> None:
        self.messages.clear()
        self.updated_at = datetime.utcnow()

    def to_dict_list(self) -> List[Dict[str, Any]]:
        return [{"role": msg.role.value, "content": msg.content} for msg in self.messages]

    class Config:
        json_encoders = {datetime: lambda v: v.isoformat()}
