# -*- coding: utf-8 -*-
"""
深度文化实体抽取主流程
从Excel备注字段抽取深层文化实体并生成知识图谱CSV文件
"""

import asyncio
import pandas as pd
import json
import logging
from pathlib import Path
from datetime import datetime
from typing import Dict, List, Any

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

from knowledge_extraction.deep_entity_extractor import DeepEntityExtractor
from data_processing.entity_normalizer import EntityNormalizer
from data_processing.relationship_builder import RelationshipBuilder


class DeepExtractionPipeline:
    """深度文化实体抽取流程"""

    def __init__(self, config_file: str):
        """
        初始化流程

        Args:
            config_file: 配置文件路径
        """
        self.config_file = config_file

        # 设置日志
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
        )
        self.logger = logging.getLogger(__name__)

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

        self.logger.info("DeepExtractionPipeline初始化完成")

    def load_excel_data(self, excel_file: str, max_rows: int = None) -> List[Dict[str, str]]:
        """
        加载Excel数据

        Args:
            excel_file: Excel文件路径
            max_rows: 最大读取行数（用于测试）

        Returns:
            项目数据列表
        """
        self.logger.info(f"读取Excel文件: {excel_file}")

        df = pd.read_excel(excel_file)

        # 限制行数
        if max_rows:
            df = df.head(max_rows)
            self.logger.info(f"限制读取行数: {max_rows}")

        # 检查必需列
        required_columns = ['总序号', '项目名称', '备注']
        missing_columns = [col for col in required_columns if col not in df.columns]

        if missing_columns:
            self.logger.error(f"Excel缺少必需列: {missing_columns}")
            raise ValueError(f"缺少必需列: {missing_columns}")

        # 构建项目数据列表
        projects_data = []
        for idx, row in df.iterrows():
            seq_num = row.get('总序号', idx + 1)
            project_name = row.get('项目名称', '')
            remark_text = row.get('备注', '')

            # 过滤空备注
            if not remark_text or len(str(remark_text).strip()) < 10:
                continue

            projects_data.append({
                'project_id': f"ICH-{int(seq_num)}",
                'project_name': str(project_name),
                'remark_text': str(remark_text),
                'row_number': idx + 2  # Excel行号（1-based + header）
            })

        self.logger.info(f"加载了{len(projects_data)}个项目的数据")

        return projects_data

    async def run_extraction(
        self,
        excel_file: str,
        output_prefix: str = "",
        max_rows: int = None
    ) -> Dict[str, Any]:
        """
        执行完整的抽取流程（支持增量保存）

        Args:
            excel_file: Excel文件路径
            output_prefix: 输出文件前缀（用于测试）
            max_rows: 最大读取行数（用于测试）

        Returns:
            处理结果统计
        """
        self.logger.info("="*60)
        self.logger.info("开始深度文化实体抽取流程（增量保存模式）")
        self.logger.info("="*60)

        start_time = datetime.now()

        # 1. 加载数据
        projects_data = self.load_excel_data(excel_file, max_rows)

        if not projects_data:
            self.logger.error("没有可处理的数据")
            return {}

        # 准备输出目录
        output_dir = Path('output')
        output_dir.mkdir(parents=True, exist_ok=True)

        # 确定输出文件名
        if output_prefix:
            nodes_file = output_dir / f"{output_prefix}_nodes.csv"
            rels_file = output_dir / f"{output_prefix}_rels.csv"
            report_file = output_dir / f"{output_prefix}_report.md"
        else:
            nodes_file = output_dir / Path(self.extractor.config['output']['nodes_file']).name
            rels_file = output_dir / Path(self.extractor.config['output']['relationships_file']).name
            report_file = output_dir / Path(self.extractor.config['output']['report_file']).name

        # 2. 定义增量保存回调函数
        async def save_progress(batch_num, total_batches, extraction_results, processed_projects):
            """每批次完成后保存进度（仅保存节点，不构建关系）"""
            try:
                self.logger.info(f"\n[增量保存] 批次 {batch_num}/{total_batches} 开始保存...")

                # 实体规范化（累积到实体注册表）
                self.normalizer.normalize_batch(extraction_results)

                # 生成节点（累积所有已处理的实体）
                entity_nodes = self.normalizer.get_entity_nodes()

                # 添加项目节点（仅当前批次）
                project_nodes = [
                    {
                        'id': item['project_id'],
                        'label': item['project_name'],
                        'type': 'ICH_Project',
                        'properties': '{}'
                    }
                    for item in processed_projects
                ]

                all_nodes = entity_nodes + project_nodes

                # 保存节点
                nodes_df = pd.DataFrame(all_nodes)
                nodes_df.to_csv(nodes_file, index=False, encoding='utf-8-sig')
                self.logger.info(f"[增量保存] 节点已保存: {nodes_file} ({len(all_nodes)}行)")

            except Exception as e:
                self.logger.error(f"[增量保存] 批次 {batch_num} 保存失败: {str(e)}", exc_info=True)

        # 3. LLM批量抽取（带增量保存回调）
        self.logger.info("\n步骤1: LLM批量抽取（增量保存模式）")
        extraction_results = await self.extractor.batch_extract(
            projects_data,
            progress_callback=save_progress
        )

        # 4. 最终数据验证和报告
        self.logger.info("\n步骤2: 最终数据验证")

        # 重新规范化所有实体（确保一致性）
        entity_id_map = self.normalizer.normalize_batch(extraction_results)

        # 重新构建所有关系
        relationships = self.builder.build_batch_relationships(
            projects_data,
            extraction_results,
            entity_id_map
        )

        # 生成最终节点
        entity_nodes = self.normalizer.get_entity_nodes()
        project_nodes = [
            {
                'id': item['project_id'],
                'label': item['project_name'],
                'type': 'ICH_Project',
                'properties': '{}'
            }
            for item in projects_data
        ]
        all_nodes = entity_nodes + project_nodes

        # 数据验证
        node_ids = set(node['id'] for node in all_nodes)
        validation_report = self.builder.validate_relationships(relationships, node_ids)

        # 保存最终结果
        self.logger.info("\n步骤3: 保存最终结果")
        nodes_df = pd.DataFrame(all_nodes)
        nodes_df.to_csv(nodes_file, index=False, encoding='utf-8-sig')
        self.logger.info(f"最终节点已保存: {nodes_file} ({len(all_nodes)}行)")

        rels_df = pd.DataFrame(relationships)
        rels_df.to_csv(rels_file, index=False, encoding='utf-8-sig')
        self.logger.info(f"最终关系已保存: {rels_file} ({len(relationships)}行)")

        # 生成报告
        self.logger.info("\n步骤4: 生成报告")
        self.generate_report(
            report_file,
            projects_data,
            all_nodes,
            relationships,
            validation_report,
            start_time
        )

        end_time = datetime.now()
        duration = (end_time - start_time).total_seconds()

        # 返回统计信息
        result_stats = {
            'total_projects': len(projects_data),
            'total_nodes': len(all_nodes),
            'total_relationships': len(relationships),
            'entity_nodes': len(entity_nodes),
            'project_nodes': len(project_nodes),
            'duration_seconds': duration,
            'validation_report': validation_report
        }

        self.logger.info("\n"+"="*60)
        self.logger.info(f"抽取完成！耗时: {duration:.2f}秒")
        self.logger.info(f"项目数: {result_stats['total_projects']}")
        self.logger.info(f"节点数: {result_stats['total_nodes']}")
        self.logger.info(f"关系数: {result_stats['total_relationships']}")
        self.logger.info("="*60)

        return result_stats

    def generate_report(
        self,
        report_file: Path,
        projects_data: List[Dict[str, str]],
        nodes: List[Dict[str, Any]],
        relationships: List[Dict[str, Any]],
        validation_report: Dict[str, Any],
        start_time: datetime
    ):
        """
       生成抽取报告

        Args:
            report_file: 报告文件路径
            projects_data: 项目数据
            nodes: 节点列表
            relationships: 关系列表
            validation_report: 验证报告
            start_time: 开始时间
        """
        report_lines = []

        # 标题
        report_lines.append("# 深度文化实体抽取报告\n")
        report_lines.append(f"**生成时间**: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
        report_lines.append(f"**耗时**: {(datetime.now() - start_time).total_seconds():.2f}秒\n")
        report_lines.append("---\n\n")

        # 1. 处理统计
        report_lines.append("## 1. 处理统计\n\n")
        report_lines.append(f"- **处理项目数**: {len(projects_data)}\n")
        report_lines.append(f"- **总节点数**: {len(nodes)}\n")
        report_lines.append(f"- **总关系数**: {len(relationships)}\n")
        report_lines.append(f"- **实体节点数**: {len([n for n in nodes if n['type'] != 'ICH_Project'])}\n")
        report_lines.append(f"- **项目节点数**: {len([n for n in nodes if n['type'] == 'ICH_Project'])}\n")

        # 2. 节点类型分布
        report_lines.append("\n## 2. 节点类型分布\n\n")
        node_types = {}
        for node in nodes:
            node_type = node['type']
            node_types[node_type] = node_types.get(node_type, 0) + 1

        report_lines.append("| 节点类型 | 数量 | 占比 |\n")
        report_lines.append("|---------|------|------|\n")
        for node_type, count in sorted(node_types.items()):
            percentage = (count / len(nodes) * 100) if len(nodes) > 0 else 0
            report_lines.append(f"| {node_type} | {count} | {percentage:.1f}% |\n")

        # 3. 关系类型分布
        report_lines.append("\n## 3. 关系类型分布\n\n")
        rel_types = {}
        for rel in relationships:
            rel_type = rel['type']
            rel_types[rel_type] = rel_types.get(rel_type, 0) + 1

        report_lines.append("| 关系类型 | 数量 | 占比 |\n")
        report_lines.append("|---------|------|------|\n")
        for rel_type, count in sorted(rel_types.items()):
            percentage = (count / len(relationships) * 100) if len(relationships) > 0 else 0
            report_lines.append(f"| {rel_type} | {count} | {percentage:.1f}% |\n")

        # 4. 数据质量
        report_lines.append("\n## 4. 数据质量\n\n")
        report_lines.append(f"- **有效关系**: {validation_report['valid_relationships']}\n")
        report_lines.append(f"- **无效关系**: {validation_report['invalid_relationships']}\n")

        if validation_report['broken_links']:
            report_lines.append(f"\n**断链警告**: {len(validation_report['broken_links'])}个\n")
            report_lines.append("```json\n")
            report_lines.append(json.dumps(validation_report['broken_links'][:10], ensure_ascii=False, indent=2))
            if len(validation_report['broken_links']) > 10:
                report_lines.append(f"\n... (还有{len(validation_report['broken_links'])-10}个)")
            report_lines.append("\n```\n")

        # 5. 项目详情（前5个）
        report_lines.append("\n## 5. 项目抽取详情（前5个）\n\n")

        for idx, (project_data, node) in enumerate(zip(projects_data[:5], nodes[:5])):
            if node['type'] != 'ICH_Project':
                continue

            report_lines.append(f"### {idx+1}. {project_data['project_name']}\n\n")
            report_lines.append(f"**项目ID**: {project_data['project_id']}\n\n")
            report_lines.append(f"**备注**: {project_data['remark_text'][:200]}...\n\n")

            # 查找相关关系
            related_rels = [r for r in relationships if r['source'] == project_data['project_id']]
            if related_rels:
                report_lines.append(f"**关系数**: {len(related_rels)}\n\n")
                report_lines.append("| 关系类型 | 目标实体 |\n")
                report_lines.append("|---------|---------|\n")
                for rel in related_rels[:10]:
                    target_node = next((n for n in nodes if n['id'] == rel['target']), None)
                    if target_node:
                        report_lines.append(f"| {rel['type']} | {target_node['label']} |\n")
                if len(related_rels) > 10:
                    report_lines.append(f"| ... | 还有{len(related_rels)-10}个关系 |\n")

            report_lines.append("\n")

        # 6. 输出文件
        report_lines.append("## 6. 输出文件\n\n")
        report_lines.append(f"- **节点文件**: `{report_file.parent / (report_file.stem.replace('_report', '') + '_nodes.csv')}`\n")
        report_lines.append(f"- **关系文件**: `{report_file.parent / (report_file.stem.replace('_report', '') + '_rels.csv')}`\n")

        # 保存报告
        with open(report_file, 'w', encoding='utf-8') as f:
            f.writelines(report_lines)

        self.logger.info(f"报告已保存: {report_file}")


async def main():
    """主函数"""
    import sys

    # 配置文件
    config_file = r"E:\Project\2026_KG_ICH\dofile\kg_project\config\deep_extraction_config.yaml"

    # 数据文件：从命令行参数获取，或使用默认完整数据文件
    if len(sys.argv) > 1:
        excel_file = sys.argv[1]
    else:
        excel_file = r"E:\Project\2026_KG_ICH\data\黑龙江国家级和省级非遗名单.xlsx"

    # 创建流程
    pipeline = DeepExtractionPipeline(config_file)

    # 执行抽取
    results = await pipeline.run_extraction(
        excel_file=excel_file,
        output_prefix="",  # 不使用前缀（完整抽取）
        max_rows=None  # 读取所有行
    )

    print("\n" + "="*60)
    print("抽取完成！")
    print("="*60)
    print(f"\n结果统计:")
    print(json.dumps(results, indent=2, ensure_ascii=False))


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