"""
DeepSeek实体识别模块
使用DeepSeek API进行非遗知识抽取
"""

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

try:
    from langchain_deepseek import ChatDeepSeek
    from langchain_core.messages import HumanMessage, SystemMessage
except ImportError:
    print("错误: 请先安装依赖包")
    print("运行: pip install langchain-deepseek langchain-core")
    raise


class DeepSeekExtractor:
    """基于DeepSeek的知识抽取器"""

    def __init__(self, config_file: str = 'config/ich_config.yaml'):
        """
        初始化抽取器

        Args:
            config_file: 配置文件路径
        """
        # 加载配置
        self.config = self._load_config(config_file)
        self.api_key = self._load_api_key()

        # 初始化模型
        self.model = ChatDeepSeek(
            model=self.config['api']['model'],
            api_key=self.api_key,
            temperature=self.config['api']['temperature'],
            max_tokens=self.config['api']['max_tokens']
        )

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

        # 加载本体配置
        self.entity_types = self.config['extraction']['entity_types']
        self.relation_types = self.config['extraction']['relation_types']

    def _setup_logger(self):
        """设置日志"""
        log_dir = Path(self.config['paths']['logs_dir'])
        log_dir.mkdir(parents=True, exist_ok=True)

        log_file = log_dir / f"extraction_{datetime.now().strftime('%Y%m%d')}.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_config(self, config_file: str) -> Dict:
        """加载配置"""
        config_path = Path(config_file)
        if not config_path.exists():
            raise FileNotFoundError(f"配置文件不存在: {config_file}")

        with open(config_path, 'r', encoding='utf-8') as f:
            return yaml.safe_load(f)

    def _load_api_key(self) -> str:
        """加载API密钥"""
        api_config_file = Path(self.config['api']['config_file'])
        if not api_config_file.exists():
            raise FileNotFoundError(f"API配置文件不存在: {api_config_file}")

        with open(api_config_file, 'r', encoding='utf-8') as f:
            api_config = yaml.safe_load(f)

        return api_config['deepseek_api_key']

    def extract_json_from_response(self, response_text: str) -> str:
        """
        从API响应中提取JSON内容

        Args:
            response_text: API响应文本

        Returns:
            str: 提取的JSON字符串
        """
        # 如果响应包含markdown代码块，提取其中的JSON
        if '```json' in response_text:
            start = response_text.find('```json') + 7
            end = response_text.find('```', start)
            if start > 6 and end > start:
                return response_text[start:end].strip()
        elif '```' in response_text:
            start = response_text.find('```') + 3
            end = response_text.find('```', start)
            if start > 2 and end > start:
                content = response_text[start:end].strip()
                if not content.startswith('json'):
                    return content
                return content[5:].strip() if content.startswith('json') else content.strip()

        # 否则直接返回原始文本
        return response_text.strip()

    async def extract_entities_from_text(
        self,
        text: str,
        project_name: str = "",
        max_retries: int = 3
    ) -> Optional[Dict[str, Any]]:
        """
        从文本中抽取实体

        Args:
            text: 输入文本
            project_name: 项目名称（可选）
            max_retries: 最大重试次数

        Returns:
            Dict: 抽取结果
        """
        # 构建提示词
        entity_types_str = "、".join(self.entity_types)

        prompt = f"""
你是一位非物质文化遗产领域的专家。请从以下文本中识别非遗相关实体，并提取属性。

项目名称：{project_name}

文本内容：
{text}

请识别以下类型的实体：
{entity_types_str}

对于每个实体，请提取以下信息：
1. 实体文本
2. 实体类型
3. 相关属性（如：民族、地点、时间、技艺特点等）

输出格式（JSON）：
{{
  "entities": [
    {{
      "text": "实体文本",
      "type": "实体类型",
      "attributes": {{
        "ethnic_group": "民族（如果适用）",
        "location": "地点（如果适用）",
        "time_period": "时期（如果适用）",
        "skill_feature": "技艺特点（如果适用）",
        "cultural_value": "文化价值（如果适用）"
      }}
    }}
  ],
  "relationships": [
    {{
      "from": "实体1",
      "to": "实体2",
      "type": "关系类型",
      "description": "关系描述"
    }}
  ]
}}

请确保输出是有效的JSON格式。
"""

        # 调用API
        for attempt in range(max_retries):
            try:
                self.logger.info(f"开始抽取实体 (尝试 {attempt + 1}/{max_retries})")

                messages = [HumanMessage(content=prompt)]
                response = await self.model.ainvoke(messages)
                result_text = response.content

                # 提取JSON
                json_text = self.extract_json_from_response(result_text)

                # 解析JSON
                try:
                    result = json.loads(json_text)

                    # 添加元数据
                    result['metadata'] = {
                        'project_name': project_name,
                        'extraction_time': datetime.now().isoformat(),
                        'text_length': len(text),
                        'model': self.config['api']['model']
                    }

                    self.logger.info(f"成功抽取 {len(result.get('entities', []))} 个实体")
                    return result

                except json.JSONDecodeError as e:
                    self.logger.warning(f"JSON解析失败: {str(e)}")
                    self.logger.debug(f"响应文本: {result_text[:500]}")
                    if attempt == max_retries - 1:
                        raise

            except Exception as e:
                self.logger.error(f"API调用失败 (尝试 {attempt + 1}/{max_retries}): {str(e)}")
                if attempt == max_retries - 1:
                    return None

                # 等待后重试
                await asyncio.sleep(self.config['processing']['retry_delay'])

        return None

    async def extract_relationships(
        self,
        entities: List[Dict],
        context: str = "",
        max_retries: int = 3
    ) -> Optional[List[Dict[str, Any]]]:
        """
        抽取实体间的关系

        Args:
            entities: 实体列表
            context: 上下文文本
            max_retries: 最大重试次数

        Returns:
            List[Dict]: 关系列表
        """
        if not entities:
            return []

        # 构建实体描述
        entity_descriptions = []
        for entity in entities:
            desc = f"- {entity.get('text', '')} ({entity.get('type', '')})"
            if entity.get('attributes'):
                attrs = ", ".join([f"{k}={v}" for k, v in entity['attributes'].items() if v])
                if attrs:
                    desc += f" [{attrs}]"
            entity_descriptions.append(desc)

        entities_str = "\n".join(entity_descriptions)
        relations_str = "、".join(self.relation_types)

        prompt = f"""
基于以下实体和上下文，识别实体间的关系：

实体列表：
{entities_str}

上下文：
{context}

可能的关系类型：
{relations_str}

对于每个关系，请提供：
1. 头实体（from）
2. 尾实体（to）
3. 关系类型
4. 关系描述

输出格式（JSON）：
{{
  "relationships": [
    {{
      "from": "实体1文本",
      "to": "实体2文本",
      "type": "关系类型",
      "description": "关系描述",
      "confidence": 0.9
    }}
  ]
}}

请确保输出是有效的JSON格式。
"""

        # 调用API
        for attempt in range(max_retries):
            try:
                messages = [HumanMessage(content=prompt)]
                response = await self.model.ainvoke(messages)
                result_text = response.content

                # 提取JSON
                json_text = self.extract_json_from_response(result_text)

                # 解析JSON
                try:
                    result = json.loads(json_text)
                    relationships = result.get('relationships', [])

                    self.logger.info(f"成功抽取 {len(relationships)} 个关系")
                    return relationships

                except json.JSONDecodeError as e:
                    self.logger.warning(f"JSON解析失败: {str(e)}")
                    if attempt == max_retries - 1:
                        raise

            except Exception as e:
                self.logger.error(f"API调用失败 (尝试 {attempt + 1}/{max_retries}): {str(e)}")
                if attempt == max_retries - 1:
                    return None

                await asyncio.sleep(self.config['processing']['retry_delay'])

        return None

    async def enrich_project_entity(
        self,
        project_data: Dict[str, Any],
        max_retries: int = 3
    ) -> Optional[Dict[str, Any]]:
        """
        丰富非遗项目实体信息

        Args:
            project_data: 项目数据
            max_retries: 最大重试次数

        Returns:
            Dict: 丰富后的实体信息
        """
        project_name = project_data.get('properties', {}).get('name', '')
        description = project_data.get('properties', {}).get('description', '')

        prompt = f"""
你是一位非物质文化遗产领域的专家。请分析以下非遗项目，提取和丰富实体信息。

项目名称：{project_name}

项目描述：
{description}

请提取以下信息：

1. **民族特色**：是否与特定少数民族相关？（满族、赫哲族、鄂伦春族等）
2. **地域特征**：体现哪些黑龙江地域特征？（寒地、冰雪、森林、江河等）
3. **技艺特点**：核心技艺特点是什么？
4. **文化价值**：有哪些重要的文化价值？
5. **传承方式**：如何传承？（家族传承、师徒制度、口传心授等）
6. **濒危状况**：是否濒危？原因是什么？
7. **保护措施**：有哪些保护措施？

输出格式（JSON）：
{{
  "ethnic_features": ["民族1", "民族2"],
  "regional_characteristics": ["特征1", "特征2"],
  "skill_features": ["技艺特点1", "技艺特点2"],
  "cultural_values": ["价值1", "价值2"],
  "transmission_methods": ["方式1", "方式2"],
  "endangerment_status": "濒危状况描述",
  "protection_measures": ["措施1", "措施2"],
  "summary": "项目总结（50字以内）"
}}

请确保输出是有效的JSON格式。
"""

        # 调用API
        for attempt in range(max_retries):
            try:
                messages = [HumanMessage(content=prompt)]
                response = await self.model.ainvoke(messages)
                result_text = response.content

                # 提取JSON
                json_text = self.extract_json_from_response(result_text)

                # 解析JSON
                try:
                    result = json.loads(json_text)

                    self.logger.info(f"成功丰富项目实体信息: {project_name}")
                    return result

                except json.JSONDecodeError as e:
                    self.logger.warning(f"JSON解析失败: {str(e)}")
                    if attempt == max_retries - 1:
                        raise

            except Exception as e:
                self.logger.error(f"API调用失败 (尝试 {attempt + 1}/{max_retries}): {str(e)}")
                if attempt == max_retries - 1:
                    return None

                await asyncio.sleep(self.config['processing']['retry_delay'])

        return None

    async def batch_extract(
        self,
        items: List[Dict[str, Any]],
        concurrent: int = 5
    ) -> List[Optional[Dict[str, Any]]]:
        """
        批量抽取

        Args:
            items: 待处理项目列表
            concurrent: 并发数

        Returns:
            List[Dict]: 抽取结果列表
        """
        results = []
        batch_size = concurrent

        for i in range(0, len(items), batch_size):
            batch = items[i:i + batch_size]
            self.logger.info(f"处理批次 {i//batch_size + 1}: {len(batch)} 个项目")

            # 并发处理
            tasks = []
            for item in batch:
                project_name = item.get('properties', {}).get('name', '')
                description = item.get('properties', {}).get('description', '')

                if description:
                    task = self.enrich_project_entity(item)
                    tasks.append(task)
                else:
                    tasks.append(asyncio.sleep(0))  # 占位任务

            batch_results = await asyncio.gather(*tasks, return_exceptions=True)
            results.extend(batch_results)

            # 显示进度
            completed = min(i + batch_size, len(items))
            self.logger.info(f"进度: {completed}/{len(items)}")

        return results


def main():
    """主函数 - 测试实体抽取"""
    import sys

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

    # 创建抽取器
    extractor = DeepSeekExtractor('config/ich_config.yaml')

    # 测试文本
    test_text = """
    桦树皮制作技艺是鄂伦春族的传统手工艺，利用桦树皮制作各种生活用品。
    传承人莫桂树2008年被评为国家级传承人。这项技艺体现了鄂伦春族对自然资源的
    巧妙利用，具有鲜明的渔猎文化特色。
    """

    # 测试实体抽取
    print("="*60)
    print("测试实体抽取")
    print("="*60)

    async def test():
        result = await extractor.extract_entities_from_text(
            text=test_text,
            project_name="桦树皮制作技艺"
        )

        if result:
            print("\n抽取结果:")
            print(json.dumps(result, indent=2, ensure_ascii=False))
        else:
            print("抽取失败")

    asyncio.run(test())


if __name__ == '__main__':
    main()
