# -*- coding: utf-8 -*-
"""
知识图谱可视化模块
使用 pyvis 生成交互式网络图，支持嵌入 Streamlit
"""
import logging
import tempfile
from pathlib import Path
from typing import Optional

import networkx as nx

logger = logging.getLogger(__name__)


# 节点类型对应的颜色和形状
NODE_STYLES = {
    "科创企业": {"color": "#4CAF50", "shape": "dot", "size": 30},
    "供应商": {"color": "#FF9800", "shape": "diamond", "size": 20},
    "客户": {"color": "#9C27B0", "shape": "triangle", "size": 20},
    "核心人员": {"color": "#2196F3", "shape": "star", "size": 15},
    "投资方": {"color": "#607D8B", "shape": "square", "size": 20},
    "企业": {"color": "#4CAF50", "shape": "dot", "size": 25},
}

# 受制裁节点的样式覆盖
SANCTIONED_STYLE = {"color": "#F44336", "size": 35}
# 受限边的样式
RESTRICTED_EDGE_STYLE = {"color": "#F44336", "dashes": True, "width": 3}


def generate_interactive_graph(
    G: nx.DiGraph,
    highlight_company: Optional[str] = None,
    output_path: Optional[str] = None,
    height: str = "600px",
    width: str = "100%",
) -> str:
    """
    生成交互式知识图谱 HTML
    """
    try:
        from pyvis.network import Network
    except ImportError:
        logger.error("pyvis 未安装，请运行: pip install pyvis")
        return _generate_fallback_html(G)

    net = Network(
        height=height,
        width=width,
        directed=True,
        notebook=False,
        bgcolor="#1a1a2e",
        font_color="white",
    )

    # 物理引擎配置
    net.set_options("""
    {
        "physics": {
            "forceAtlas2Based": {
                "gravitationalConstant": -50,
                "centralGravity": 0.01,
                "springLength": 150,
                "springConstant": 0.08
            },
            "solver": "forceAtlas2Based",
            "stabilization": {"iterations": 100}
        },
        "interaction": {
            "hover": true,
            "tooltipDelay": 200,
            "navigationButtons": true
        }
    }
    """)

    # 添加节点
    for node, data in G.nodes(data=True):
        node_type = data.get("node_type", "企业")
        style = NODE_STYLES.get(node_type, NODE_STYLES["企业"]).copy()

        # 受制裁节点特殊样式
        if data.get("is_sanctioned", False):
            style.update(SANCTIONED_STYLE)

        # 高亮选中企业
        if highlight_company and node == highlight_company:
            style["color"] = "#FFD700"  # 金色
            style["size"] = 45
            style["borderWidth"] = 3

        # 构建标签和悬浮提示
        label = node.split("@")[0] if "@" in node else node
        title_parts = [f"<b>{label}</b>", f"类型: {node_type}"]
        if data.get("is_sanctioned"):
            title_parts.append("⚠️ <b>已被制裁</b>")
        if data.get("stock_code"):
            title_parts.append(f"代码: {data['stock_code']}")
        if data.get("industry"):
            title_parts.append(f"行业: {data['industry']}")
        if data.get("title"):
            title_parts.append(f"职位: {data['title']}")
        if data.get("status"):
            title_parts.append(f"状态: {data['status']}")

        net.add_node(
            node,
            label=label,
            title="<br>".join(title_parts),
            color=style["color"],
            shape=style["shape"],
            size=style["size"],
        )

    # 添加边
    for u, v, data in G.edges(data=True):
        edge_style = {
            "color": "#666666",
            "width": 1,
            "arrows": "to",
        }

        # 受限边特殊样式
        if data.get("status") == "受限":
            edge_style.update(RESTRICTED_EDGE_STYLE)
        elif data.get("critical", False):
            edge_style["width"] = 2
            edge_style["color"] = "#FFC107"  # 关键边用黄色

        relation = data.get("relation", "")
        net.add_edge(
            u, v,
            title=relation,
            label=relation if len(relation) <= 8 else "",
            **edge_style,
        )

    # 输出 HTML
    if output_path is None:
        output_path = str(Path(tempfile.gettempdir()) / "kg_visualization.html")

    net.save_graph(output_path)

    # 读取 HTML 内容
    with open(output_path, "r", encoding="utf-8") as f:
        return f.read()


def _generate_fallback_html(G: nx.DiGraph) -> str:
    """当 pyvis 不可用时的备用 HTML 可视化"""
    nodes_info = []
    for node, data in G.nodes(data=True):
        label = node.split("@")[0] if "@" in node else node
        is_sanctioned = data.get("is_sanctioned", False)
        nodes_info.append(f"<li style='color: {'red' if is_sanctioned else 'green'}'>{label} ({data.get('node_type', '未知')})</li>")

    return f"""
    <html><body style='background: #1a1a2e; color: white; padding: 20px;'>
    <h2>📊 知识图谱节点列表（pyvis 未安装，使用简化视图）</h2>
    <p>节点数: {G.number_of_nodes()} | 边数: {G.number_of_edges()}</p>
    <ul>{''.join(nodes_info[:50])}</ul>
    <p style='color: #888;'>安装 pyvis 以获得交互式可视化: pip install pyvis</p>
    </body></html>
    """


def get_subgraph_for_company(G: nx.DiGraph, company: str, depth: int = 2) -> nx.DiGraph:
    """
    提取以指定企业为中心的子图（上下游 N 层）
    """
    if company not in G:
        return nx.DiGraph()

    # 收集相关节点
    related_nodes = {company}

    # 上游（前驱）
    current_layer = {company}
    for _ in range(depth):
        next_layer = set()
        for node in current_layer:
            next_layer.update(G.predecessors(node))
        related_nodes.update(next_layer)
        current_layer = next_layer

    # 下游（后继）
    current_layer = {company}
    for _ in range(depth):
        next_layer = set()
        for node in current_layer:
            next_layer.update(G.successors(node))
        related_nodes.update(next_layer)
        current_layer = next_layer

    return G.subgraph(related_nodes).copy()