#!/usr/bin/env python
"""
完整流程脚本：从文档到知识图谱构建和查询
示例用法：python scripts/complete_pipeline.py --data-dir ../data/1法律 --output-dir ./output
"""

import argparse
import logging
import sys
from pathlib import Path

# 修复：移除conda环境的路径，确保使用虚拟环境的包
sys.path = [p for p in sys.path if 'anaconda3' not in p.lower() and 'conda' not in p.lower()]
# 添加src到路径
sys.path.insert(0, str(Path(__file__).parent.parent))

from src.preprocessing.document_parser import DocumentParser
from src.kg_builder.indexer import GraphIndexer
from src.kg_builder.graph import KnowledgeGraph
from src.utils.config import Config, load_config

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


def main():
    parser = argparse.ArgumentParser(description="完整的知识图谱构建流程")
    parser.add_argument(
        "--data-dir",
        type=str,
        required=True,
        help="法规数据目录路径"
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        default="output",
        help="输出目录"
    )
    parser.add_argument(
        "--max-docs",
        type=int,
        default=None,
        help="最大处理文档数（用于测试）"
    )
    parser.add_argument(
        "--skip-index",
        action="store_true",
        help="跳过索引构建（使用已有索引）"
    )
    
    args = parser.parse_args()
    
    # 加载配置
    config = load_config()
    
    # 创建输出目录
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    kg_file = output_dir / "knowledge_graph.json"
    
    # 步骤1: 构建知识图谱索引
    if not args.skip_index:
        logger.info("=" * 60)
        logger.info("步骤1: 解析文档并构建知识图谱")
        logger.info("=" * 60)
        
        # 解析文档
        logger.info(f"解析目录: {args.data_dir}")
        parser_obj = DocumentParser()
        docs = parser_obj.parse_directory(args.data_dir)
        
        if args.max_docs:
            docs = docs[:args.max_docs]
        
        logger.info(f"找到 {len(docs)} 个文档")
        
        # 切分为TextUnit
        logger.info("切分文档为TextUnit...")
        all_textunits = []
        for doc in docs:
            textunits = parser_obj.split_into_textunits(
                doc,
                max_length=config.MAX_TEXTUNIT_LENGTH
            )
            for tu in textunits:
                tu["doc_id"] = doc.get("file_path", "")
                tu["id"] = f"{doc.get('file_path', '')}_{tu.get('paragraph_index', 0)}"
            all_textunits.extend(textunits)
        
        logger.info(f"共生成 {len(all_textunits)} 个TextUnit")
        
        # 构建知识图谱
        logger.info("开始构建知识图谱...")
        indexer = GraphIndexer(config=config)
        kg_graph = indexer.index(all_textunits, use_verification=True)
        
        # 保存图谱
        kg_obj = KnowledgeGraph()
        kg_obj.graph = kg_graph
        kg_obj.save(str(kg_file), format="json")
        
        logger.info(f"知识图谱已保存到: {kg_file}")
    else:
        logger.info(f"跳过索引构建，从文件加载: {kg_file}")
        kg_obj = KnowledgeGraph()
        kg_obj.load(str(kg_file), format="json")
    
    # 步骤2: 展示图谱统计信息
    logger.info("=" * 60)
    logger.info("步骤2: 知识图谱统计信息")
    logger.info("=" * 60)
    
    if kg_obj.graph:
        metadata = kg_obj.graph.graph.get("metadata", {})
        logger.info(f"节点数: {kg_obj.graph.number_of_nodes()}")
        logger.info(f"边数: {kg_obj.graph.number_of_edges()}")
        logger.info(f"三元组数: {metadata.get('num_triplets', 0)}")
        logger.info(f"社区数: {metadata.get('num_communities', 0)}")
        
        # 展示一些社区摘要
        summaries = metadata.get("summaries", {})
        if summaries:
            logger.info("\n前5个社区摘要：")
            for i, (comm_id, comm_data) in enumerate(list(summaries.items())[:5], 1):
                logger.info(f"{i}. {comm_id}: {comm_data.get('summary', '无摘要')[:100]}...")
    
    logger.info("\n知识图谱构建流程完成！")


if __name__ == "__main__":
    main()