"""
黑龙江省非物质文化遗产知识图谱构建 - 主处理脚本
"""

import json
import asyncio
import yaml
from pathlib import Path
from typing import Dict, List, Any
import logging
from datetime import datetime
import sys

# 添加项目路径
sys.path.insert(0, str(Path(__file__).parent))

from knowledge_extraction.llm_extractor import DeepSeekExtractor


class ICKnowledgeGraphBuilder:
    """非遗知识图谱构建器"""

    def __init__(self, config_file: str = 'config/ich_config.yaml'):
        """初始化构建器"""
        # 加载配置
        with open(config_file, 'r', encoding='utf-8') as f:
            self.config = yaml.safe_load(f)

        # 设置路径
        self.project_root = Path(__file__).parent.parent
        self.data_dir = self.project_root / self.config['paths']['data_dir']
        self.output_dir = self.project_root / self.config['paths']['output_dir']
        self.logs_dir = self.project_root / self.config['paths']['logs_dir']

        # 创建目录
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.logs_dir.mkdir(parents=True, exist_ok=True)

        # 设置日志
        self.logger = self._setup_logger()

        # 初始化DeepSeek抽取器
        self.extractor = DeepSeekExtractor(config_file)

        # 加载进度
        self.progress_file = self.output_dir / 'progress.json'
        self.progress = self._load_progress()

    def _setup_logger(self):
        """设置日志"""
        log_file = self.logs_dir / f"build_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"
        logging.basicConfig(
            level=getattr(logging, self.config['logging']['level']),
            format=self.config['logging']['format'],
            handlers=[
                logging.FileHandler(log_file, encoding='utf-8'),
                logging.StreamHandler()
            ]
        )
        return logging.getLogger(__name__)

    def _load_progress(self) -> Dict:
        """加载进度"""
        if self.progress_file.exists():
            with open(self.progress_file, 'r', encoding='utf-8') as f:
                return json.load(f)
        return {
            'total': 0,
            'processed': [],
            'successful': [],
            'failed': [],
            'last_update': None
        }

    def _save_progress(self):
        """保存进度"""
        self.progress['last_update'] = datetime.now().isoformat()
        with open(self.progress_file, 'w', encoding='utf-8') as f:
            json.dump(self.progress, f, indent=2, ensure_ascii=False)

    def load_nodes(self) -> List[Dict[str, Any]]:
        """加载节点数据"""
        nodes_file = self.output_dir / 'nodes.json'
        if not nodes_file.exists():
            self.logger.error(f"节点文件不存在: {nodes_file}")
            return []

        with open(nodes_file, 'r', encoding='utf-8') as f:
            nodes = json.load(f)

        self.logger.info(f"加载了 {len(nodes)} 个节点")
        self.progress['total'] = len(nodes)

        return nodes

    async def enrich_single_node(self, node: Dict[str, Any]) -> Dict[str, Any]:
        """丰富单个节点"""
        node_id = node.get('id')
        project_name = node.get('properties', {}).get('name', '')

        self.logger.info(f"开始处理节点: {node_id} - {project_name}")

        try:
            # 调用DeepSeek丰富信息
            enriched_data = await self.extractor.enrich_project_entity(node)

            if enriched_data:
                # 合并到原节点
                node['properties']['enriched_data'] = enriched_data
                node['properties']['enriched_at'] = datetime.now().isoformat()

                self.logger.info(f"成功丰富节点: {node_id}")
                return {
                    'node_id': node_id,
                    'status': 'success',
                    'enriched_data': enriched_data
                }
            else:
                self.logger.warning(f"丰富节点失败（无返回数据）: {node_id}")
                return {
                    'node_id': node_id,
                    'status': 'failed',
                    'error': 'No data returned'
                }

        except Exception as e:
            self.logger.error(f"丰富节点失败: {node_id} - {str(e)}")
            return {
                'node_id': node_id,
                'status': 'failed',
                'error': str(e)
            }

    async def batch_enrich_nodes(
        self,
        nodes: List[Dict[str, Any]],
        max_nodes: int = None,
        concurrent: int = 5
    ):
        """批量丰富节点"""
        # 过滤已处理的节点
        pending_nodes = [
            node for node in nodes
            if node.get('id') not in self.progress['processed']
        ]

        # 限制处理数量（用于测试）
        if max_nodes:
            pending_nodes = pending_nodes[:max_nodes]

        self.logger.info(f"开始批量处理: {len(pending_nodes)} 个节点待处理")

        # 分批处理
        batch_size = concurrent
        for i in range(0, len(pending_nodes), batch_size):
            batch = pending_nodes[i:i + batch_size]
            batch_num = i // batch_size + 1
            total_batches = (len(pending_nodes) + batch_size - 1) // batch_size

            self.logger.info(f"处理批次 {batch_num}/{total_batches}: {len(batch)} 个节点")

            # 并发处理
            tasks = [self.enrich_single_node(node) for node in batch]
            results = await asyncio.gather(*tasks, return_exceptions=True)

            # 更新进度
            for result in results:
                if isinstance(result, Exception):
                    self.logger.error(f"处理异常: {str(result)}")
                    continue

                node_id = result.get('node_id')
                status = result.get('status')

                if status == 'success':
                    self.progress['successful'].append(node_id)
                else:
                    self.progress['failed'].append(node_id)

                self.progress['processed'].append(node_id)

            # 保存进度
            self._save_progress()

            # 显示进度
            self._print_progress()

    def _print_progress(self):
        """打印进度"""
        total = self.progress['total']
        processed = len(self.progress['processed'])
        successful = len(self.progress['successful'])
        failed = len(self.progress['failed'])

        print("\n" + "="*60)
        print("处理进度")
        print("="*60)
        print(f"总节点数: {total}")
        print(f"已处理: {processed} ({processed/total*100:.1f}%)")
        print(f"成功: {successful}")
        print(f"失败: {failed}")
        print(f"成功率: {successful/processed*100:.1f}%" if processed > 0 else "成功率: N/A")
        print(f"最后更新: {self.progress['last_update']}")
        print("="*60 + "\n")

    def save_enriched_nodes(self, nodes: List[Dict[str, Any]]):
        """保存丰富后的节点"""
        output_file = self.output_dir / 'nodes_enriched.json'

        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(nodes, f, indent=2, ensure_ascii=False)

        self.logger.info(f"丰富后的节点已保存到: {output_file}")

        # 统计
        enriched_count = sum(
            1 for node in nodes
            if 'enriched_data' in node.get('properties', {})
        )

        print(f"\n统计信息:")
        print(f"总节点数: {len(nodes)}")
        print(f"已丰富节点数: {enriched_count}")
        print(f"丰富率: {enriched_count/len(nodes)*100:.1f}%")


