"""
graph_builder.py - 知识图谱构建的核心流程控制器
包含完整的构建流程和交互逻辑
"""
import os
import json
try:
    import faiss
except ImportError:  # pragma: no cover - depends on deployment environment
    faiss = None
try:
    import numpy as np
except ImportError:  # pragma: no cover - depends on deployment environment
    np = None
import time
from datetime import datetime
from pathlib import Path
import logging
from tqdm import tqdm
from collections import defaultdict

from models import GraphNode, GraphEdge, NodeType, EdgeType
from code_parser import (
    parser, CPP_EXTENSIONS, IGNORE_DIRS,
    safe_decode, extract_comment_before, extract_code_snippet,
    extract_full_function_name, collect_call_expressions,
    extract_class_name, extract_member_function_name, extract_base_classes
)
from llm_processor import generate_code_logic, generate_code_summary, get_qwen_embedding
from config import PROJECT_ROOT, MAX_CODE_LENGTH, QWEN_EMBEDDING_MODEL

logger = logging.getLogger(__name__)

# 全局图数据结构
FUNCTION_CALL_GRAPH = defaultdict(list)
FILE_DEPENDENCIES = defaultdict(list)
CALLED_BY_GRAPH = defaultdict(list)


def _as_node_type(value):
    return value.value if hasattr(value, "value") else value


def _normalize_output_dir(output_dir):
    if output_dir is None:
        return Path(".")
    path = Path(output_dir)
    path.mkdir(parents=True, exist_ok=True)
    return path


def _qualified_name(node):
    if node.id.startswith("Function:"):
        return node.id[len("Function:"):]
    return node.name


def _has_real_embedding(embedding):
    if not embedding:
        return False
    try:
        return any(abs(float(value)) > 1e-12 for value in embedding)
    except (TypeError, ValueError):
        return False


def _function_metadata(node, embedding_dim=None):
    raw = node.raw_attributes or {}
    embedding_available = raw.get("embedding_available")
    if embedding_available is None:
        embedding_available = _has_real_embedding(node.embedding)
    return {
        "node_id": node.id,
        "name": node.name,
        "function_name": node.name,
        "qualified_name": raw.get("qualified_name") or _qualified_name(node),
        "file": node.file_path,
        "file_path": node.file_path,
        "start_line": node.start_line,
        "end_line": node.end_line,
        "signature": node.signature,
        "summary": node.summary or "",
        "logic_flow": node.logic_flow or "",
        "code_snippet": raw.get("code_snippet", ""),
        "calls": raw.get("calls", []),
        "called_by": raw.get("called_by", []),
        "includes": raw.get("includes", []),
        "embedding_model": QWEN_EMBEDDING_MODEL,
        "embedding_dim": embedding_dim or (len(node.embedding) if node.embedding else 0),
        "embedding_available": bool(embedding_available),
        "embedding_error": raw.get("embedding_error", ""),
    }


class SymbolResolver:
    """符号解析器：将代码符号映射为图谱节点ID"""

    def __init__(self):
        self._symbol_table = {}
        self._file_to_entities = defaultdict(list)

    def register_entity(self, node: GraphNode):
        """注册新发现的实体到符号表"""
        self._symbol_table[node.name] = node.id
        if node.file_path:
            self._file_to_entities[node.file_path].append(node.id)

    def resolve_name(self, name: str, current_context: dict) -> str:
        """
        解析代码中的名称到节点ID
        :param name: 代码中的名称
        :param current_context: 当前上下文（文件、类等）
        :return: 节点ID或None
        """
        if name in self._symbol_table:
            return self._symbol_table[name]

        for qualified_name, node_id in self._symbol_table.items():
            if qualified_name.endswith('::' + name) or qualified_name == name:
                return node_id

        return None


