"""向量搜索 API 数据模型。

用于 Knowledge Graph 向量搜索接口的请求和响应数据模型。
"""

from typing import Generic, TypeVar

from pydantic import BaseModel, Field


class VectorSearchRequest(BaseModel):
    """向量搜索请求模型。

    用于向量库的语义搜索查询。

    Attributes
    ----------
    question : str
        搜索问题，根据模型的最大输入长度限制
    limit : int
        返回结果数量限制，范围 1-50，默认 10
    score_threshold : float
        相似度阈值，范围 0.0-1.0，默认 0.5
    offset : int
        结果偏移量，用于分页，默认 0
    enable_entity_extraction : bool
        是否启用实体抽取，true 为实体抽取后查询，false 为直接根据问题内容向量查询，默认 true

    Examples
    --------
    >>> request = VectorSearchRequest(
    ...     question="菜品有哪些",
    ...     limit=10,
    ...     score_threshold=0.5
    ... )
    >>> print(request.question)
    '菜品有哪些'
    """

    question: str = Field(..., description="搜索问题")
    limit: int = Field(default=10, ge=1, le=50, description="返回结果数量限制")
    score_threshold: float = Field(default=0.5, ge=0.0, le=1.0, description="相似度阈值")
    offset: int = Field(default=0, ge=0, description="结果偏移量，用于分页")
    enable_entity_extraction: bool = Field(
        default=True,
        description="是否启用实体抽取",
    )


class EntityMetadata(BaseModel):
    """实体集合元数据。

    知识图谱中的实体节点信息。

    Attributes
    ----------
    entity_id : str
        实体唯一标识
    entity_type : str
        实体类型 (如: Category-品类)
    name : str
        实体名称
    description : str
        实体描述
    doc_ids : list[str]
        关联的文档ID列表
    source_chunk_ids : list[str]
        来源文本块ID列表
    source_picture_ids : list[str]
        来源图片ID列表
    update_at : str
        更新时间
    file_type : str
        文件类型
    """

    batch_id: str | None = Field(default=None, description="批次ID")
    vector_id: str | None = Field(default=None, description="向量ID")
    content_id: str | None = Field(default=None, description="内容ID")
    text: str | None = Field(default=None, description="文本内容")
    text_length: int | None = Field(default=None, description="文本长度")
    vector_dimension: int | None = Field(default=None, description="向量维度")
    entity_id: str = Field(..., description="实体唯一标识")
    entity_type: str = Field(..., description="实体类型")
    name: str = Field(..., description="实体名称")
    aliases: list[str] = Field(default_factory=list, description="实体别名列表")
    doc_ids: list[str] = Field(default_factory=list, description="关联的文档ID列表")
    source_chunk_ids: list[str] = Field(default_factory=list, description="来源文本块ID列表")
    source_picture_ids: list[str] = Field(default_factory=list, description="来源图片ID列表")
    description: str = Field(default="", description="实体描述")
    update_at: str | None = Field(default=None, description="更新时间")
    status: str | None = Field(default=None, description="状态")
    entity_name: str | None = Field(default=None, description="实体名称")
    source_file: str | None = Field(default=None, description="源文件路径")
    file_type: str = Field(..., description="文件类型")


class RelationshipMetadata(BaseModel):
    """关系集合元数据。

    实体之间的关系信息。

    Attributes
    ----------
    relation_id : str
        关系唯一标识
    relationship_type : str
        关系类型
    source : str
        源实体名称
    source_id : str
        源实体ID
    target : str
        目标实体名称
    target_id : str
        目标实体ID
    description : str
        关系描述
    weight : float
        关系权重
    """

    batch_id: str | None = Field(default=None, description="批次ID")
    vector_id: str | None = Field(default=None, description="向量ID")
    content_id: str | None = Field(default=None, description="内容ID")
    text: str | None = Field(default=None, description="文本内容")
    text_length: int | None = Field(default=None, description="文本长度")
    vector_dimension: int | None = Field(default=None, description="向量维度")
    relation_id: str = Field(..., description="关系唯一标识")
    relationship_type: str = Field(default="", description="关系类型")
    description: str = Field(default="", description="关系描述")
    weight: float | int = Field(..., description="关系权重")
    source_chunk_ids: list[str] = Field(default_factory=list, description="来源文本块ID列表")
    source_picture_ids: list[str] = Field(default_factory=list, description="来源图片ID列表")
    source: str = Field(..., description="源实体名称")
    source_id: str = Field(..., description="源实体ID")
    target: str = Field(..., description="目标实体名称")
    target_id: str = Field(..., description="目标实体ID")
    update_at: str | None = Field(default=None, description="更新时间")
    status: str | None = Field(default=None, description="状态")
    source_file: str | None = Field(default=None, description="源文件路径")
    file_type: str = Field(..., description="文件类型")


