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

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


# 数据库表模型
class HistoryMessage(SQLModel, table=True):
    id: int = SQLField(default_factory=int, primary_key=True)
    chat_id: str
    new_message: str
    timestamp: int

    @staticmethod
    def adapt(chat_id: str, new_message: List[ModelMessage]) -> "HistoryMessage":
        return HistoryMessage(
            chat_id=chat_id,
            new_message=ModelMessagesTypeAdapter.dump_json(new_message).decode("utf-8"),
            timestamp=time_ns() // 1000,  # 微秒级时间戳
        )


""" 
聊天、对话和消息关系：
一次聊天包含若干轮对话，每轮对话包含一条输入消息（input_message）和若干条输出消息（output_messages）。其中，输入消息和输出消息合称消息（message）。
"""


class Type_(StrEnum):
    """消息类型类"""

    THINKING = "thinking"
    TEXT = "text"
    CALL = "call"
    TOOL_ARGS = "tool_args"
    TOOL_RETURN = "tool_return"
    RESULT = "result"
    ERROR = "error"


# 前缀映射表（动态生成）
PREFIX_MAPING = {f"{i:02d}:": t for i, t in enumerate(Type_)}


class Message(BaseModel):
    """消息类"""

    id: str = Field(default_factory=lambda: uuid4().hex, description="消息唯一标识")
    type_: Type_ = Field(..., description="消息类型")
    content: str = Field(default="", description="消息内容")


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

    id: str = Field(default_factory=lambda: uuid4().hex, description="对话唯一标识")
    input_: str = Field(..., description="输入消息")
    output: List[Message] = Field(default_factory=list, description="输出消息")


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

    description: str = Field(default="新聊天", description="描述")
    is_streaming: bool = Field(
        default=False,
        description="流式输出状态，True 表示正在流式输出，False 表示非正在流式输出",
    )
    dialogs: List[Dialog] = Field(default_factory=list, description="对话列表")