def main():
    """主函数"""
    import argparse

    parser = argparse.ArgumentParser(description='黑龙江非遗知识图谱构建')
    parser.add_argument('--config', default='config/ich_config.yaml', help='配置文件路径')
    parser.add_argument('--max-nodes', type=int, help='最大处理节点数（用于测试）')
    parser.add_argument('--status', action='store_true', help='查看进度')
    parser.add_argument('--resume', action='store_true', help='断点续传')

    args = parser.parse_args()

    # 创建构建器
    builder = ICKnowledgeGraphBuilder(args.config)

    if args.status:
        # 显示进度
        builder._print_progress()
        return

    # 加载节点数据
    print("加载节点数据...")
    nodes = builder.load_nodes()

    if not nodes:
        print("没有找到节点数据")
        return

    # 批量处理
    print("开始批量处理节点...")
    asyncio.run(builder.batch_enrich_nodes(
        nodes,
        max_nodes=args.max_nodes,
        concurrent=builder.config['processing']['concurrent_requests']
    ))

    # 加载并保存丰富后的节点
    print("保存处理结果...")
    nodes_file = builder.output_dir / 'nodes.json'
    with open(nodes_file, 'r', encoding='utf-8') as f:
        all_nodes = json.load(f)

    builder.save_enriched_nodes(all_nodes)

    print("\n处理完成！")
    builder._print_progress()


if __name__ == '__main__':
    main()
