# -*- coding: utf-8 -*-
"""
重新抽取失败项目的脚本
"""

import asyncio
import pandas as pd
import yaml
import json
from pathlib import Path
from datetime import datetime
import sys

# 添加模块路径
sys.path.append(str(Path(__file__).parent / 'src'))

from knowledge_extraction.deep_entity_extractor import DeepEntityExtractor

async def retry_failed_projects():
    """重新抽取失败的项目"""

    # 加载配置
    config_file = Path(__file__).parent / 'config' / 'deep_extraction_config.yaml'
    with open(config_file, 'r', encoding='utf-8') as f:
        config = yaml.safe_load(f)

    # 读取原始数据
    data_file = Path(__file__).parent.parent.parent / 'data' / '黑龙江国家级和省级非遗名单.xlsx'
    df = pd.read_excel(data_file, engine='openpyxl')

    # 找到失败的项目
    # 第0列是序号，需要拼接成ICH-xxx格式
    failed_data = df[df.iloc[:, 0].isin([118, 229])]

    print(f"找到 {len(failed_data)} 个失败项目")
    print("=" * 60)

    # 初始化抽取器
    extractor = DeepEntityExtractor(str(config_file))

    results = []

    for idx, row in failed_data.iterrows():
        project_num = int(row.iloc[0])  # 序号（118, 229）
        project_id = f'ICH-{project_num}'  # 拼接成ICH-xxx格式
        project_name = row.iloc[3]  # 项目名称（第4列）
        description = row.iloc[7] if len(row) > 7 else ""  # 完整描述（第8列备注）

        print(f"\n正在抽取: {project_id} - {project_name}")
        print(f"描述长度: {len(description)} 字符")

        # 准备输入数据
        input_data = {
            'project_id': project_id,
            'project_name': project_name,
            'description': description
        }

        # 抽取实体和关系
        try:
            result = await extractor.extract_from_remark(
                project_id=project_id,
                project_name=project_name,
                remark_text=description
            )

            if result:
                results.append({
                    'project_id': project_id,
                    'project_name': project_name,
                    'extraction_result': result
                })
                print(f"[OK] 抽取成功: {len(result.get('entities', []))} 个实体, {len(result.get('relationships', []))} 条关系")
            else:
                print(f"[FAIL] 抽取失败")

        except Exception as e:
            print(f"[ERROR] 抽取异常: {str(e)}")

    # 保存结果
    output_file = Path(__file__).parent / 'output' / 'retry_projects.json'
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(results, f, ensure_ascii=False, indent=2)

    print(f"\n结果已保存到: {output_file}")
    print(f"成功: {len(results)}/{len(failed_data)}")

    return results


if __name__ == '__main__':
    asyncio.run(retry_failed_projects())