"""
消息数据模型
定义聊天消息的结构
"""
from datetime import datetime
from typing import Optional, Dict, Any, List
from enum import Enum
from pydantic import Field
from beanie import Document


class MessageType(str, Enum):
    """消息类型枚举"""
    TEXT = "text"               # 纯文本
    IMAGE = "image"             # 图片
    FILE = "file"               # 文件
    SYSTEM = "system"           # 系统消息
    ACTION = "action"           # 动作消息（如调用工具）


class MessageAttachment:
    """消息附件"""
    attachment_type: str        # 附件类型: image, file
    url: str                    # 资源URL
    name: str                   # 文件名
    size: int = 0               # 文件大小(字节)
    mime_type: str = ""         # MIME类型
    
    def __init__(
        self,
        attachment_type: str,
        url: str,
        name: str,
        size: int = 0,
        mime_type: str = ""
    ):
        self.attachment_type = attachment_type
        self.url = url
        self.name = name
        self.size = size
        self.mime_type = mime_type
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典"""
        return {
            "attachment_type": self.attachment_type,
            "url": self.url,
            "name": self.name,
            "size": self.size,
            "mime_type": self.mime_type
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "MessageAttachment":
        """从字典创建"""
        return cls(
            attachment_type=data.get("attachment_type", ""),
            url=data.get("url", ""),
            name=data.get("name", ""),
            size=data.get("size", 0),
            mime_type=data.get("mime_type", "")
        )


class Message(Document):
    """
    消息文档模型
    存储聊天消息
    """
    message_id: str = Field(..., description="唯一标识")
    room_id: str = Field(..., description="聊天室ID")
    discussion_id: str = Field(..., description="讨论ID")
    agent_id: Optional[str] = Field(default=None, description="发送Agent ID(系统消息为空)")
    
    # 消息内容
    content: str = Field(..., description="消息内容")
    message_type: str = Field(default=MessageType.TEXT.value, description="消息类型")
    attachments: List[Dict[str, Any]] = Field(default_factory=list, description="附件列表")
    
    # 元数据
    round: int = Field(default=0, description="所属轮次")
    token_count: int = Field(default=0, description="token数量")
    
    # 工具调用相关
    tool_calls: List[Dict[str, Any]] = Field(default_factory=list, description="工具调用记录")
    tool_results: List[Dict[str, Any]] = Field(default_factory=list, description="工具调用结果")
    
    # 时间戳
    created_at: datetime = Field(default_factory=datetime.utcnow)
    
    class Settings:
        name = "messages"
        indexes = [
            [("room_id", 1), ("created_at", 1)],
            [("discussion_id", 1)],
            [("agent_id", 1)],
        ]
    
    def get_attachments(self) -> List[MessageAttachment]:
        """获取附件对象列表"""
        return [MessageAttachment.from_dict(a) for a in self.attachments]
    
    def is_from_agent(self, agent_id: str) -> bool:
        """检查消息是否来自指定Agent"""
        return self.agent_id == agent_id
    
    def is_system_message(self) -> bool:
        """检查是否为系统消息"""
        return self.message_type == MessageType.SYSTEM.value
    
    class Config:
        json_schema_extra = {
            "example": {
                "message_id": "msg-001",
                "room_id": "product-design-room",
                "discussion_id": "disc-001",
                "agent_id": "product-manager",
                "content": "我认为登录系统应该支持多种认证方式...",
                "message_type": "text",
                "round": 1,
                "token_count": 150
            }
        }