# 标准库导入
import time
from functools import lru_cache
from typing import Dict, List, Any, Optional

# 第三方库导入
from langchain.chains.base import Chain
from pydantic import Field

# 本地导入
from utils.logger_config import setup_logger
from utils.prompt_config import prompt_router
from utils.rag_config import BaseMethod, ConfigManager

# 配置
logger = setup_logger(__name__)

class IntentChain(Chain):
    """
    意图链，用于处理用户问题并返回答案
    继承自Chain，并组合BaseMethod的功能
    """
    
    # 声明Pydantic字段
    config_manager: ConfigManager = Field(default_factory=ConfigManager)
    base_method: BaseMethod = Field(default_factory=BaseMethod)
    prompt_template: Any = None

    def __init__(self, config_path=None):
        """
        初始化意图链
        
        Args:
            config_path: 配置文件路径，如果为None则使用默认路径
        """
        super().__init__()
        # 使用指定的配置文件初始化
        if config_path:
            self.config_manager = ConfigManager(config_path)
            self.base_method = BaseMethod(config_path)
        self._init_prompt_template()
    
    def _init_prompt_template(self) -> None:
        """初始化提示词模板"""
        try:
            self.prompt_template = prompt_router("intent_prompt")
        except Exception as e:
            logger.error(f"初始化提示词模板失败: {e}")
            raise

    @property
    def input_keys(self) -> List[str]:
        """定义输入键"""
        return ["question"]

    @property
    def output_keys(self) -> List[str]:
        """定义输出键"""
        return ["answer"]

    def _format_documents(self, documents: List[str]) -> str:
        """
        格式化检索到的文档
        
        Args:
            documents: 文档列表
            
        Returns:
            格式化后的文档字符串
        """
        retriever_text = " ".join([doc for doc in documents])
        return retriever_text
    
    def _retrieve_documents(self, question: str) -> List[str]:
        """
        检索相关文档
        
        Args:
            question: 用户问题
            
        Returns:
            相关文档内容列表
            
        Raises:
            Exception: 检索失败时抛出异常
        """
        try:
            retrieved_docs = self.base_method.csv_retriever(question)
            if not retrieved_docs:
                logger.warning(f"未找到相关文档: {question}")
                return []
            return [doc.page_content for doc in retrieved_docs]
        except Exception as e:
            logger.error(f"文档检索失败: {e}")
            raise

    def _generate_answer(self, context: str, question: str, history: str) -> str:
        """
        生成答案
        
        Args:
            context: 上下文信息
            question: 用户问题
            history: 历史对话记录
            
        Returns:
            生成的答案
            
        Raises:
            Exception: 生成答案失败时抛出异常
        """
        try:
            prompt = self.prompt_template.format(
                history=history,
                context=context,
                question=question
            )
            return self.base_method.model_config.llm(prompt).content
        except Exception as e:
            logger.error(f"生成答案失败: {e}")
            raise

    def _call(self, inputs: Dict[str, Any]) -> Dict[str, str]:
        """
        处理用户输入并返回答案
        
        Args:
            inputs: 包含用户问题的字典
            history: 历史对话记录
            
        Returns:
            包含答案的字典
            
        Raises:
            KeyError: 输入缺少必要字段
            Exception: 处理过程中的其他异常
        """
        try:
            # 参数验证
            if "question" not in inputs:
                raise KeyError("输入缺少'question'字段")
            
            question = inputs["question"]
            if not isinstance(question, str) or not question.strip():
                raise ValueError("问题不能为空")
            
            if "history" not in inputs:
                raise KeyError("输入缺少'history'字段")

            history = inputs["history"]
            logger.info(f"intent_chain history message: {history}")

            # 检索文档
            search_start_time = time.time()
            documents = self._retrieve_documents(question)
            if not documents:
                documents = "NaN"
            logger.info(f"意图识别向量库-检索耗时: {time.time() - search_start_time} 秒, 意图识别检索到文档: {documents}")

            # 格式化文档
            context = self._format_documents(documents)
            logger.info(f"intent_chain retriever context: {context}")

            # 生成答案
            generate_start_time = time.time()
            answer = self._generate_answer(context, question, history)
            logger.info(f"意图识别耗时: {time.time() - generate_start_time} 秒, 意图识别结果: {answer}")

            # 返回结果
            return {"answer": answer}
            
        except (KeyError, ValueError) as e:
            logger.error(f"输入参数错误: {e}")
            raise
        except Exception as e:
            logger.error(f"意图识别失败: {str(e)}")
            return {"answer": "意图识别过程出现错误"}

if __name__ == "__main__":
    # 使用示例
    chain = IntentChain()
    history = "你好，请问有什么可以帮助你？"
    try:
        result = chain.invoke({"question": "你好，请问有什么可以帮助你？", "history": history})
        print(f"回答: {result['answer']}")
    except Exception as e:
        logger.error(f"处理失败: {e}")