#!/usr/bin/env python
"""
快速测试：只处理1个文档的前2个TextUnit
"""

import sys
import logging
from pathlib import Path

sys.path = [p for p in sys.path if 'anaconda3' not in p.lower() and 'conda' not in p.lower()]
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.utils.config import load_config

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

def main():
    logger.info("=" * 60)
    logger.info("快速测试：处理1个文档的前2个TextUnit")
    logger.info("=" * 60)
    
    config = load_config()
    
    # 解析1个文档
    parser = DocumentParser()
    doc_path = "../data/1法律/10-中华人民共和国乡村振兴促进法.docx"
    doc = parser.parse_docx(doc_path)
    
    # 切分为TextUnit，只取前2个
    textunits = parser.split_into_textunits(doc, max_length=config.MAX_TEXTUNIT_LENGTH)[:2]
    for tu in textunits:
        tu["doc_id"] = doc.get("file_path", "")
        tu["id"] = f"{doc.get('file_path', '')}_{tu.get('paragraph_index', 0)}"
    
    logger.info(f"准备处理 {len(textunits)} 个TextUnit")
    
    # 构建知识图谱（只处理这2个）
    logger.info("开始构建知识图谱...")
    indexer = GraphIndexer(config=config)
    kg_graph = indexer.index(textunits, use_verification=True)
    
    logger.info("=" * 60)
    logger.info("快速测试完成！")
    logger.info(f"节点数: {kg_graph.number_of_nodes()}")
    logger.info(f"边数: {kg_graph.number_of_edges()}")
    logger.info("=" * 60)

if __name__ == "__main__":
    main()