#!/usr/bin/env python
"""
知识图谱构建脚本
示例用法：python scripts/build_kg.py --data-dir ../data/1法律 --output ./output/kg.json
"""

import argparse
import json
import logging
import sys
from pathlib import Path

# 添加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",
        type=str,
        default="output/kg.json",
        help="输出文件路径"
    )
    parser.add_argument(
        "--max-docs",
        type=int,
        default=None,
        help="最大处理文档数（用于测试）"
    )
    parser.add_argument(
        "--no-verification",
        action="store_true",
        help="不使用二次对话验证（加快速度）"
    )
    
    args = parser.parse_args()
    
    # 加载配置
    config = load_config()
    
    # 确保输出目录存在
    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    
    # 解析文档
    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=not args.no_verification
    )
    
    # 保存图谱
    kg_obj = KnowledgeGraph()
    kg_obj.graph = kg_graph
    kg_obj.save(str(output_path), format="json")
    
    # 保存统计信息
    stats = {
        "num_docs": len(docs),
        "num_textunits": len(all_textunits),
        "num_nodes": kg_graph.number_of_nodes(),
        "num_edges": kg_graph.number_of_edges(),
        "metadata": kg_graph.graph.get("metadata", {})
    }
    
    stats_path = output_path.parent / f"{output_path.stem}_stats.json"
    with open(stats_path, "w", encoding="utf-8") as f:
        json.dump(stats, f, ensure_ascii=False, indent=2)
    
    logger.info(f"知识图谱构建完成！")
    logger.info(f"  - 节点数: {stats['num_nodes']}")
    logger.info(f"  - 边数: {stats['num_edges']}")
    logger.info(f"  - 结果已保存到: {output_path}")
    logger.info(f"  - 统计信息已保存到: {stats_path}")


if __name__ == "__main__":
    main()