"""GraphRAG 融合查询的请求/响应数据模型。"""

from __future__ import annotations

from pydantic import BaseModel, Field


class RetrievalStrategy(BaseModel):
    """检索策略配置。"""

    vector_top_k: int = Field(default=5, ge=1, le=50, description="向量检索返回数")
    graph_depth: int = Field(default=2, ge=1, le=5, description="图谱扩展深度")
    graph_max_entities: int = Field(default=20, ge=1, le=100, description="图谱最大实体数")
    vector_weight: float = Field(default=0.6, ge=0.0, le=1.0, description="向量分数权重")
    graph_weight: float = Field(default=0.4, ge=0.0, le=1.0, description="图谱相关性权重")
    enable_graph: bool = Field(default=True, description="是否启用图谱检索")
    enable_vector: bool = Field(default=True, description="是否启用向量检索")


class GraphRAGQueryRequest(BaseModel):
    """GraphRAG 查询请求。"""

    query: str = Field(
        ...,
        min_length=1,
        max_length=2000,
        description="用户查询",
    )
    knowledge_base_id: str = Field(
        ...,
        min_length=1,
        max_length=64,
        description="知识库 ID，用于权限校验（由上游 Java 后端传入）",
    )
    collection_name: str = Field(
        ...,
        min_length=1,
        max_length=256,
        pattern=r"^[a-zA-Z0-9_\-\u4e00-\u9fff]+$",
        description="Milvus collection 名称（= 知识库名），仅允许字母、数字、下划线、连字符和中文",
    )
    graph_id: str = Field(
        ...,
        pattern=r"^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$",
        description="Neo4j 图谱 ID（UUID 格式）",
    )
    strategy: RetrievalStrategy = Field(
        default_factory=RetrievalStrategy,
        description="可选策略覆盖",
    )


class VectorChunk(BaseModel):
    """向量检索到的文档片段。"""

    id: str
    text: str
    score: float
    metadata: dict[str, object] = Field(default_factory=dict)


class EntitySummary(BaseModel):
    """实体摘要。"""

    id: str
    name: str
    type: str
    description: str = ""


class RelationSummary(BaseModel):
    """关系摘要。"""

    source_name: str
    source_type: str
    target_name: str
    target_type: str
    relation_type: str


class GraphContext(BaseModel):
    """图谱上下文。"""

    entities: list[EntitySummary] = Field(default_factory=list)
    relations: list[RelationSummary] = Field(default_factory=list)
    textualized: str = ""


class RetrievalContext(BaseModel):
    """检索上下文（检索结果的结构化表示）。"""

    vector_chunks: list[VectorChunk] = Field(default_factory=list)
    graph_context: GraphContext = Field(default_factory=GraphContext)
    merged_text: str = ""


class GraphRAGQueryResponse(BaseModel):
    """GraphRAG 查询响应。"""

    answer: str = Field(..., description="LLM 生成的回答")
    context: RetrievalContext = Field(..., description="检索上下文")
    model: str = Field(..., description="使用的 LLM 模型名")