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

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


# 消息历史数据表模型
class MessageHistory(SQLModel, table=True):
    id: str = SqlField(
        default_factory=lambda: str(uuid7()),
        primary_key=True,
        description="消息唯一标识",
    )
    conversation_id: str = SqlField(index=True, description="对话唯一标识")
    run_id: str = SqlField(index=True, description="运行唯一标识")
    new_message: str = SqlField(description="新增消息")

    @staticmethod
    def adapt(
        conversation_id: str, run_id: str, new_message: List[ModelMessage]
    ) -> "MessageHistory":
        """
        适配新增消息为消息历史记录
        :param conversation_id: 对话唯一标识
        :param run_id: 运行唯一标识
        :param new_message: 新增消息
        :return: 消息历史记录
        """
        return MessageHistory(
            conversation_id=conversation_id,
            run_id=run_id,
            new_message=ModelMessagesTypeAdapter.dump_json(new_message).decode(
                "utf-8"
            ),  # 序列化为 JSON 字符串
        )


class ReasoningKind(StrEnum):
    """推理种类"""

    THINKING = "thinking"
    TOOL_SEARCH = "tool-search"
    CAPABILITY_LOAD = "load-capability"
    TOOL_CALL = "tool-call"
    TEXT = "text"
    TOOL_RETURN = "tool-return"
    RETRY_PROMPT = "retry-prompt"
    RUN_RETURN = "run-return"


class Reasoning(BaseModel):
    """推理类"""

    kind: ReasoningKind = Field(..., description="推理种类")
    content: str = Field(default="", description="推理内容")


class Run(BaseModel):
    """运行类"""

    user_prompt: str = Field(..., description="用户提示词")
    reasonings: Dict[int, Reasoning] = Field(
        default_factory=dict, description="推理字典"
    )
    is_reasoning: bool = Field(
        default=False,
        description="推理状态，True 表示正在推理，False 表示推理完成",
    )
    is_reasoning_panel_open: bool = Field(
        default=False,
        description="推理面板展开状态，True 表示推理面板展开，False 表示推理面板折叠",
    )
    assistant_content: str = Field(default="", description="回复正文")
    is_running: bool = Field(
        default=False,
        description="运行状态，True 表示正在运行，False 表示运行完成",
    )


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

    description: str = Field(default="新对话", description="描述")
    runs: Dict[str, Run] = Field(default_factory=dict, description="运行字典")