class ChunkMetadata(BaseModel):
    """文本块集合元数据。

    文档分块后的文本内容信息。

    Attributes
    ----------
    chunk_id : str
        内容块唯一标识
    doc_id : str
        所属文档ID
    doc_name : str
        文档名称
    page : int
        页码信息
    text_length : int
        文本长度
    """

    batch_id: str | None = Field(default=None, description="批次ID")
    vector_id: str | None = Field(default=None, description="向量ID")
    content_id: str | None = Field(default=None, description="内容ID")
    text: str | None = Field(default=None, description="文本内容")
    text_length: int = Field(..., description="文本长度")
    vector_dimension: int | None = Field(default=None, description="向量维度")
    chunk_id: str = Field(..., description="内容块唯一标识")
    doc_id: str = Field(..., description="所属文档ID")
    doc_name: str = Field(..., description="文档名称")
    page: int = Field(..., description="页码信息")
    update_at: str | None = Field(default=None, description="更新时间")
    status: str | None = Field(default=None, description="状态")
    source_file: str | None = Field(default=None, description="源文件路径")
    file_type: str = Field(..., description="文件类型")


class PictureMetadata(BaseModel):
    """图片集合元数据。

    文档中的图片及其描述信息。

    Attributes
    ----------
    pic_id : str
        图片唯一标识
    description : str
        图片描述文本
    filename : str
        图片文件路径
    doc_id : str
        所属文档ID
    doc_name : str
        文档名称
    page : int
        页码信息
    """

    batch_id: str | None = Field(default=None, description="批次ID")
    vector_id: str | None = Field(default=None, description="向量ID")
    content_id: str | None = Field(default=None, description="内容ID")
    text: str | None = Field(default=None, description="文本内容")
    text_length: int | None = Field(default=None, description="文本长度")
    vector_dimension: int | None = Field(default=None, description="向量维度")
    pic_id: str = Field(..., description="图片唯一标识")
    description: str = Field(default="", description="图片描述文本")
    filename: str = Field(..., description="图片文件路径")
    doc_id: str = Field(..., description="所属文档ID")
    doc_name: str = Field(..., description="文档名称")
    page: int = Field(..., description="页码信息")
    update_at: str | None = Field(default=None, description="更新时间")
    status: str | None = Field(default=None, description="状态")
    source_file: str | None = Field(default=None, description="源文件路径")
    file_type: str = Field(..., description="文件类型")


class CommunityFinding(BaseModel):
    """社区发现和分析结果。

    Attributes
    ----------
    summary : str
        发现摘要
    explanation : str
        详细解释
    """

    summary: str = Field(..., description="发现摘要")
    explanation: str = Field(..., description="详细解释")


class CommunityMetadata(BaseModel):
    """社区集合元数据。

    知识图谱社区分析结果。

    Attributes
    ----------
    community_id : str
        社区唯一标识
    title : str
        社区标题
    summary : str
        社区摘要
    entity_ids : list[str]
        包含的实体ID列表
    relationship_ids : list[str]
        包含的关系ID列表
    rating : float
        社区评分
    findings : list[CommunityFinding]
        社区发现和分析结果
    """

    batch_id: str | None = Field(default=None, description="批次ID")
    vector_id: str | None = Field(default=None, description="向量ID")
    content_id: str | None = Field(default=None, description="内容ID")
    text: str | None = Field(default=None, description="文本内容")
    text_length: int | None = Field(default=None, description="文本长度")
    vector_dimension: int | None = Field(default=None, description="向量维度")
    community_id: str = Field(..., description="社区唯一标识")
    entity_ids: list[str] = Field(default_factory=list, description="包含的实体ID列表")
    relationship_ids: list[str] = Field(default_factory=list, description="包含的关系ID列表")
    title: str = Field(..., description="社区标题")
    summary: str = Field(..., description="社区摘要")
    parent_id: str | None = Field(default=None, description="父社区ID")
    children_ids: list[str] = Field(default_factory=list, description="子社区ID列表")
    rating: float = Field(..., description="社区评分")
    rating_explanation: str | None = Field(default=None, description="评分说明")
    findings: list[CommunityFinding] = Field(
        default_factory=list,
        description="社区发现和分析结果",
    )
    update_at: str | None = Field(default=None, description="更新时间")
    status: str | None = Field(default=None, description="状态")
    source_file: str | None = Field(default=None, description="源文件路径")
    file_type: str = Field(..., description="文件类型")


# 定义泛型类型变量
MetadataType = TypeVar("MetadataType", bound=BaseModel)