def process_single_file(file_path, all_nodes, all_edges, symbol_resolver, function_raw_info_map):
    """处理单个C/C++文件，提取代码实体和结构关系"""
    logger.info(f"处理单个文件: {file_path}")

    try:
        with open(file_path, 'rb') as f:
            source_bytes = f.read()
        source_lines = source_bytes.split(b'\n')
        tree = parser.parse(source_bytes)

        # 1. 创建文件节点
        file_node_id = f"File:{file_path}"
        file_node = GraphNode(
            id=file_node_id,
            type=NodeType.FILE,
            name=Path(file_path).name,
            file_path=file_path
        )
        all_nodes.append(file_node)
        symbol_resolver.register_entity(file_node)

        # 2. 收集头文件依赖
        includes = []
        for node in tree.root_node.named_children:
            if node.type == "preproc_include":
                path_node = node.child_by_field_name("path")
                if path_node:
                    include_text = path_node.text.decode('utf-8').strip('"<>')
                    includes.append(include_text)
        FILE_DEPENDENCIES[file_path] = includes

        # 3. 遍历文件内的顶级声明
        for node in tree.root_node.named_children:
            # 3.1 处理函数定义
            if node.type == "function_definition":
                func_name = extract_full_function_name(node)
                if func_name == "<unknown>":
                    logger.warning(f"    跳过无法解析函数名的节点: {node.text[:50]}...")
                    continue

                # 创建函数节点
                func_node_id = f"Function:{func_name}"
                func_node = GraphNode(
                    id=func_node_id,
                    type=NodeType.FUNCTION,
                    name=func_name,
                    signature=node.text.decode('utf-8', errors='ignore').split('{')[0].strip() + ' {...}',
                    file_path=file_path,
                    start_line=node.start_point[0] + 1,
                    end_line=node.end_point[0] + 1,
                    raw_attributes={'ast_node_type': 'function_definition'}
                )
                all_nodes.append(func_node)
                symbol_resolver.register_entity(func_node)

                # 存储原始信息
                comment = extract_comment_before(node, source_lines)
                code_snippet = extract_code_snippet(tree, node, source_lines)
                function_raw_info_map[func_node_id] = {
                    'ast_node': node,
                    'comment': comment,
                    'code_snippet': code_snippet,
                    'source_lines': source_lines,
                    'tree': tree
                }

                # 添加 FILE -> CONTAINS -> FUNCTION 边
                all_edges.append(GraphEdge(
                    source_id=file_node_id,
                    target_id=func_node_id,
                    type=EdgeType.CONTAINS
                ))

                # 收集函数调用关系
                callees = collect_call_expressions(node)
                FUNCTION_CALL_GRAPH[func_name] = callees

            # 3.2 处理类/结构体定义
            elif node.type in ["class_specifier", "struct_specifier"]:
                class_name = extract_class_name(node)
                if not class_name or class_name == "<unknown>":
                    logger.debug(f"    跳过无法解析类名的节点")
                    continue

                # 创建类节点
                class_node_id = f"Class:{class_name}"
                class_node = GraphNode(
                    id=class_node_id,
                    type=NodeType.CLASS,
                    name=class_name,
                    signature=node.text.decode('utf-8', errors='ignore').split('{')[0].strip() + ' {...}',
                    file_path=file_path,
                    start_line=node.start_point[0] + 1,
                    end_line=node.end_point[0] + 1,
                    raw_attributes={
                        'ast_type': node.type,
                        'is_struct': node.type == "struct_specifier"
                    }
                )
                all_nodes.append(class_node)
                symbol_resolver.register_entity(class_node)

                # 添加 FILE -> CONTAINS -> CLASS 边
                all_edges.append(GraphEdge(
                    source_id=file_node_id,
                    target_id=class_node_id,
                    type=EdgeType.CONTAINS
                ))

                # 处理基类/继承关系
                base_clause = node.child_by_field_name("base_clause")
                if base_clause:
                    base_classes = extract_base_classes(base_clause)
                    for base_class_info in base_classes:
                        base_name = base_class_info.get('name')
                        access_type = base_class_info.get('access', 'public')
                        edge_type = EdgeType.EXTENDS if access_type in ['public', 'protected',
                                                                        'private'] else EdgeType.IMPLEMENTS

                        if base_name:
                            base_id = symbol_resolver.resolve_name(
                                base_name,
                                {'current_file': file_path, 'current_class': class_name}
                            )
                            if base_id:
                                all_edges.append(GraphEdge(
                                    source_id=class_node_id,
                                    target_id=base_id,
                                    type=edge_type,
                                    properties={'access': access_type}
                                ))

                # 处理类内方法
                for member in node.named_children:
                    if member.type == "function_definition":
                        member_func_name = extract_member_function_name(member, class_name)
                        if not member_func_name or member_func_name == "<unknown>":
                            continue

                        member_node_id = f"Function:{class_name}::{member_func_name}"
                        existing_member_node = next((n for n in all_nodes if n.id == member_node_id), None)

                        if not existing_member_node:
                            member_node = GraphNode(
                                id=member_node_id,
                                type=NodeType.FUNCTION,
                                name=member_func_name,
                                signature=member.text.decode('utf-8', errors='ignore').split('{')[0].strip() + ' {...}',
                                file_path=file_path,
                                start_line=member.start_point[0] + 1,
                                end_line=member.end_point[0] + 1,
                                raw_attributes={
                                    'is_member': True,
                                    'parent_class': class_name,
                                    'ast_node_type': 'member_function_definition'
                                }
                            )
                            all_nodes.append(member_node)
                            symbol_resolver.register_entity(member_node)

                            # 存储原始信息
                            member_comment = extract_comment_before(member, source_lines)
                            member_code = extract_code_snippet(tree, member, source_lines)
                            function_raw_info_map[member_node_id] = {
                                'ast_node': member,
                                'comment': member_comment,
                                'code_snippet': member_code,
                                'source_lines': source_lines,
                                'tree': tree
                            }

                            # 收集成员函数调用
                            member_callees = collect_call_expressions(member)
                            FUNCTION_CALL_GRAPH[f"{class_name}::{member_func_name}"] = member_callees

                        # 添加类包含边
                        all_edges.append(GraphEdge(
                            source_id=class_node_id,
                            target_id=member_node_id,
                            type=EdgeType.CONTAINS,
                            properties={'member_type': 'method'}
                        ))

    except Exception as e:
        logger.error(f"解析文件 {file_path} 失败: {e}", exc_info=True)
        raise


