#!/usr/bin/env python
"""
知识抽取脚本
示例用法：python scripts/extract.py --input data/sample.txt --output output/entities.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.extraction.ner import NERExtractor
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(
        "--input",
        type=str,
        required=True,
        help="输入文件路径（.docx或.txt）"
    )
    parser.add_argument(
        "--output",
        type=str,
        default="output/entities.json",
        help="输出文件路径"
    )
    parser.add_argument(
        "--use-verification",
        action="store_true",
        default=True,
        help="使用二次对话验证"
    )
    
    args = parser.parse_args()
    
    # 加载配置
    config = load_config()
    
    # 确保输出目录存在
    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    
    # 解析文档
    input_path = Path(args.input)
    if input_path.suffix == ".docx":
        parser = DocumentParser()
        doc_data = parser.parse_docx(str(input_path))
        text = doc_data["text"]
    else:
        # 假设是纯文本文件
        with open(input_path, "r", encoding="utf-8") as f:
            text = f.read()
    
    # 执行NER
    logger.info("开始实体识别...")
    ner_extractor = NERExtractor(config=config)
    entities = ner_extractor.extract(text, use_verification=args.use_verification)
    
    # 保存结果
    result = {
        "input_file": str(input_path),
        "entities": entities,
        "statistics": {
            entity_type: len(entity_list)
            for entity_type, entity_list in entities.items()
        }
    }
    
    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(result, f, ensure_ascii=False, indent=2)
    
    logger.info(f"实体识别完成，结果已保存到: {output_path}")
    logger.info(f"统计信息: {result['statistics']}")


if __name__ == "__main__":
    main()