class SearchResultItem(BaseModel, Generic[MetadataType]):
    """搜索结果项通用结构（泛型）。

    使用泛型保持强类型的元数据。

    Parameters
    ----------
    MetadataType : TypeVar
        元数据的具体类型，可以是 EntityMetadata, RelationshipMetadata 等

    Attributes
    ----------
    content : str
        匹配的内容文本
    score : float
        相似度分数 (0-1)
    metadata : MetadataType
        详细的元数据信息，类型由泛型参数决定
    source : str
        来源集合名称 (entities, relationships, chunks, pictures, communities)

    Examples
    --------
    >>> metadata = EntityMetadata(
    ...     entity_id="19c2c8ed",
    ...     entity_type="Category-品类",
    ...     name="西北菜",
    ...     file_type="entity"
    ... )
    >>> item = SearchResultItem[EntityMetadata](
    ...     content="西北菜是中国西北地区的传统菜肴",
    ...     score=0.59961027,
    ...     metadata=metadata,
    ...     source="entities"
    ... )
    >>> print(item.metadata.entity_id)
    '19c2c8ed'
    """

    content: str = Field(..., description="匹配的内容文本")
    score: float = Field(..., ge=0.0, le=1.0, description="相似度分数 (0-1)")
    metadata: MetadataType = Field(..., description="详细的元数据信息")
    source: str = Field(..., description="来源集合名称")


class SearchResults(BaseModel):
    """按集合分组的搜索结果。

    使用强类型的泛型列表，每个集合对应其特定的元数据类型。

    Attributes
    ----------
    entities : list[SearchResultItem[EntityMetadata]]
        实体搜索结果列表，元数据为 EntityMetadata 类型
    relationships : list[SearchResultItem[RelationshipMetadata]]
        关系搜索结果列表，元数据为 RelationshipMetadata 类型
    chunks : list[SearchResultItem[ChunkMetadata]]
        文本块搜索结果列表，元数据为 ChunkMetadata 类型
    pictures : list[SearchResultItem[PictureMetadata]]
        图片搜索结果列表，元数据为 PictureMetadata 类型
    communities : list[SearchResultItem[CommunityMetadata]]
        社区搜索结果列表，元数据为 CommunityMetadata 类型

    Examples
    --------
    >>> entity_meta = EntityMetadata(
    ...     entity_id="test",
    ...     entity_type="Category",
    ...     name="测试",
    ...     file_type="entity"
    ... )
    >>> entity_item = SearchResultItem[EntityMetadata](
    ...     content="测试内容",
    ...     score=0.8,
    ...     metadata=entity_meta,
    ...     source="entities"
    ... )
    >>> results = SearchResults(entities=[entity_item])
    >>> print(results.entities[0].metadata.entity_type)
    'Category'
    """

    entities: list[SearchResultItem[EntityMetadata]] = Field(
        default_factory=list,
        description="实体搜索结果列表",
    )
    relationships: list[SearchResultItem[RelationshipMetadata]] = Field(
        default_factory=list,
        description="关系搜索结果列表",
    )
    chunks: list[SearchResultItem[ChunkMetadata]] = Field(
        default_factory=list,
        description="文本块搜索结果列表",
    )
    pictures: list[SearchResultItem[PictureMetadata]] = Field(
        default_factory=list,
        description="图片搜索结果列表",
    )
    communities: list[SearchResultItem[CommunityMetadata]] = Field(
        default_factory=list,
        description="社区搜索结果列表",
    )


class VectorSearchResponse(BaseModel):
    """向量搜索响应模型。

    Attributes
    ----------
    question : str
        原始搜索问题
    search_results : SearchResults
        按集合分组的搜索结果

    Examples
    --------
    >>> entity_meta = EntityMetadata(
    ...     entity_id="test",
    ...     entity_type="Category",
    ...     name="西北菜",
    ...     file_type="entity"
    ... )
    >>> entity_item = SearchResultItem[EntityMetadata](
    ...     content="西北菜",
    ...     score=0.6,
    ...     metadata=entity_meta,
    ...     source="entities"
    ... )
    >>> response = VectorSearchResponse(
    ...     question="菜品有哪些",
    ...     search_results=SearchResults(entities=[entity_item])
    ... )
    >>> print(response.question)
    '菜品有哪些'
    >>> print(response.search_results.entities[0].metadata.name)
    '西北菜'
    """

    question: str = Field(..., description="原始搜索问题")
    search_results: SearchResults = Field(..., description="按集合分组的搜索结果")


__all__ = [
    # Request
    "VectorSearchRequest",
    # Response
    "VectorSearchResponse",
    "SearchResults",
    "SearchResultItem",
    # Metadata models
    "EntityMetadata",
    "RelationshipMetadata",
    "ChunkMetadata",
    "PictureMetadata",
    "CommunityMetadata",
    "CommunityFinding",
]