def process_project_files(project_root, all_nodes, all_edges, symbol_resolver, function_raw_info_map):
    """处理整个项目目录"""
    logger.info(f"处理项目目录: {project_root}")

    for root, dirs, files in os.walk(project_root):
        # 跳过忽略目录
        dirs[:] = [d for d in dirs if d not in IGNORE_DIRS]

        for file in files:
            file_path_obj = Path(file)
            if file_path_obj.suffix.lower() not in CPP_EXTENSIONS:
                continue

            full_file_path = os.path.join(root, file)
            process_single_file(full_file_path, all_nodes, all_edges, symbol_resolver, function_raw_info_map)

    return len(files)


def semantic_enhancement(all_nodes, all_edges, function_raw_info_map):
    """阶段三：语义增强，为函数节点生成摘要与向量嵌入"""
    logger.info("阶段三：语义增强，为函数节点生成摘要与向量嵌入...")

    # 1. 过滤出需要处理的函数节点
    function_nodes = [n for n in all_nodes if n.type == NodeType.FUNCTION]
    nodes_with_raw_info = []
    nodes_without_raw_info = []

    for func_node in function_nodes:
        if func_node.id in function_raw_info_map:
            nodes_with_raw_info.append(func_node)
        else:
            nodes_without_raw_info.append(func_node)

    logger.info(f"有原始信息的节点: {len(nodes_with_raw_info)} 个")
    if nodes_without_raw_info:
        logger.warning(f"无原始信息的节点: {len(nodes_without_raw_info)} 个")

    # 2. 处理有原始信息的节点
    if nodes_with_raw_info:
        print(f"\n开始为 {len(nodes_with_raw_info)} 个函数生成摘要和向量...")
        print("这可能需要一些时间，请耐心等待...")

        processed_count = 0
        for func_node in tqdm(nodes_with_raw_info, desc="生成函数摘要", unit="个"):
            func_node_id = func_node.id
            func_name = func_node.name

            # 获取原始信息
            raw_info = function_raw_info_map[func_node_id]
            comment = raw_info['comment']
            code_snippet = raw_info['code_snippet']
            file_path = func_node.file_path
            func_node.raw_attributes.update({
                "comment": comment,
                "code_snippet": code_snippet,
                "includes": FILE_DEPENDENCIES.get(file_path, []),
                "qualified_name": _qualified_name(func_node),
            })

            # 1. 生成流程逻辑
            try:
                logic = generate_code_logic(func_name, comment, code_snippet, file_path)
            except Exception as e:
                logger.error(f"为函数 {func_name} 生成逻辑失败: {e}")
                logic = "逻辑生成失败。"

            # 2. 生成功能摘要
            try:
                summary = generate_code_summary(
                    func_name, comment, logic, code_snippet, file_path
                )
            except Exception as e:
                logger.error(f"为函数 {func_name} 生成摘要失败: {e}")
                summary = "摘要生成失败。"

            # 3. 更新节点信息
            func_node.summary = summary
            func_node.logic_flow = logic

            # 获取调用关系
            called_functions = FUNCTION_CALL_GRAPH.get(func_name, [])
            caller_functions = CALLED_BY_GRAPH.get(func_name, [])
            func_node.raw_attributes.update({
                'calls': called_functions,
                'called_by': caller_functions,
                'includes': FILE_DEPENDENCIES.get(file_path, [])
            })

            # 4. 生成向量嵌入
            context_text_for_embedding = (
                f"函数名: {func_name}\n"
                f"功能摘要: {summary}\n"
                f"流程逻辑: {logic}\n"
                f"所在文件: {os.path.basename(file_path)}"
            )

            try:
                embedding_vector = get_qwen_embedding(context_text_for_embedding)
                embedding_values = (
                    embedding_vector.tolist()
                    if hasattr(embedding_vector, "tolist")
                    else list(embedding_vector)
                )
                if not _has_real_embedding(embedding_values):
                    raise RuntimeError("Embedding API returned an empty or zero vector.")
                func_node.embedding = embedding_values
                func_node.raw_attributes["embedding_available"] = True
                func_node.raw_attributes["embedding_error"] = ""
            except Exception as e:
                logger.error(f"为函数 {func_name} 生成嵌入失败: {e}")
                func_node.embedding = [0.0] * 1024
                func_node.raw_attributes["embedding_available"] = False
                func_node.raw_attributes["embedding_error"] = str(e)[:500]

            processed_count += 1
            # 每处理5个函数显示一次进度
            if processed_count % 5 == 0:
                print(f"  已处理 {processed_count}/{len(nodes_with_raw_info)} 个函数")

            # 添加延迟避免API限流
            time.sleep(0.1)

        print(f"✓ 完成 {len(nodes_with_raw_info)} 个函数的语义增强")
    else:
        logger.warning("没有找到任何有原始信息的函数节点，跳过摘要生成。")

    # 3. 处理无原始信息的节点
    if nodes_without_raw_info:
        print(f"\n为 {len(nodes_without_raw_info)} 个无原始信息的节点生成基础信息...")
        for func_node in nodes_without_raw_info:
            func_node.summary = f"函数 {func_node.name}，位于 {func_node.file_path}"
            func_node.logic_flow = "无法生成逻辑流程（缺少原始代码信息）"
            func_node.embedding = [0.0] * 1024
            func_node.raw_attributes.update({
                "calls": FUNCTION_CALL_GRAPH.get(func_node.name, []),
                "called_by": CALLED_BY_GRAPH.get(func_node.name, []),
                "includes": FILE_DEPENDENCIES.get(func_node.file_path, []),
                "qualified_name": _qualified_name(func_node),
                "embedding_available": False,
                "embedding_error": "Missing source information.",
            })


