"""实现层 - 异步提示链执行器"""
import json
import logging
from typing import Dict, Any, List
from ...interfaces.llm_client import LLMClient
from ...interfaces.prompt_loader import PromptLoader
from ...interfaces.chain_executor import ChainExecutor


logger = logging.getLogger(__name__)


class AsyncChainExecutor(ChainExecutor):
    """异步提示链执行器"""

    def __init__(
        self,
        llm_client: LLMClient,
        prompt_loader: PromptLoader,
        default_temperature: float = 0.3,
        default_max_tokens: int = 1024
    ):
        """
        初始化执行器

        Args:
            llm_client: 大模型API客户端
            prompt_loader: 提示词加载器
            default_temperature: 默认温度参数
            default_max_tokens: 默认最大token数
        """
        self.llm_client = llm_client
        self.prompt_loader = prompt_loader
        self.default_temperature = default_temperature
        self.default_max_tokens = default_max_tokens

    async def execute_chain(
        self,
        chain_steps: List[Dict[str, Any]],
        initial_input: Dict[str, Any]
    ) -> Dict[str, Any]:
        """
        执行提示链

        Args:
            chain_steps: 提示链步骤列表
            initial_input: 初始输入变量

        Returns:
            最终结果和中间结果
        """
        # 存储所有步骤的结果
        step_results = []
        context = initial_input.copy()

        for step in chain_steps:
            step_name = step.get("name", "unknown")
            prompt_file = step.get("prompt_file", "")
            output_key = step.get("output_key", "")
            input_from = step.get("input_from", None)

            logger.info(f"执行步骤: {step_name}")

            # 加载提示词
            prompt_data = self.prompt_loader.load_prompt(prompt_file)

            # 准备变量
            variables = context.copy()

            # 如果指定了从上一步获取输入
            if input_from and input_from in context:
                variables["input"] = context[input_from]

            # 渲染用户提示词
            user_prompt = self.prompt_loader.render_template(
                prompt_data["user_template"],
                variables
            )

            # 构建消息列表
            messages = []
            if prompt_data["system"]:
                messages.append({"role": "system", "content": prompt_data["system"]})
            messages.append({"role": "user", "content": user_prompt})
            # # 【推荐位置】在这里添加保存代码
            # with open('prompts.txt', 'w', encoding='utf-8') as f:
            #     contents = [msg["content"] for msg in messages]
            #     f.write('\n'.join(contents))

            # 调用API
            response = await self.llm_client.chat_completion(
                messages,
                temperature=self.default_temperature,
                max_tokens=self.default_max_tokens
            )

            # 解析响应
            content = response.get("content", "")
            parsed_result = self._parse_response(content, prompt_data.get("output_parser", {}))

            # 存储结果
            step_result = {
                "name": step_name,
                "prompt_file": prompt_file,
                "result": parsed_result,
                "raw_content": content,
                "usage": response.get("usage", {})
            }
            step_results.append(step_result)

            # 更新上下文
            if output_key:
                context[output_key] = parsed_result

        # 返回最终结果
        final_result = step_results[-1]["result"] if step_results else {}

        return {
            "final_result": final_result,
            "steps": step_results,
            "context": context
        }

    def _parse_response(
        self,
        content: str,
        output_parser: Dict[str, Any]
    ) -> Any:
        """
        解析API响应

        Args:
            content: API返回的内容
            output_parser: 输出解析器配置

        Returns:
            解析后的结果
        """
        parser_type = output_parser.get("type", "text")

        if parser_type == "json":
            try:
                # 尝试提取JSON（可能包含在代码块中）
                json_content = self._extract_json(content)
                return json.loads(json_content)
            except json.JSONDecodeError as e:
                logger.warning(f"JSON解析失败: {e}, 返回原始内容")
                return content
        else:
            return content

    def _extract_json(self, content: str) -> str:
        """
        从内容中提取JSON

        Args:
            content: 可能包含JSON的文本

        Returns:
            提取的JSON字符串
        """
        # 尝试找到JSON代码块
        import re

        # 匹配 ```json ... ``` 格式
        json_pattern = r'```json\s*(.*?)\s*```'
        match = re.search(json_pattern, content, re.DOTALL)
        if match:
            return match.group(1)

        # 匹配 ``` ... ``` 格式
        code_pattern = r'```\s*(.*?)\s*```'
        match = re.search(code_pattern, content, re.DOTALL)
        if match:
            return match.group(1)

        # 如果没有找到代码块，尝试匹配整个内容
        return content.strip()
