#!/usr/bin/env python
"""
测试文档解析（仅测试3个文件，不调用LLM）
"""

import sys
import json
import logging
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()]
sys.path.insert(0, str(Path(__file__).parent.parent))

from src.preprocessing.document_parser import DocumentParser
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("测试文档解析（3个文件）")
    logger.info("=" * 60)
    
    # 配置
    config = load_config()
    data_dir = Path("../data/1法律")
    output_dir = Path("./output")
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # 解析文档
    logger.info(f"解析目录: {data_dir}")
    parser = DocumentParser()
    docs = parser.parse_directory(str(data_dir))
    
    # 只处理前3个
    docs = docs[:3]
    logger.info(f"处理 {len(docs)} 个文档")
    
    # 切分为TextUnit
    logger.info("切分文档为TextUnit...")
    all_textunits = []
    for i, doc in enumerate(docs, 1):
        logger.info(f"处理文档 {i}/{len(docs)}: {Path(doc.get('file_path', '')).name}")
        textunits = parser.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(textunits)} 个TextUnit")
    
    logger.info(f"共生成 {len(all_textunits)} 个TextUnit")
    
    # 保存结果
    output_file = output_dir / "test_textunits.json"
    result = {
        "num_docs": len(docs),
        "num_textunits": len(all_textunits),
        "docs": [
            {
                "file_path": doc.get("file_path", ""),
                "title": doc.get("title", ""),
                "text_length": len(doc.get("text", "")),
            }
            for doc in docs
        ],
        "sample_textunits": [
            {
                "id": tu.get("id", ""),
                "text": tu.get("text", "")[:200] + "..." if len(tu.get("text", "")) > 200 else tu.get("text", ""),
                "paragraph_index": tu.get("paragraph_index", 0),
            }
            for tu in all_textunits[:10]  # 只保存前10个作为示例
        ]
    }
    
    with open(output_file, "w", encoding="utf-8") as f:
        json.dump(result, f, ensure_ascii=False, indent=2)
    
    logger.info("=" * 60)
    logger.info("测试完成！")
    logger.info(f"结果已保存到: {output_file}")
    logger.info(f"统计信息:")
    logger.info(f"  - 文档数: {result['num_docs']}")
    logger.info(f"  - TextUnit数: {result['num_textunits']}")
    logger.info("=" * 60)


if __name__ == "__main__":
    main()