def build_call_edges(all_nodes, all_edges, symbol_resolver):
    """阶段二：构建函数调用边"""
    logger.info("阶段二：解析符号，构建函数调用图 (CALLS 边)...")

    # 构建反向调用图
    CALLED_BY_GRAPH.clear()
    for caller_name, callee_names in FUNCTION_CALL_GRAPH.items():
        caller_node_id = f"Function:{caller_name}"
        # 确保调用者节点已存在
        if not any(n.id == caller_node_id for n in all_nodes):
            continue

        for callee_name in callee_names:
            # 通过符号解析器将名称解析为节点ID
            callee_node_id = symbol_resolver.resolve_name(
                callee_name,
                {'current_file': next((n.file_path for n in all_nodes if n.id == caller_node_id), None)}
            )
            if callee_node_id:
                # 添加 CALLS 边
                all_edges.append(GraphEdge(
                    source_id=caller_node_id,
                    target_id=callee_node_id,
                    type=EdgeType.CALLS
                ))
                # 更新反向图
                if callee_name not in CALLED_BY_GRAPH:
                    CALLED_BY_GRAPH[callee_name] = []
                CALLED_BY_GRAPH[callee_name].append(caller_name)

    logger.info(f"阶段二完成。已构建边数: {len(all_edges)}")


