# -*- 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
from data_processing.entity_normalizer import EntityNormalizer
from data_processing.relationship_builder import RelationshipBuilder

async def retry_and_merge():
    """重新抽取失败项目并直接合并到现有文件"""

    # 加载配置
    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')

    # 找到失败的项目
    failed_data = df[df.iloc[:, 0].isin([118, 229])]

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

    # 初始化组件
    extractor = DeepEntityExtractor(str(config_file))
    ontology_file = Path(__file__).parent / 'config' / 'entity_ontology.yaml'
    normalizer = EntityNormalizer(str(ontology_file))
    builder = RelationshipBuilder(str(ontology_file))

    # 读取现有的实体注册表（如果存在）
    existing_nodes_file = Path(__file__).parent / 'output' / 'nodes_llm.csv'
    existing_rels_file = Path(__file__).parent / 'output' / 'rels_llm.csv'

    existing_nodes_df = pd.read_csv(existing_nodes_file, encoding='utf-8-sig')
    existing_rels_df = pd.read_csv(existing_rels_file, encoding='utf-8-sig')

    print(f"现有节点数: {len(existing_nodes_df)}")
    print(f"现有关系数: {len(existing_rels_df)}")

    # 将现有实体加载到规范化器中
    print("\n=== 加载现有实体到规范化器 ===")
    for idx, row in existing_nodes_df.iterrows():
        if row['type'] != 'ICH_Project':
            # 将现有实体添加到规范化器的注册表
            entity_type = row['type']
            entity_text = row['label']
            entity_id = row['id']

            # 规范化文本
            normalized_text = normalizer.normalize_text(entity_text)

            # 添加到映射表
            if normalized_text not in normalizer.text_to_id_map:
                normalizer.text_to_id_map[normalized_text] = entity_id
                normalizer.entity_registry[entity_id] = {
                    'type': entity_type,
                    'text': entity_text,
                    'normalized_text': normalized_text,
                    'canonical_name': normalized_text  # 添加这个字段
                }

    print(f"已加载 {len(normalizer.entity_registry)} 个现有实体")

    # 重新抽取失败的项目
    extraction_results = []

    for idx, row in failed_data.iterrows():
        project_num = int(row.iloc[0])
        project_id = f'ICH-{project_num}'
        project_name = row.iloc[3]
        description = row.iloc[7] if len(row) > 7 else ""

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

        try:
            result = await extractor.extract_from_remark(
                project_id=project_id,
                project_name=project_name,
                remark_text=description
            )

            if result:
                extraction_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)}")

    if not extraction_results:
        print("\n没有成功抽取的项目")
        return

    # 规范化新抽取的实体（会自动去重）
    print("\n=== 规范化新抽取的实体 ===")
    for item in extraction_results:
        result = item['extraction_result']
        entities = result.get('entities', [])

        for entity in entities:
            # 实体名称可能在name字段或attributes.name字段
            entity_name = entity.get('name') or entity.get('attributes', {}).get('name', '')
            entity_type = entity.get('type', '')

            # 规范化实体（会自动去重）
            normalizer.normalize_entity(entity, similarity_threshold=0.85)

    # 获取所有实体节点（包括新增的）
    all_entity_nodes = normalizer.get_entity_nodes()
    print(f"规范化后实体节点数: {len(all_entity_nodes)}")

    # 构建关系
    print("\n=== 构建关系 ===")
    all_new_relationships = []

    for item in extraction_results:
        project_id = item['project_id']
        result = item['extraction_result']

        entities = result.get('entities', [])
        relationships = result.get('relationships', [])

        # 构建实体名称到ID的映射
        entity_name_to_id = {}
        for entity in entities:
            entity_name = entity.get('name') or entity.get('attributes', {}).get('name', '')
            entity_type = entity.get('type', '')
            normalized_name = normalizer.normalize_text(entity_name)
            entity_id = normalizer.text_to_id_map.get(normalized_name)
            if entity_id:
                entity_name_to_id[entity_name] = entity_id

        # 手动构建关系
        for rel in relationships:
            source_name = rel.get('source_entity') or rel.get('source')
            target_name = rel.get('target_entity') or rel.get('target')
            rel_type = rel.get('type')
            rel_props = rel.get('properties', {})

            # 查找源实体ID
            if source_name == project_id:
                source_id = project_id
            else:
                normalized_name = normalizer.normalize_text(source_name)
                source_id = normalizer.text_to_id_map.get(normalized_name)

            # 查找目标实体ID
            normalized_name = normalizer.normalize_text(target_name)
            target_id = normalizer.text_to_id_map.get(normalized_name)

            # 如果都找到了，添加关系
            if source_id and target_id:
                all_new_relationships.append({
                    'source': source_id,
                    'target': target_id,
                    'type': rel_type,
                    'properties': json.dumps(rel_props, ensure_ascii=False) if rel_props else '{}'
                })

        print(f"项目 {project_id}: {len([r for r in all_new_relationships if r['source'] == project_id])} 条关系")

    # 合并节点和关系
    print("\n=== 合并数据 ===")

    # 项目节点：只添加缺失的项目
    existing_project_ids = set(existing_nodes_df[existing_nodes_df['type'] == 'ICH_Project']['id'].tolist())
    new_project_nodes = [
        {'id': item['project_id'], 'label': item['project_name'], 'type': 'ICH_Project', 'properties': '{}'}
        for item in extraction_results
        if item['project_id'] not in existing_project_ids
    ]

    # 合并所有节点
    all_nodes_df = pd.concat([
        existing_nodes_df[existing_nodes_df['type'] != 'ICH_Project'],  # 现有实体节点
        pd.DataFrame(all_entity_nodes),  # 所有实体节点（包括新增和去重后的现有）
        pd.DataFrame(new_project_nodes),  # 新增项目节点
        existing_nodes_df[existing_nodes_df['type'] == 'ICH_Project']  # 现有项目节点
    ], ignore_index=True)

    # 去重节点（按ID）
    all_nodes_df = all_nodes_df.drop_duplicates(subset=['id'], keep='first')

    # 合并关系
    all_rels_df = pd.concat([
        existing_rels_df,
        pd.DataFrame(all_new_relationships)
    ], ignore_index=True)

    # 去重关系
    all_rels_df = all_rels_df.drop_duplicates(subset=['source', 'target', 'type'], keep='first')

    print(f"合并后节点数: {len(all_nodes_df)} (新增 {len(all_entity_nodes) - len(existing_nodes_df[existing_nodes_df['type'] != 'ICH_Project'])} 个实体)")
    print(f"合并后关系数: {len(all_rels_df)} (新增 {len(all_new_relationships)} 条)")

    # 保存结果
    print("\n=== 保存结果 ===")
    all_nodes_df.to_csv(existing_nodes_file, index=False, encoding='utf-8-sig')
    all_rels_df.to_csv(existing_rels_file, index=False, encoding='utf-8-sig')

    print(f"节点已保存: {existing_nodes_file}")
    print(f"关系已保存: {existing_rels_file}")

    # 验证
    print("\n=== 验证 ===")
    for item in extraction_results:
        project_id = item['project_id']
        rel_count = len(all_rels_df[all_rels_df['source'] == project_id])
        print(f"{project_id}: {rel_count} 条关系")

    print("\n完成！")

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


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