"""知识图谱图查询数据模型。

定义基于图数据库的查询结果数据结构，
用于实体关系探索、子图查询、路径发现等场景。
"""

from __future__ import annotations

from pydantic import BaseModel, ConfigDict, Field

from hezor_common.utilities.text_helpers import to_camel

__all__ = [
    "CoOccurrence",
    "CoOccurrenceList",
    "GraphEdge",
    "GraphNode",
    "GraphPath",
    "GraphQueryResult",
    "GraphStatistics",
    "PathResult",
    "SubGraph",
]


# ─────────────────── 图节点与关系 ───────────────────


class GraphNode(BaseModel):
    """图数据库节点。

    表示图中的一个节点（Entity 或 Community），
    包含标签和属性信息。

    Parameters
    ----------
    id : str
        节点 ID（Neo4j 内部 ID 或业务 ID）。
    labels : list[str]
        节点标签列表（如 ``["Entity"]``）。
    name : str
        节点名称。
    entity_type : str
        实体类型（如 ``"Category-品类"``）。
    description : str
        节点描述。
    properties : dict
        节点的全部属性。
    """

    model_config = ConfigDict(
        alias_generator=to_camel,
        populate_by_name=True,
    )

    id: str = Field(..., description="节点 ID")
    labels: list[str] = Field(default_factory=list, description="节点标签列表")
    name: str = Field(default="", description="节点名称")
    entity_type: str = Field(default="", description="实体类型")
    description: str = Field(default="", description="节点描述")
    properties: dict = Field(default_factory=dict, description="节点全部属性")


class GraphEdge(BaseModel):
    """图数据库关系（边）。

    表示图中两个节点之间的一条有向关系，
    包含关系类型、权重和属性。

    Parameters
    ----------
    id : str
        关系 ID。
    source_id : str
        起始节点 ID。
    target_id : str
        目标节点 ID。
    source_name : str
        起始节点名称。
    target_name : str
        目标节点名称。
    relationship_type : str
        关系类型（如 ``"RELATED_TO"``）。
    weight : float
        关系权重。
    description : str
        关系描述。
    properties : dict
        关系全部属性。
    """

    model_config = ConfigDict(
        alias_generator=to_camel,
        populate_by_name=True,
    )

    id: str = Field(default="", description="关系 ID")
    source_id: str = Field(..., description="起始节点 ID")
    target_id: str = Field(..., description="目标节点 ID")
    source_name: str = Field(default="", description="起始节点名称")
    target_name: str = Field(default="", description="目标节点名称")
    relationship_type: str = Field(default="", description="关系类型")
    weight: float = Field(default=1.0, description="关系权重")
    description: str = Field(default="", description="关系描述")
    properties: dict = Field(default_factory=dict, description="关系全部属性")


# ─────────────────── 子图 ───────────────────


class SubGraph(BaseModel):
    """子图结构。

    包含一组节点和连接它们的关系（边），
    用于前端可视化渲染。

    Parameters
    ----------
    nodes : list[GraphNode]
        子图中的所有节点。
    edges : list[GraphEdge]
        子图中的所有关系。
    center_node_id : str
        中心节点 ID（可选，用于定位视图）。
    """

    model_config = ConfigDict(
        alias_generator=to_camel,
        populate_by_name=True,
    )

    nodes: list[GraphNode] = Field(default_factory=list, description="节点列表")
    edges: list[GraphEdge] = Field(default_factory=list, description="关系列表")
    center_node_id: str = Field(default="", description="中心节点 ID")


# ─────────────────── 路径 ───────────────────


class GraphPath(BaseModel):
    """两个节点之间的路径。

    表示图中从起点到终点的一条路径，
    由交替的节点和关系组成。

    Parameters
    ----------
    nodes : list[GraphNode]
        路径上的节点序列。
    edges : list[GraphEdge]
        路径上的关系序列。
    length : int
        路径长度（跳数）。
    """

    model_config = ConfigDict(
        alias_generator=to_camel,
        populate_by_name=True,
    )

    nodes: list[GraphNode] = Field(default_factory=list, description="路径节点序列")
    edges: list[GraphEdge] = Field(default_factory=list, description="路径关系序列")
    length: int = Field(default=0, description="路径长度")


