"""
提示词加载器
从YAML配置文件中加载系统提示词和用户提示词模板
基于新版LangChain ChatPromptTemplate实现
"""

import yaml
import os
from typing import Dict, Any, List
from foundation.observability.logger.loggering import review_logger as logger
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.messages import SystemMessage, HumanMessage

class PromptLoader:
    """基于LangChain的提示词管理器（单例模式）"""

    _instance = None
    _initialized = False

    def __new__(cls, prompt_config_dir: str = None, auto_preload: bool = True):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance

    def __init__(self, prompt_config_dir: str = None, auto_preload: bool = True):
        """
        初始化提示词加载器

        Args:
            prompt_config_dir: 提示词配置文件目录路径
            auto_preload: 是否自动预加载所有提示词
        """
        # 避免重复初始化
        if self._initialized:
            return

        if prompt_config_dir is None:
            # 默认路径 - 指向prompt目录
            current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
            prompt_config_dir = os.path.join(current_dir, 'prompt')

        self.prompt_config_dir = prompt_config_dir
        self._cache = {}  # 统一缓存

        # 自动预加载所有提示词
        if auto_preload:
            logger.info("开始预加载所有提示词...")
            stats = self._preload_all_prompts()
            logger.info(f"预加载完成: {stats}")
            logger.info(f"总缓存数量: {len(self._cache)}")

        # 标记已初始化
        self._initialized = True

    @classmethod
    def reset_instance(cls):
        """重置单例实例，用于重新初始化"""
        cls._instance = None
        cls._initialized = False

    def list_available_prompts(self, reviewer_type: str) -> List[str]:
        """
        列出指定审查器类型的所有可用提示词名称

        Args:
            reviewer_type: 审查器类型 (basic, technical, rag, ai)

        Returns:
            List[str]: 可用的提示词名称列表
        """
        try:
            # 构建配置文件路径
            if reviewer_type == 'ai':
                config_file = os.path.join(self.prompt_config_dir, "ai_suggestion.yaml")
            elif reviewer_type == 'outline':
                config_file = os.path.join(self.prompt_config_dir, "outline_reviewers.yaml")
            elif reviewer_type == 'query_extract':
                config_file = os.path.join(self.prompt_config_dir, "query_extract.yaml")
            else:
                config_file = os.path.join(self.prompt_config_dir, f"{reviewer_type}_reviewers.yaml")

            # 加载YAML配置
            with open(config_file, 'r', encoding='utf-8') as f:
                config = yaml.safe_load(f)

            return list(config.keys()) if config else []

        except Exception as e:
            logger.error(f"列出可用提示词失败: {reviewer_type}, 错误: {str(e)}")
            return []

    def _load_prompt(self, reviewer_type: str, prompt_name: str) -> Dict[str, str]:
        """
        加载指定审查器的提示词

        Args:
            reviewer_type: 审查器类型 (basic, technical, rag, ai)
            prompt_name: 提示词名称

        Returns:
            Dict[str, str]: 包含 system_prompt 和 user_prompt_template 的字典
        """
        cache_key = f"{reviewer_type}_{prompt_name}"

        # 检查缓存
        if cache_key in self._cache:
            return self._cache[cache_key]

        try:
            # 构建配置文件路径
            if reviewer_type == 'ai':
                config_file = os.path.join(self.prompt_config_dir, "ai_suggestion.yaml")
            elif reviewer_type == 'outline':
                config_file = os.path.join(self.prompt_config_dir, "outline_reviewers.yaml")
            elif reviewer_type == 'query_extract':
                config_file = os.path.join(self.prompt_config_dir, "query_extract.yaml")
            else:
                config_file = os.path.join(self.prompt_config_dir, f"{reviewer_type}_reviewers.yaml")

            # 加载YAML配置
            with open(config_file, 'r', encoding='utf-8') as f:
                config = yaml.safe_load(f)

            # 获取指定提示词
            if prompt_name not in config:
                raise KeyError(f"提示词 '{prompt_name}' 在配置文件中不存在")

            prompt_config = config[prompt_name]

            # 处理 query_extract.yaml 的特殊格式
            if reviewer_type == 'query_extract':
                if 'system' in prompt_config and 'user_template' in prompt_config:
                    # 转换双花括号为单花括号，以便LangChain能够识别变量
                    user_template = prompt_config['user_template'].replace('{{review_content}}', '{review_content}')
                    result = {
                        'system_prompt': prompt_config['system'],
                        'user_prompt_template': user_template
                    }
                else:
                    raise ValueError(f"query_extract 配置缺少必要字段: {prompt_name}")
            else:
                # 验证必要的字段
                if 'system_prompt' not in prompt_config or 'user_prompt_template' not in prompt_config:
                    raise ValueError(f"提示词配置缺少必要字段: {prompt_name}")

                result = {
                    'system_prompt': prompt_config['system_prompt'],
                    'user_prompt_template': prompt_config['user_prompt_template']
                }

            # 缓存结果
            self._cache[cache_key] = result
            return result

        except Exception as e:
            logger.error(f"加载提示词配置失败: {reviewer_type}/{prompt_name}, 错误: {str(e)}")
            # 返回默认提示词
            return {
                'system_prompt': f"你是专业的施工方案审查专家，负责进行{prompt_name}审查。",
                'user_prompt_template': "请审查：{review_content}"
            }

    def get_prompt_template(self, reviewer_type: str, prompt_name: str =None, **kwargs) -> ChatPromptTemplate:
        """
        获取ChatPromptTemplate实例（从缓存获取，避免重复I/O）

        Args:
            reviewer_type: 审查器类型 (basic, technical, rag, ai, query_extract)
            prompt_name: 提示词名称，对于特殊类型可以省略
            **kwargs: 模板变量，如 content, review_references, review_content 等

        Returns:
            ChatPromptTemplate: LangChain ChatPromptTemplate实例
        """
        # 特殊处理：对于query_extract类型，只有一个键
        if reviewer_type == "query_extract":
            cache_key = "query_extract_query_extract"
        elif prompt_name is None:
            cache_key = f"{reviewer_type}"
        else:
            cache_key = f"{reviewer_type}_{prompt_name}"

        try:
            prompt_config = self._cache[cache_key]

            # 创建ChatPromptTemplate
            template = ChatPromptTemplate.from_messages([
                ("system", prompt_config['system_prompt']),
                ("user", prompt_config['user_prompt_template']+" /no_think")
            ])

            # 如果有变量，使用partial方法预先填充
            if kwargs:
                try:
                    template = template.partial(**kwargs)
                    logger.info(f"模板变量填充成功: {list(kwargs.keys())}")
                except Exception as e:
                    # 如果partial失败，记录错误并返回原始模板
                    logger.error(f"模板变量填充失败: {kwargs}, 错误: {str(e)}")
                    pass

            return template

        except Exception as e:
            logger.error(f"创建ChatPromptTemplate失败: {reviewer_type}/{prompt_name}, 错误: {str(e)}")
            # 返回默认模板
            display_name = prompt_name if prompt_name else reviewer_type
            return ChatPromptTemplate.from_messages([
                ("system", f"你是专业的施工方案审查专家，负责进行{display_name}审查。"),
                ("user", "请审查：{review_content}")
            ])
  
    def manage_cache(self, reviewer_type: str = None, prompt_name: str = None, action: str = "clear"):
        """
        缓存管理：清除或重新加载

        Args:
            reviewer_type: 审查器类型
            prompt_name: 提示词名称
            action: "clear" 或 "reload"
        """
        if action == "clear":
            if reviewer_type is None:
                self._cache.clear()
                logger.info("已清除所有提示词缓存")
            elif prompt_name is None:
                keys_to_remove = [k for k in self._cache.keys() if k.startswith(f"{reviewer_type}_")]
                for key in keys_to_remove:
                    del self._cache[key]
                logger.info(f"已清除{reviewer_type}类型的所有提示词缓存")
            else:
                cache_key = f"{reviewer_type}_{prompt_name}"
                if cache_key in self._cache:
                    del self._cache[cache_key]
                logger.info(f"已清除提示词缓存: {reviewer_type}/{prompt_name}")

        elif action == "reload" and prompt_name:
            cache_key = f"{reviewer_type}_{prompt_name}"
            if cache_key in self._cache:
                del self._cache[cache_key]
            self._load_prompt(reviewer_type, prompt_name)
            logger.info(f"已重新加载提示词: {reviewer_type}/{prompt_name}")

    def _preload_all_prompts(self) -> Dict[str, Any]:
        """
        预加载所有提示词到缓存中（内部方法）

        Args:
            reviewer_types: 要预加载的审查器类型列表，None表示加载所有类型

        Returns:
            Dict[str, Any]: 加载结果统计
        """

        reviewer_types = ['basic', 'technical', 'rag', 'ai', 'outline', 'query_extract']

        stats = {
            'loaded_types': [],
            'total_loaded': 0,
            'failed_items': [],
            'by_type': {},
            'cache_size_mb': 0
        }

        for reviewer_type in reviewer_types:
            prompts = self.list_available_prompts(reviewer_type)
            if not prompts:
                logger.warning(f"未找到{reviewer_type}类型的提示词")
                continue

            type_stats = {
                'total': len(prompts),
                'loaded': 0,
                'failed': []
            }

            for prompt_name in prompts:
                try:
                    self._load_prompt(reviewer_type, prompt_name)
                    type_stats['loaded'] += 1
                    stats['total_loaded'] += 1
                except Exception as e:
                    type_stats['failed'].append(prompt_name)
                    stats['failed_items'].append(f"{reviewer_type}/{prompt_name}: {str(e)}")

            stats['by_type'][reviewer_type] = type_stats
            stats['loaded_types'].append(reviewer_type)

        # 计算缓存大小
        stats['cache_size_mb'] = round(
            sum(len(str(v)) for v in self._cache.values()) / (1024 * 1024), 2
        )

        return stats




# 全局提示词加载器实例（自动预加载）
prompt_loader = PromptLoader()