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

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


# 消息历史数据表模型
# 需使用 reflex db init 初始化数据库表，若重新初始化需手动删除 alembic 相关配置和文件夹
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="运行唯一标识")
    run_new_message: str = SqlField(description="运行新增消息")

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


""" 
会话、运行和消息关系：
一次会话（conversation）包含若干次运行（run），每次运行包含用户提示词（user_prompt）和输出（output）。其中，输出包含推理和回答
"""


class EventKind(StrEnum):
    """事件类型枚举"""

    PART_START = "part_start"
    PART_DELTA = "part_delta"
    PART_END = "part_end"
    FUNCTION_TOOL_CALL = "function_tool_call"
    FUNCTION_TOOL_RESULT = "function_tool_result"
    RUN_START = "run_start"
    RUN_END = "run_end"


class PartKind(StrEnum):
    """分片类型枚举"""

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


class Event(BaseModel):
    """
    事件类
    """

    event_kind: EventKind = Field(..., description="事件类型")
    event_content: str = Field(default="", description="事件内容")
    part_index: Optional[int] = Field(default=None, description="分片索引")
    previous_part_kind: Optional[PartKind] = Field(
        default=None, description="上个分片类型"
    )
    part_kind: Optional[PartKind] = Field(default=None, description="分片类型")
    next_part_kind: Optional[PartKind] = Field(default=None, description="下个分片类型")
    tool_name: Optional[str] = Field(default=None, description="工具名称")
    run_id: str = Field(..., description="运行唯一标识")
    run_new_messages: List[ModelMessage] = Field(
        default=[], description="运行新增消息列表"
    )


class ReasoningKind(StrEnum):
    """推理类型枚举"""

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


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

    reasoning_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_expanded: bool = Field(
        default=False,
        description="推理折叠面板展开状态，True 表示展开，False 表示折叠",
    )
    answer: str = Field(default="", description="回答")
    is_streaming: bool = Field(
        default=False,
        description="运行流式输出状态，True 表示正在流式输出，False 表示非正在流式输出",
    )


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

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