# =============================================================================
# RAGFlow API 数据模型
# =============================================================================
# 说明：定义 RAGFlow API 请求/响应的 Pydantic 数据模型
# 参考：https://ragflow.io/docs/http_api_reference

from typing import Any, Dict, List, Optional

from pydantic import BaseModel, Field


class RetrievalChunk(BaseModel):
    """检索返回的单个文本块。

    Attributes:
        id: 块唯一ID
        content: 块内容文本
        document_id: 所属文档ID
        document_keyword: 所属文档名称
        similarity: 综合相似度分数
        term_similarity: 关键词相似度
        vector_similarity: 向量相似度
        highlight: 高亮标记的内容（可选）
    """
    id: str = Field(default="", description="块唯一ID")
    content: str = Field(default="", description="块内容文本")
    document_id: str = Field(default="", description="所属文档ID")
    document_keyword: str = Field(default="", description="所属文档名称")
    similarity: float = Field(default=0.0, description="综合相似度分数")
    term_similarity: float = Field(default=0.0, description="关键词相似度")
    vector_similarity: float = Field(default=0.0, description="向量相似度")
    highlight: Optional[str] = Field(default=None, description="高亮标记的内容")

    model_config = {"from_attributes": True}


class DocAggregate(BaseModel):
    """文档聚合统计。

    Attributes:
        doc_id: 文档ID
        doc_name: 文档名称
        count: 命中的块数量
    """
    doc_id: str = Field(default="", description="文档ID")
    doc_name: str = Field(default="", description="文档名称")
    count: int = Field(default=0, description="命中块数量")

    model_config = {"from_attributes": True}


class RetrievalResult(BaseModel):
    """检索结果。

    Attributes:
        chunks: 命中的文本块列表
        doc_aggs: 按文档聚合统计
        total: 命中总数
    """
    chunks: List[RetrievalChunk] = Field(default_factory=list, description="命中文本块列表")
    doc_aggs: List[DocAggregate] = Field(default_factory=list, description="文档聚合统计")
    total: int = Field(default=0, description="命中总数")

    model_config = {"from_attributes": True}


class DatasetInfo(BaseModel):
    """数据集（知识库）信息。

    Attributes:
        id: 数据集ID
        name: 数据集名称
        chunk_method: 分块方法
        permission: 权限
        document_count: 文档数量
        embedding_model: 向量模型
        create_time: 创建时间
        update_time: 更新时间
    """
    id: str = Field(default="", description="数据集ID")
    name: str = Field(default="", description="数据集名称")
    chunk_method: str = Field(default="naive", description="分块方法")
    permission: str = Field(default="me", description="权限")
    document_count: int = Field(default=0, description="文档数量")
    embedding_model: str = Field(default="", description="向量模型")
    create_time: Optional[str] = Field(default=None, description="创建时间")
    update_time: Optional[str] = Field(default=None, description="更新时间")

    model_config = {"from_attributes": True}


class DocumentInfo(BaseModel):
    """文档信息。

    Attributes:
        id: 文档ID
        name: 文档名称
        chunk_method: 分块方法
        chunk_count: 块数量
        create_time: 创建时间
        update_time: 更新时间
    """
    id: str = Field(default="", description="文档ID")
    name: str = Field(default="", description="文档名称")
    chunk_method: str = Field(default="naive", description="分块方法")
    chunk_count: int = Field(default=0, description="块数量")
    create_time: Optional[str] = Field(default=None, description="创建时间")
    update_time: Optional[str] = Field(default=None, description="更新时间")

    model_config = {"from_attributes": True}