def save_results(all_nodes, all_edges, base_name, target_path, mode, output_dir=None):
    """阶段四：持久化存储结果"""
    logger.info("阶段四：持久化图谱与向量索引...")

    output_path = _normalize_output_dir(output_dir)
    graph_data_path = str(output_path / f"{base_name}_code_knowledge_graph.json")
    vector_db_path = str(output_path / f"{base_name}_rag.faiss")
    metadata_path = str(output_path / f"{base_name}_rag_metadata.json")

    # 1. 保存为图谱文件
    nodes_dict = [node.dict() for node in all_nodes]
    edges_dict = [edge.dict() for edge in all_edges]

    graph_data = {
        "metadata": {
            "project_root": target_path if mode == "full_project" else "single_file",
            "file_path": target_path if mode == "single_file" else None,
            "mode": mode,
            "generated_at": datetime.now().isoformat(),
            "total_nodes": len(all_nodes),
            "total_edges": len(all_edges)
        },
        "nodes": nodes_dict,
        "edges": edges_dict
    }

    with open(graph_data_path, 'w', encoding='utf-8') as f:
        json.dump(graph_data, f, indent=2, ensure_ascii=False, default=str)
    logger.info(f"    图谱结构已保存至: {graph_data_path}")
    print(f"✓ 图谱结构已保存: {graph_data_path}")

    # 2. 构建并保存向量索引
    function_nodes = [node for node in all_nodes if node.type == NodeType.FUNCTION and node.embedding]
    if function_nodes:
        dimension = len(function_nodes[0].embedding)
        metadata_for_faiss = []

        vectors = []
        for node in function_nodes:
            vectors.append(node.embedding)
            metadata_for_faiss.append(_function_metadata(node, dimension))

        if vectors:
            if faiss is not None:
                if np is None:
                    raise RuntimeError("numpy is required when writing a FAISS index.")
                vectors_np = np.array(vectors, dtype='float32')
                index = faiss.IndexFlatL2(dimension)
                index.add(vectors_np)
                faiss.write_index(index, vector_db_path)
            else:
                with open(vector_db_path, 'w', encoding='utf-8') as f:
                    json.dump(
                        {
                            "format": "simple_l2_vector_index",
                            "dimension": dimension,
                            "vectors": [[float(value) for value in vector] for vector in vectors],
                        },
                        f,
                        ensure_ascii=False,
                    )
            with open(metadata_path, 'w', encoding='utf-8') as f:
                json.dump(metadata_for_faiss, f, indent=2, ensure_ascii=False)
            logger.info(f"    向量索引已保存。函数节点数: {len(function_nodes)}")
            logger.info(f"    向量数据库: {vector_db_path}")
            logger.info(f"    元数据: {metadata_path}")
            print(f"✓ 向量索引已保存: {vector_db_path}")
    else:
        logger.warning("    没有可生成向量的函数节点，跳过FAISS索引构建。")
        print("⚠ 无向量数据，跳过FAISS索引构建")

    return graph_data_path, vector_db_path, metadata_path


def build_code_knowledge_base(target_path, output_dir=None, semantic=True, base_name=None, embedding_dim=1024):
    """Build a code knowledge base without interactive prompts."""
    target_path = str(Path(target_path).resolve())
    if not os.path.exists(target_path):
        raise FileNotFoundError(f"Target path does not exist: {target_path}")

    all_nodes = []
    all_edges = []
    symbol_resolver = SymbolResolver()
    function_raw_info_map = {}

    FUNCTION_CALL_GRAPH.clear()
    FILE_DEPENDENCIES.clear()
    CALLED_BY_GRAPH.clear()

    if os.path.isdir(target_path):
        mode = "full_project"
        process_project_files(target_path, all_nodes, all_edges, symbol_resolver, function_raw_info_map)
        base_name = base_name or "project"
    else:
        mode = "single_file"
        process_single_file(target_path, all_nodes, all_edges, symbol_resolver, function_raw_info_map)
        base_name = base_name or f"file_{Path(target_path).stem}"

    build_call_edges(all_nodes, all_edges, symbol_resolver)

    if semantic:
        semantic_enhancement(all_nodes, all_edges, function_raw_info_map)
    else:
        for node in all_nodes:
            if node.type == NodeType.FUNCTION:
                raw_info = function_raw_info_map.get(node.id, {})
                node.summary = node.summary or f"Function {node.name}, located at {node.file_path}"
                node.logic_flow = node.logic_flow or "Semantic enhancement was skipped."
                node.embedding = node.embedding or [0.0] * embedding_dim
                node.raw_attributes.update({
                    "comment": raw_info.get("comment", ""),
                    "code_snippet": raw_info.get("code_snippet", ""),
                    "calls": FUNCTION_CALL_GRAPH.get(node.name, []),
                    "called_by": CALLED_BY_GRAPH.get(node.name, []),
                    "includes": FILE_DEPENDENCIES.get(node.file_path, []),
                    "qualified_name": _qualified_name(node),
                    "embedding_available": False,
                    "embedding_error": "Semantic enhancement was skipped.",
                })

    return save_results(all_nodes, all_edges, base_name, target_path, mode, output_dir=output_dir)


