import json
from .base import BaseModel
from . import db


class Message(BaseModel):
    __tablename__ = "messages"

    conversation_id = db.Column(
        db.BigInteger,
        db.ForeignKey("conversations.id"),
        nullable=False,
        comment="会话ID",
    )
    role = db.Column(
        db.String(20), nullable=False, comment="角色（user, assistant, system）"
    )
    content = db.Column(db.Text, nullable=False, comment="消息内容")
    tokens = db.Column(db.Integer, nullable=True, comment="Token数")
    model = db.Column(db.String(50), nullable=True, comment="使用的模型")
    tool_calls = db.Column(db.JSON, nullable=True, comment="工具调用记录")
    metadata = db.Column(db.JSON, nullable=True, comment="元数据")

    def to_dict(self):
        result = super().to_dict()
        if result.get("tool_calls"):
            result["tool_calls"] = (
                json.loads(result["tool_calls"])
                if isinstance(result["tool_calls"], str)
                else result["tool_calls"]
            )
        if result.get("metadata"):
            result["metadata"] = (
                json.loads(result["metadata"])
                if isinstance(result["metadata"], str)
                else result["metadata"]
            )
        return result