class PathResult(BaseModel):
    """路径查询结果。

    Parameters
    ----------
    source_name : str
        起始实体名称。
    target_name : str
        目标实体名称。
    paths : list[GraphPath]
        找到的所有路径。
    """

    model_config = ConfigDict(
        alias_generator=to_camel,
        populate_by_name=True,
    )

    source_name: str = Field(default="", description="起始实体名称")
    target_name: str = Field(default="", description="目标实体名称")
    paths: list[GraphPath] = Field(default_factory=list, description="路径列表")


# ─────────────────── 图统计 ───────────────────


class GraphStatistics(BaseModel):
    """图数据库统计信息。

    Parameters
    ----------
    entity_count : int
        实体节点总数。
    community_count : int
        社区节点总数。
    relationship_count : int
        关系总数。
    entity_type_distribution : dict[str, int]
        实体类型分布。
    relationship_type_distribution : dict[str, int]
        关系类型分布。
    community_level_distribution : dict[str, int]
        社区层级分布。
    avg_relationships_per_entity : float
        每个实体的平均关系数。
    """

    model_config = ConfigDict(
        alias_generator=to_camel,
        populate_by_name=True,
    )

    entity_count: int = Field(default=0, description="实体节点总数")
    community_count: int = Field(default=0, description="社区节点总数")
    relationship_count: int = Field(default=0, description="关系总数")
    entity_type_distribution: dict[str, int] = Field(
        default_factory=dict, description="实体类型分布"
    )
    relationship_type_distribution: dict[str, int] = Field(
        default_factory=dict, description="关系类型分布"
    )
    community_level_distribution: dict[str, int] = Field(
        default_factory=dict, description="社区层级分布"
    )
    avg_relationships_per_entity: float = Field(default=0.0, description="每个实体的平均关系数")


# ─────────────────── 实体共现 ───────────────────


class CoOccurrence(BaseModel):
    """实体共现关系。

    表示两个实体通过共享关系或共同社区形成的
    关联强度。

    Parameters
    ----------
    entity_name : str
        关联实体名称。
    entity_type : str
        关联实体类型。
    co_occurrence_count : int
        共现次数。
    shared_relationship_types : list[str]
        共享的关系类型列表。
    """

    model_config = ConfigDict(
        alias_generator=to_camel,
        populate_by_name=True,
    )

    entity_name: str = Field(..., description="关联实体名称")
    entity_type: str = Field(default="", description="关联实体类型")
    co_occurrence_count: int = Field(default=0, description="共现次数")
    shared_relationship_types: list[str] = Field(
        default_factory=list, description="共享关系类型列表"
    )


class CoOccurrenceList(BaseModel):
    """实体共现列表。

    Parameters
    ----------
    center_entity : str
        中心实体名称。
    items : list[CoOccurrence]
        共现实体列表。
    """

    model_config = ConfigDict(
        alias_generator=to_camel,
        populate_by_name=True,
    )

    center_entity: str = Field(..., description="中心实体名称")
    items: list[CoOccurrence] = Field(default_factory=list, description="共现列表")


class GraphQueryResult(BaseModel):
    """图查询统一响应模型。

    根据 ``query_type`` 的不同，仅对应字段有值，
    其余字段为 None。

    Parameters
    ----------
    query_type : str
        执行的查询类型。
    statistics : GraphStatistics | None
        图统计信息（``graph_statistics``）。
    nodes : list[GraphNode] | None
        节点列表（``entity_search``、
        ``entity_communities``、``related_communities``）。
    subgraph : SubGraph | None
        子图（``entity_relationships``、
        ``entity_subgraph``、``community_subgraph``）。
    path_result : PathResult | None
        路径结果（``find_paths``）。
    co_occurrences : CoOccurrenceList | None
        共现列表（``entity_co_occurrence``）。
    """

    model_config = ConfigDict(
        alias_generator=to_camel,
        populate_by_name=True,
    )

    query_type: str = Field(..., description="查询类型")
    statistics: GraphStatistics | None = Field(
        default=None,
        description="图统计信息",
    )
    nodes: list[GraphNode] | None = Field(
        default=None,
        description="节点列表",
    )
    subgraph: SubGraph | None = Field(
        default=None,
        description="子图",
    )
    path_result: PathResult | None = Field(
        default=None,
        description="路径结果",
    )
    co_occurrences: CoOccurrenceList | None = Field(
        default=None,
        description="共现列表",
    )
