"""
Cognitive Weaver 的 AI 推理模块
处理与 AI 模型的通信以进行关系提取
"""

import asyncio
import json
from typing import Optional
from openai import OpenAI
from .parser import LinkData

class AIInferenceEngine:
    """处理关系提取的 AI 推理"""
    
    def __init__(self, config):
        """
        使用配置初始化 AI 推理引擎。
        
        参数:
            config: 包含 AI 模型设置的配置对象。
        """
        self.config = config
        self.client = None
        self.initialize_client()
    
    def initialize_client(self):
        """根据配置初始化 AI 客户端"""
        ai_config = self.config.ai_model
        
        try:
            if ai_config.provider.lower() == "deepseek":
                # 初始化 DeepSeek 客户端
                self.client = OpenAI(
                    api_key=ai_config.api_key or "sk-3bde3d12ab464212aec4be3113016b33",
                    base_url=ai_config.base_url or "https://api.deepseek.com"
                )
            elif ai_config.provider.lower() == "openai":
                # 初始化 OpenAI 客户端
                self.client = OpenAI(
                    api_key=ai_config.api_key,
                    base_url=ai_config.base_url or "https://api.openai.com/v1"
                )
            else:
                raise ValueError(f"不支持的 AI 提供商: {ai_config.provider}")
        except Exception as e:
            print(f"警告: 无法初始化 AI 客户端: {e}")
            print("使用模拟模式进行测试。")
            self.client = None
    
    async def infer_relation(self, link_data: LinkData) -> Optional[str]:
        """
        使用 AI 推断两个笔记之间的关系
        返回关系链接（例如 "[[支撑观点]]"）或失败时返回 None
        """
        try:
            # 使用提示词模板准备提示词
            prompt = self._build_prompt(link_data)
            
            # 调用 AI 模型
            response = await self._call_ai_model(prompt)
            
            # 提取并验证关系链接
            relation_link = self._extract_relation_link(response)
            
            if relation_link and self._is_valid_relation(relation_link):
                return relation_link
            else:
                print(f"收到无效的关系链接: {relation_link}")
                return None
                
        except Exception as e:
            print(f"AI 推理错误: {e}")
            # 用于测试目的，返回模拟关系
            print("使用模拟关系进行测试: [[简单提及]]")
            return "[[简单提及]]"
    
    def _build_prompt(self, link_data: LinkData) -> str:
        """使用模板构建 AI 推理的提示词"""
        return f"""
源笔记:《{link_data.source_note}》
目标笔记:《{link_data.target_note}》
上下文:"...{link_data.context_text}..."

请判断关系并生成链接。
"""
    
    async def generate_response(self, prompt: str, system_prompt: str = None) -> str:
        """
        为通用提示词从 AI 模型生成响应。
        
        参数:
            prompt (str): 发送给 AI 模型的用户提示词。
            system_prompt (str, optional): 设置上下文的系统提示词。
                默认为通用助手提示词。
        
        返回:
            str: 从 AI 模型生成的响应。
        """
        if system_prompt is None:
            system_prompt = "你是一位有帮助的AI助手，擅长文本分析和关键词提取。"
        
        # 在事件循环中运行同步API调用
        loop = asyncio.get_event_loop()
        response = await loop.run_in_executor(
            None,
            lambda: self.client.chat.completions.create(
                model=self.config.ai_model.model_name,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": prompt}
                ],
                temperature=self.config.ai_model.temperature,
                max_tokens=200
            )
        )
        
        return response.choices[0].message.content

    async def _call_ai_model(self, prompt: str) -> str:
        """使用准备好的提示词调用 AI 模型"""
        # 如果客户端不可用（模拟模式），返回模拟响应
        if self.client is None:
            print("使用模拟 AI 响应进行测试")
            # 检查是否是关键词相似性验证请求
            if "是否指向同一个心理学概念" in prompt or "相同概念" in prompt:
                # 对于关键词相似性验证，返回肯定回答以便测试
                return "是"
            else:
                # 对于关系推理，返回关系链接
                return "[[简单提及]]"
        
        # 来自提示词文件的系统提示词
        system_prompt = """你是一位专注于知识图谱分析的AI助手，对Obsidian的链接哲学有深刻理解。你的核心任务是：

1. 分析上下文: 阅读提供的、包含一个链接的文本片段（上下文）。
2. 判断关系: 根据上下文，从预定义关系列表中，选择一个最能描述"源笔记"与"目标笔记"之间关系。
3. 生成链接: 将选定的关系名称封装成一个标准的Obsidian wiki链接 `[[关系名称]]`。
4. 严格输出: 你的最终回答必须且只能是一个单一的、无任何多余文本的Obsidian wiki链接。不要包含任何解释、问候、标点或额外的文字。

预定义关系列表：
- 支撑观点
- 反驳观点
- 举例说明
- 定义概念
- 属于分类
- 包含部分
- 引出主题
- 简单提及"""
        
        # 在事件循环中运行同步API调用
        loop = asyncio.get_event_loop()
        response = await loop.run_in_executor(
            None,
            lambda: self.client.chat.completions.create(
                model=self.config.ai_model.model_name,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": prompt}
                ],
                temperature=self.config.ai_model.temperature,
                max_tokens=50
            )
        )
        
        return response.choices[0].message.content
    
    def _extract_relation_link(self, response: str) -> Optional[str]:
        """从 AI 响应中提取关系链接"""
            # 查找 Obsidian wiki 链接模式
        import re
        match = re.search(r'\[\[(.*?)\]\]', response)
        if match:
            return f"[[{match.group(1)}]]"
        return None
    
    def _is_valid_relation(self, relation_link: str) -> bool:
        """检查提取的关系是否有效"""
            # 从链接中提取关系名称
        import re
        match = re.search(r'\[\[(.*?)\]\]', relation_link)
        if not match:
            return False
        
        relation_name = match.group(1)
        
            # 检查是否在预定义关系中
        predefined_relations = self.config.relations.predefined_relations
        return relation_name in predefined_relations
