# -*- coding: utf-8 -*-
"""
深度文化实体抽取器
利用LLM从非遗项目备注中抽取深层文化实体和关系
"""

import json
import asyncio
import logging
from typing import Dict, List, Any, Optional
from pathlib import Path
import yaml
from langchain_deepseek import ChatDeepSeek
from langchain_core.messages import HumanMessage, SystemMessage


class DeepEntityExtractor:
    """深度文化实体抽取器"""

    def __init__(self, config_file: str):
        """
        初始化抽取器

        Args:
            config_file: 配置文件路径
        """
        # 加载配置
        with open(config_file, 'r', encoding='utf-8') as f:
            self.config = yaml.safe_load(f)

        # 设置日志
        log_config = self.config.get('logging', {})
        log_file = log_config.get('log_file', 'logs/deep_extraction.log')
        Path(log_file).parent.mkdir(parents=True, exist_ok=True)

        # 创建logger
        self.logger = logging.getLogger(__name__)
        self.logger.setLevel(getattr(logging, log_config.get('level', 'INFO')))

        # 清除已有的handlers
        self.logger.handlers.clear()

        # 文件handler（实时刷新）
        file_handler = logging.FileHandler(log_file, encoding='utf-8')
        file_handler.setLevel(getattr(logging, log_config.get('level', 'INFO')))
        file_formatter = logging.Formatter(log_config.get('format', '%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
        file_handler.setFormatter(file_formatter)
        # 强制实时刷新
        file_handler.flush = lambda: file_handler.stream.flush()
        self.logger.addHandler(file_handler)

        # 控制台handler
        console_handler = logging.StreamHandler()
        console_handler.setLevel(getattr(logging, log_config.get('level', 'INFO')))
        console_formatter = logging.Formatter(log_config.get('format', '%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
        console_handler.setFormatter(console_formatter)
        self.logger.addHandler(console_handler)

        # 加载实体本体配置
        ontology_file = Path(config_file).parent / 'entity_ontology.yaml'
        with open(ontology_file, 'r', encoding='utf-8') as f:
            self.ontology = yaml.safe_load(f)

        # 初始化LLM
        self._init_llm()

        # 构建提示词
        self.system_prompt = self._build_system_prompt()

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

    def _init_llm(self):
        """初始化LLM模型"""
        llm_config = self.config['llm']

        # 读取API密钥（从现有的api_keys.yaml）
        api_keys_file = Path(__file__).parent.parent.parent / 'config' / 'api_keys.yaml'
        with open(api_keys_file, 'r', encoding='utf-8') as f:
            api_keys = yaml.safe_load(f)

        # 适配API密钥格式
        api_key = api_keys.get('deepseek_api_key', '')
        if not api_key:
            # 尝试嵌套格式
            api_key = api_keys.get('deepseek', {}).get('api_key', '')

        self.llm = ChatDeepSeek(
            model=llm_config['model'],
            temperature=llm_config['temperature'],
            max_tokens=llm_config['max_tokens'],
            api_key=api_key
        )

        self.logger.info(f"LLM初始化完成: {llm_config['model']}")

    def _build_system_prompt(self) -> str:
        """构建系统提示词"""
        entity_prompts = self.config.get('entity_type_prompts', {})
        relationship_prompts = self.config.get('relationship_type_prompts', {})

        prompt = """你是一位非物质文化遗产领域的专家，擅长从项目描述中识别深层文化实体和关系。

【任务】
请从非遗项目描述中识别以下9类深层文化实体和9种深层关系：

=== 实体类型 ===
"""

        # 添加实体类型说明
        for entity_type, description in entity_prompts.items():
            prompt += f"{entity_type}：{description}\n"

        prompt += "\n=== 关系类型 ===\n"

        # 添加关系类型说明
        for rel_type, description in relationship_prompts.items():
            prompt += f"{rel_type}：{description}\n"

        prompt += """
【输出要求】
1. 严格按JSON格式输出，不要包含任何其他文本
2. 只输出明确的、文本中提到的实体和关系，不要臆测
3. 实体text必须从原文中提取，不要自行改写
4. 同一个实体只识别一次，避免重复
5. 关系必须基于文本中的明确描述
6. 如果某类实体或关系不存在，相应数组为空
7. 确保JSON格式正确，可以被Python解析

【JSON输出格式】
```json
{
  "entities": [
    {
      "text": "实体文本（如：鱼皮）",
      "type": "实体类型（如：Material）",
      "attributes": {
        "name": "标准名称",
        "category": "类别（如：动物材料）",
        "description": "详细描述"
      }
    }
  ],
  "relationships": [
    {
      "source_entity": "ICH-{项目ID}",
      "target_entity": "目标实体文本",
      "type": "关系类型",
      "properties": {
        "description": "关系描述",
        "context": "上下文信息"
      }
    }
  ],
  "summary": {
    "total_entities": 0,
    "total_relationships": 0
  }
}
```
"""

        return prompt

    async def extract_from_remark(
        self,
        project_id: str,
        project_name: str,
        remark_text: str
    ) -> Optional[Dict[str, Any]]:
        """
        从单条备注抽取实体和关系

        Args:
            project_id: 项目ID（如ICH-1）
            project_name: 项目名称
            remark_text: 备注文本

        Returns:
            抽取结果（包含entities, relationships, summary）
        """
        if not remark_text or len(remark_text.strip()) < 10:
            self.logger.warning(f"项目 {project_id} 备注文本过短，跳过抽取")
            return None

        # 构建用户提示词
        user_prompt = f"""【项目信息】
项目名称：{project_name}
项目ID：{project_id}

【描述文本】
{remark_text}

请从上述描述中识别深层文化实体和关系，严格按JSON格式输出。"""

        try:
            # 调用LLM
            messages = [
                SystemMessage(content=self.system_prompt),
                HumanMessage(content=user_prompt)
            ]

            max_retries = self.config['llm']['max_retries']
            retry_delay = self.config['llm']['retry_delay']
            request_timeout = self.config['llm'].get('request_timeout', 120)

            for attempt in range(max_retries):
                try:
                    self.logger.info(f"项目 {project_id} 开始LLM调用（尝试{attempt+1}/{max_retries}）")

                    # 添加超时控制
                    response = await asyncio.wait_for(
                        self.llm.ainvoke(messages),
                        timeout=request_timeout
                    )
                    result_text = response.content

                    self.logger.info(f"项目 {project_id} LLM调用成功，开始提取JSON")

                    # 提取JSON
                    extraction_result = self._extract_json(result_text)

                    if extraction_result:
                        # 验证结果
                        if self._validate_extraction(extraction_result):
                            self.logger.info(f"项目 {project_id} 抽取成功：{extraction_result['summary']['total_entities']}个实体，{extraction_result['summary']['total_relationships']}个关系")
                            return extraction_result
                        else:
                            self.logger.warning(f"项目 {project_id} 抽取结果验证失败")
                    else:
                        self.logger.warning(f"项目 {project_id} JSON提取失败")

                except asyncio.TimeoutError:
                    self.logger.error(f"项目 {project_id} LLM调用超时（{request_timeout}秒）（尝试{attempt+1}/{max_retries}）")
                    if attempt < max_retries - 1:
                        await asyncio.sleep(retry_delay)
                    else:
                        self.logger.error(f"项目 {project_id} 达到最大重试次数，放弃抽取")
                        raise
                except Exception as e:
                    self.logger.error(f"项目 {project_id} LLM调用失败（尝试{attempt+1}/{max_retries}）: {type(e).__name__}: {str(e)}")
                    if attempt < max_retries - 1:
                        await asyncio.sleep(retry_delay)
                    else:
                        raise

        except Exception as e:
            self.logger.error(f"项目 {project_id} 抽取失败: {str(e)}", exc_info=True)
            return None

    def _extract_json(self, text: str) -> Optional[Dict[str, Any]]:
        """
        从文本中提取JSON

        Args:
            text: LLM返回的文本

        Returns:
            解析后的JSON对象，失败返回None
        """
        # 尝试直接解析
        try:
            return json.loads(text)
        except json.JSONDecodeError:
            pass

        # 尝试提取JSON块
        import re
        json_pattern = r'```json\s*(.*?)\s*```'
        match = re.search(json_pattern, text, re.DOTALL)

        if match:
            try:
                return json.loads(match.group(1))
            except json.JSONDecodeError:
                pass

        # 尝试提取花括号内容
        brace_pattern = r'\{.*\}'
        match = re.search(brace_pattern, text, re.DOTALL)

        if match:
            try:
                return json.loads(match.group(0))
            except json.JSONDecodeError:
                pass

        self.logger.error(f"无法从文本中提取有效JSON: {text[:200]}...")
        return None

    def _validate_extraction(self, result: Dict[str, Any]) -> bool:
        """
        验证抽取结果

        Args:
            result: 抽取结果

        Returns:
            是否有效
        """
        # 检查必需字段
        required_fields = ['entities', 'relationships', 'summary']
        for field in required_fields:
            if field not in result:
                self.logger.error(f"缺少必需字段: {field}")
                return False

        # 检查entities格式
        if not isinstance(result['entities'], list):
            self.logger.error("entities必须是列表")
            return False

        for entity in result['entities']:
            if not all(k in entity for k in ['text', 'type', 'attributes']):
                self.logger.error(f"实体缺少必需字段: {entity}")
                return False

            # 检查实体类型是否合法
            entity_type = entity['type']
            if entity_type not in self.ontology['entity_types']:
                self.logger.warning(f"未知实体类型: {entity_type}")

        # 检查relationships格式
        if not isinstance(result['relationships'], list):
            self.logger.error("relationships必须是列表")
            return False

        for rel in result['relationships']:
            if not all(k in rel for k in ['source_entity', 'target_entity', 'type', 'properties']):
                self.logger.error(f"关系缺少必需字段: {rel}")
                return False

            # 检查关系类型是否合法
            rel_type = rel['type']
            if rel_type not in self.ontology['relationship_types']:
                self.logger.warning(f"未知关系类型: {rel_type}")

        return True

    async def batch_extract(
        self,
        projects_data: List[Dict[str, str]],
        batch_size: int = None,
        progress_callback=None
    ) -> List[Optional[Dict[str, Any]]]:
        """
        批量抽取实体和关系

        Args:
            projects_data: 项目数据列表，每项包含project_id, project_name, remark_text
            batch_size: 批次大小（从配置文件读取）
            progress_callback: 进度回调函数，每批次完成后调用

        Returns:
            抽取结果列表
        """
        if batch_size is None:
            batch_size = self.config['batch_processing']['batch_size']

        batch_timeout = self.config['llm'].get('batch_timeout', 300)
        incremental_save = self.config['batch_processing'].get('incremental_save', False)

        results = []
        total = len(projects_data)

        self.logger.info(f"="*60)
        self.logger.info(f"开始批量抽取，共{total}个项目，批次大小{batch_size}")
        self.logger.info(f"增量保存: {'启用' if incremental_save else '禁用'}")
        self.logger.info(f"请求超时: {self.config['llm'].get('request_timeout', 120)}秒")
        self.logger.info(f"批次超时: {batch_timeout}秒")
        self.logger.info(f"="*60)

        for i in range(0, total, batch_size):
            batch = projects_data[i:i+batch_size]
            batch_num = i // batch_size + 1
            total_batches = (total + batch_size - 1) // batch_size

            self.logger.info(f"-"*60)
            self.logger.info(f"开始处理批次 {batch_num}/{total_batches}，包含{len(batch)}个项目")
            for item in batch:
                self.logger.info(f"  - {item['project_id']}: {item['project_name']}")

            # 并发处理当前批次
            batch_tasks = [
                self.extract_from_remark(
                    item['project_id'],
                    item['project_name'],
                    item['remark_text']
                )
                for item in batch
            ]

            try:
                # 添加批次级别的超时控制
                batch_results = await asyncio.wait_for(
                    asyncio.gather(*batch_tasks, return_exceptions=True),
                    timeout=batch_timeout
                )
                results.extend(batch_results)

                # 统计本批次结果
                success_count = sum(1 for r in batch_results if r is not None and not isinstance(r, Exception))
                error_count = len(batch_results) - success_count

                self.logger.info(f"批次 {batch_num}/{total_batches} 完成 - 成功: {success_count}, 失败: {error_count}")

                # 调用进度回调（用于增量保存）
                if progress_callback:
                    await progress_callback(batch_num, total_batches, results, projects_data[:i+len(batch)])

            except asyncio.TimeoutError:
                self.logger.error(f"批次 {batch_num}/{total_batches} 超时（{batch_timeout}秒），部分请求失败")
                # 将未完成的结果标记为None
                for item in batch[len(results):]:
                    results.append(None)

        self.logger.info(f"="*60)
        self.logger.info(f"批量抽取完成，共处理{len(results)}个项目")
        self.logger.info(f"="*60)

        return results