def build_rag_database_interactive():
    """交互式知识库构建主流程"""
    logger.info("=== 交互式知识图谱构建流程 ===")

    # 显示模式选择菜单
    print("\n" + "=" * 50)
    print("知识库构建模式选择")
    print("=" * 50)
    print("1. 全项目模式 - 处理整个项目目录")
    print("2. 单文件模式 - 处理单个文件")
    print("3. 退出")
    print("=" * 50)

    choice = input("请选择模式 (1-3): ").strip()

    if choice == "3":
        logger.info("用户选择退出")
        return None, None

    if choice == "1":
        # 全项目模式 - 直接引导用户输入
        print("\n" + "=" * 60)
        print("全项目知识库构建模式")
        print("=" * 60)

        # 显示当前config.py中配置的默认路径
        current_default_path = str(PROJECT_ROOT) if PROJECT_ROOT else "未配置"
        if not current_default_path or current_default_path == ".":
            current_default_path = "当前工作目录"

        # 直接要求用户输入
        while True:
            project_prompt = f"请直接输入要构建知识库的项目根目录完整路径\n(当前config.py中的默认路径: {current_default_path}): "
            user_input = input(project_prompt).strip()

            if not user_input:
                # 用户没有输入，尝试使用config中的默认路径
                if PROJECT_ROOT and os.path.exists(PROJECT_ROOT) and os.path.isdir(PROJECT_ROOT):
                    target_path = PROJECT_ROOT
                    print(f"使用config.py中的默认路径: {target_path}")
                else:
                    print("⚠ 您没有输入路径，且config.py中的默认路径无效。")
                    retry = input("是否重新输入? (y/n): ").strip().lower()
                    if retry != 'y':
                        print("知识库构建取消。")
                        return None, None
                    continue
            else:
                # 使用用户输入的路径
                target_path = Path(user_input)

            # 验证路径
            if not os.path.exists(target_path):
                print(f"⚠ 路径不存在: {target_path}")
                retry = input("是否重新输入? (y/n): ").strip().lower()
                if retry != 'y':
                    print("知识库构建取消。")
                    return None, None
                continue

            if not os.path.isdir(target_path):
                print(f"⚠ 路径不是目录: {target_path}")
                retry = input("是否重新输入? (y/n): ").strip().lower()
                if retry != 'y':
                    print("知识库构建取消。")
                    return None, None
                continue

            break  # 路径有效，退出循环

        mode = "full_project"
        print(f"\n✓ 将处理整个项目: {target_path}")
        print("开始处理...")

    elif choice == "2":
        # 单文件模式
        while True:
            file_path = input("\n请输入要处理的文件完整路径: ").strip()
            if not file_path:
                print("输入为空，退出单文件模式")
                return None, None

            if not os.path.exists(file_path):
                print(f"错误: 文件不存在: {file_path}")
                continue

            file_ext = Path(file_path).suffix.lower()
            if file_ext not in CPP_EXTENSIONS:
                print(f"警告: 文件类型 {file_ext} 可能不是C/C++文件，继续吗? (y/n): ", end="")
                confirm = input().strip().lower()
                if confirm != 'y':
                    continue

            target_path = file_path
            mode = "single_file"
            print(f"\n将处理单个文件: {target_path}")
            print("开始处理...")
            break
    else:
        print("无效的选择，退出")
        return None, None

    # 确认开始
    print("\n即将开始构建知识库，确认继续? (y/n): ", end="")
    confirm = input().strip().lower()
    if confirm != 'y':
        print("操作已取消")
        return None, None

    # --- 初始化数据结构 ---
    all_nodes = []  # 存储 GraphNode 对象
    all_edges = []  # 存储 GraphEdge 对象
    symbol_resolver = SymbolResolver()  # 符号解析器
    function_raw_info_map = {}  # 临时存储函数的原始信息

    # 清空全局图
    FUNCTION_CALL_GRAPH.clear()
    FILE_DEPENDENCIES.clear()
    CALLED_BY_GRAPH.clear()

    # ========================
    # 阶段一：静态分析，提取节点与硬性边
    # ========================
    logger.info("阶段一：静态分析，提取代码实体与结构关系...")

    if choice == "1":
        # 全项目模式
        processed_files_count = process_project_files(target_path, all_nodes, all_edges,
                                                      symbol_resolver, function_raw_info_map)
        logger.info(f"已处理 {processed_files_count} 个文件")
    else:
        # 单文件模式
        process_single_file(target_path, all_nodes, all_edges, symbol_resolver, function_raw_info_map)
        logger.info(f"已处理单个文件: {target_path}")

    logger.info(f"阶段一完成。初步识别节点数: {len(all_nodes)}")

    # ========================
    # 阶段二：解析符号，构建准确的调用边
    # ========================
    build_call_edges(all_nodes, all_edges, symbol_resolver)

    # ========================
    # 阶段三：语义增强，生成摘要与向量
    # ========================
    print("\n" + "=" * 50)
    print("阶段三：语义增强")
    print("=" * 50)
    print("此阶段将调用LLM API为函数生成摘要和向量嵌入")
    print(f"需要处理的函数节点数: {len([n for n in all_nodes if n.type == NodeType.FUNCTION])}")
    print("这可能会消耗API调用额度并需要一些时间")
    print("-" * 50)

    proceed = input("是否继续阶段三? (y/n): ").strip().lower()
    if proceed != 'y':
        logger.info("用户跳过阶段三（语义增强）")
        print("跳过阶段三，直接进入阶段四")

        # 为未处理的函数节点设置默认值
        for node in all_nodes:
            if node.type == NodeType.FUNCTION and not node.summary:
                node.summary = f"函数 {node.name}，位于 {node.file_path}"
                node.logic_flow = "跳过语义增强阶段"
                node.embedding = [0.0] * 1024
                raw_info = function_raw_info_map.get(node.id, {})
                node.raw_attributes.update({
                    "comment": raw_info.get("comment", ""),
                    "code_snippet": raw_info.get("code_snippet", ""),
                    "calls": FUNCTION_CALL_GRAPH.get(node.name, []),
                    "called_by": CALLED_BY_GRAPH.get(node.name, []),
                    "includes": FILE_DEPENDENCIES.get(node.file_path, []),
                    "qualified_name": _qualified_name(node),
                    "embedding_available": False,
                    "embedding_error": "Semantic enhancement was skipped.",
                })
    else:
        semantic_enhancement(all_nodes, all_edges, function_raw_info_map)

    # ========================
    # 阶段四：持久化存储
    # ========================
    print("\n" + "=" * 50)
    print("阶段四：持久化存储")
    print("=" * 50)

    save_choice = input("是否保存结果? (y/n): ").strip().lower()
    if save_choice != 'y':
        logger.info("用户取消保存，退出构建流程")
        print("构建流程已取消")
        return all_nodes, all_edges

    # 确定保存文件名
    if choice == "1":
        base_name = "project"
    else:
        file_name = Path(target_path).stem
        base_name = f"file_{file_name}"

    graph_data_path, vector_db_path, metadata_path = save_results(
        all_nodes, all_edges, base_name, target_path, mode
    )

    print("\n" + "=" * 50)
    print("构建完成！")
    print("=" * 50)
    print(f"总节点数: {len(all_nodes)}")
    print(f"总边数: {len(all_edges)}")
    print(f"图谱文件: {graph_data_path}")
    if os.path.exists(vector_db_path):
        print(f"向量文件: {vector_db_path}")

    logger.info("=== 知识图谱构建流程全部完成 ===")
    return all_nodes, all_edges
