# -*- coding: utf-8 -*-
"""
数据模型
"""
from enum import StrEnum
from time import time_ns
from typing import List, Dict
from uuid import uuid4

from pydantic import BaseModel, Field
from pydantic_ai.messages import ModelMessage, ModelMessagesTypeAdapter
from sqlmodel import Field as SqlField, SQLModel


# 消息历史数据表模型
# 需使用 reflex db init 初始化数据库表，若重新初始化需手动删除 alembic 相关配置和文件夹
class MessageHistory(SQLModel, table=True):
    id: str = SqlField(
        default_factory=lambda: uuid4().hex,
        primary_key=True,
        description="消息唯一标识",
    )
    chat_id: str = SqlField(index=True, description="聊天唯一标识")
    new_message: str = SqlField(description="新消息")
    create_at: int = SqlField(
        index=True,
        description="创建时间戳（毫秒级）",
    )

    @staticmethod
    def adapt(chat_id: str, new_message: List[ModelMessage]) -> "MessageHistory":
        return MessageHistory(
            chat_id=chat_id,
            new_message=ModelMessagesTypeAdapter.dump_json(new_message).decode(
                "utf-8"
            ),  # 序列化为 JSON 字符串
            create_at=time_ns() // 1_000_000,  # 毫秒级时间戳
        )


""" 
聊天、对话和消息关系：
一次聊天包含若干轮对话，每轮对话包含用户提示词（user_prompt）和输出（output）
其中，
输出包含若干片段（part）
"""


class PartType(StrEnum):
    """片段类型（适配前端渲染）"""

    THINKING = "thinking"
    TEXT = "text"
    TOOL_NAME = "tool_name"
    TOOL_ARGS = "tool_args"
    TOOL_RETURN = "tool_return"
    FINISHED = "finished"
    ERROR = "error"


# 动态生成前缀和片段类型映射表
PREFIX_MAPING = {f"{i:02d}": t for i, t in enumerate(PartType)}


class Part(BaseModel):
    """片段类（仅就输出消息）"""

    part_type: PartType = Field(..., description="片段类型")
    content: str = Field(default="", description="片段内容")

    is_streaming: bool = Field(
        default=False,
        description="流式输出状态，True 表示正在流式输出，False 表示非正在流式输出",
    )

    is_open: bool = Field(
        default=False,
        description="折叠面板打开状态，True表示打开，False表示关闭",
    )


class Dialog(BaseModel):
    """对话类"""

    user_prompt: str = Field(..., description="用户提示词")
    output: Dict[str, Part] = Field(
        default_factory=dict, description="输出，键为片段唯一标识，值为片段对象"
    )


class Chat(BaseModel):
    """聊天类"""

    description: str = Field(default="新聊天", description="聊天描述")
    is_streaming: bool = Field(
        default=False,
        description="流式输出状态，True 表示正在流式输出，False 表示非正在流式输出",
    )
    dialogs: Dict[str, Dialog] = Field(
        default_factory=dict, description="对话列表，键为对话唯一标识，值为对话对象"
    )