# -*- coding: utf-8 -*-
"""
知识图谱构建引擎
基于 NetworkX 构建科创企业供应链风险传染图谱
节点类型：企业、供应商、客户、核心人员、制裁实体
边类型：供应关系、客户关系、任职关系、投资关系
"""
import json
import logging
from pathlib import Path

import networkx as nx

logger = logging.getLogger(__name__)

DATA_DIR = Path(__file__).parent.parent / "data"


def build_graph() -> nx.DiGraph:
    """
    构建完整的科创企业供应链风险知识图谱
    """
    G = nx.DiGraph()

    # 1. 加载企业数据，添加企业节点
    _add_company_nodes(G)

    # 2. 加载供应链数据，添加关系边
    _add_supply_chain_edges(G)

    # 3. 加载实体清单，标记受制裁节点
    _mark_sanctioned_nodes(G)

    logger.info(f"图谱构建完成: {G.number_of_nodes()} 节点, {G.number_of_edges()} 边")
    return G


def _add_company_nodes(G: nx.DiGraph):
    """添加科创板企业节点及其关联的核心人员节点"""
    filepath = DATA_DIR / "sample_companies.json"
    with open(filepath, "r", encoding="utf-8") as f:
        companies = json.load(f)

    for company in companies:
        name = company["short_name"]
        G.add_node(
            name,
            node_type="科创企业",
            stock_code=company["stock_code"],
            industry=company["industry"],
            sector=company["sector"],
            is_sanctioned=company["compliance"]["entity_list_status"] != "未列入",
            risk_level="正常",
            color="#4CAF50",  # 默认绿色
        )

        # 添加核心技术人员节点
        for person in company.get("core_tech_personnel", []):
            person_id = f"{person['name']}@{name}"
            G.add_node(
                person_id,
                node_type="核心人员",
                real_name=person["name"],
                title=person["title"],
                status=person["status"],
                importance=person["importance"],
                company=name,
                color="#2196F3",  # 蓝色
            )
            G.add_edge(
                person_id, name,
                relation="任职于",
                edge_type="personnel",
            )

        # 添加供应商节点
        for supplier in company.get("supply_chain", {}).get("key_suppliers", []):
            supplier_name = supplier.split("（")[0].split("(")[0].strip()
            if not G.has_node(supplier_name):
                G.add_node(
                    supplier_name,
                    node_type="供应商",
                    is_sanctioned=False,
                    risk_level="正常",
                    color="#FF9800",  # 橙色
                )
            G.add_edge(
                supplier_name, name,
                relation="供应",
                detail=supplier,
                edge_type="supply",
            )

        # 添加客户节点
        for customer in company.get("supply_chain", {}).get("key_customers", []):
            customer_name = customer.split("（")[0].split("(")[0].strip()
            if not G.has_node(customer_name):
                G.add_node(
                    customer_name,
                    node_type="客户",
                    is_sanctioned=False,
                    risk_level="正常",
                    color="#9C27B0",  # 紫色
                )
            G.add_edge(
                name, customer_name,
                relation="供货给",
                edge_type="customer",
            )


def _add_supply_chain_edges(G: nx.DiGraph):
    """从供应链关系文件添加更详细的边"""
    filepath = DATA_DIR / "supply_chain.json"
    with open(filepath, "r", encoding="utf-8") as f:
        data = json.load(f)

    # 添加供应关系
    for rel in data.get("supply_relations", []):
        from_node = rel["from"]
        to_node = rel["to"]

        # 确保节点存在
        if not G.has_node(from_node):
            G.add_node(from_node, node_type="供应商", is_sanctioned=False,
                       risk_level="正常", color="#FF9800")
        if not G.has_node(to_node):
            G.add_node(to_node, node_type="企业", is_sanctioned=False,
                       risk_level="正常", color="#4CAF50")

        G.add_edge(
            from_node, to_node,
            relation=rel["relation"],
            critical=rel.get("critical", False),
            status=rel.get("status", "正常"),
            edge_type="supply",
        )

    # 添加投资关系
    for rel in data.get("investment_relations", []):
        from_node = rel["from"]
        to_node = rel["to"]

        if not G.has_node(from_node):
            G.add_node(from_node, node_type="投资方", is_sanctioned=False,
                       risk_level="正常", color="#607D8B")

        G.add_edge(
            from_node, to_node,
            relation=rel["relation"],
            share_ratio=rel.get("share_ratio", 0),
            edge_type="investment",
        )


def _mark_sanctioned_nodes(G: nx.DiGraph):
    """标记受制裁的节点，并向下游传播风险"""
    filepath = DATA_DIR / "entity_list.json"
    with open(filepath, "r", encoding="utf-8") as f:
        entities = json.load(f)

    # 收集所有受制裁实体的名称和别名
    sanctioned_names = set()
    for entity in entities:
        sanctioned_names.add(entity["entity_name"])
        for alias in entity.get("aliases", []):
            sanctioned_names.add(alias)

    # 标记图谱中的受制裁节点
    for node in G.nodes():
        for sname in sanctioned_names:
            if node in sname or sname in node:
                G.nodes[node]["is_sanctioned"] = True
                G.nodes[node]["risk_level"] = "高危"
                G.nodes[node]["color"] = "#F44336"  # 红色
                break


def get_node_info(G: nx.DiGraph, node_name: str) -> dict:
    """获取节点详细信息"""
    if node_name not in G:
        return {"error": f"节点 '{node_name}' 不存在"}

    node_data = dict(G.nodes[node_name])
    predecessors = list(G.predecessors(node_name))
    successors = list(G.successors(node_name))

    return {
        "name": node_name,
        "attributes": node_data,
        "upstream": predecessors,
        "downstream": successors,
        "degree": G.degree(node_name),
    }


def get_graph_stats(G: nx.DiGraph) -> dict:
    """获取图谱统计信息"""
    node_types = {}
    for _, data in G.nodes(data=True):
        t = data.get("node_type", "未知")
        node_types[t] = node_types.get(t, 0) + 1

    sanctioned_count = sum(
        1 for _, data in G.nodes(data=True) if data.get("is_sanctioned", False)
    )

    return {
        "total_nodes": G.number_of_nodes(),
        "total_edges": G.number_of_edges(),
        "node_types": node_types,
        "sanctioned_nodes": sanctioned_count,
        "density": nx.density(G),
    }