from enum import Enum
from pydantic import BaseModel
from uuid import UUID, uuid4
from datetime import datetime
from typing import List, Optional, Union

class MessageType(str, Enum):
    TEXT = "text"
    CHART = "chart"
    SYSTEM = "system"

class MessageContent(BaseModel):
    text: Optional[str] = None
    chart_image_id: Optional[str] = None  # Reference to chart image ID
    
    @classmethod
    def create_text(cls, text: str):
        return cls(text=text)
    
    @classmethod
    def create_chart(cls, chart_image_id: str):
        return cls(chart_image_id=chart_image_id)

class ConversationMessage(BaseModel):
    id: str = str(uuid4())
    conversation_id: str
    sender_id: str  # user_id or system
    message_type: MessageType
    content: MessageContent
    timestamp: datetime = datetime.utcnow()
    metadata: dict = {}  # For additional data like AI analysis results

class Conversation(BaseModel):
    id: str = str(uuid4())
    user_id: str
    title: str
    created_at: datetime = datetime.utcnow()
    updated_at: datetime = datetime.utcnow()
    is_active: bool = True
    messages: List[ConversationMessage] = []
    referenced_charts: List[str] = []  # All chart image IDs used in this conversation