#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
AI 代理模块
负责与大模型进行交互，生成文本内容
"""

import os
import time
import logging
import traceback
from openai import AsyncOpenAI, APITimeoutError, APIConnectionError, RateLimitError, APIStatusError
import tiktoken
from typing import Optional, Tuple

from core.config import AIModelConfig
from core.exception import AIModelError, RetryableError, NonRetryableError

logger = logging.getLogger(__name__)


class AIAgent:
    """
    AI代理类，负责与AI模型交互生成文本内容
    """

    def __init__(self, config: AIModelConfig):
        """
        初始化 AI 代理

        Args:
            config: AI模型配置
        """
        self.config = config
        self.client = AsyncOpenAI(
            api_key=self.config.api_key,
            base_url=self.config.api_url,
            timeout=self.config.timeout
        )
        # try:
        #     self.tokenizer = tiktoken.encoding_for_model(self.config.model)
        # except KeyError:
        #     logger.warning(f"模型 '{self.config.model}' 没有找到对应的tokenizer，将使用 'cl100k_base'")
        #     self.tokenizer = tiktoken.get_encoding("cl100k_base")

    async def generate_text(
        self, system_prompt: str, user_prompt: str, use_stream: bool = False,
        temperature: Optional[float] = None, top_p: Optional[float] = None, 
        presence_penalty: Optional[float] = None, stage: str = ""
    ) -> Tuple[str, int, int, float]:
        """
        生成文本 (支持流式和非流式)

        Args:
            system_prompt: 系统提示
            user_prompt: 用户提示
            use_stream: 是否流式返回
            temperature: 温度参数，控制随机性
            top_p: Top-p采样参数
            presence_penalty: 存在惩罚参数
            stage: 当前所处阶段，用于日志记录

        Returns:
            一个元组 (generated_text, input_tokens, output_tokens, time_cost)
        """
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ]
        
        # 使用传入的参数或默认配置
        temp = temperature if temperature is not None else self.config.temperature
        tp = top_p if top_p is not None else self.config.top_p
        pp = presence_penalty if presence_penalty is not None else self.config.presence_penalty
        
        # 记录使用的模型参数
        stage_info = f"[{stage}]" if stage else ""
        logger.info(f"{stage_info} 使用模型参数: temperature={temp:.2f}, top_p={tp:.2f}, presence_penalty={pp:.2f}")
        
        input_tokens = self.count_tokens(system_prompt + user_prompt)
        logger.info(f"{stage_info} 开始生成任务... 输入token数: {input_tokens}")

        last_exception = None
        backoff_time = 1.0  # Start with 1 second
        start_time = time.time()

        for attempt in range(self.config.max_retries):
            try:
                response = await self.client.chat.completions.create(
                    model=self.config.model,
                    messages=messages,
                    temperature=temp,
                    top_p=tp,
                    presence_penalty=pp,
                    stream=use_stream
                )
                
                time_cost = time.time() - start_time

                if use_stream:
                    # 流式处理需要异步迭代
                    full_text = await self._process_stream(response)
                    output_tokens = self.count_tokens(full_text)
                    logger.info(f"{stage_info} 任务完成，耗时 {time_cost:.2f} 秒. 输出token数: {output_tokens}")
                    return full_text, input_tokens, output_tokens, time_cost
                else:
                    output_text = response.choices[0].message.content.strip()
                    output_tokens = self.count_tokens(output_text)
                    logger.info(f"{stage_info} 任务完成，耗时 {time_cost:.2f} 秒. 输出token数: {output_tokens}")
                    return output_text, input_tokens, output_tokens, time_cost

            except (APITimeoutError, APIConnectionError) as e:
                last_exception = RetryableError(f"AI模型连接或超时错误: {e}")
                logger.warning(f"{stage_info} 尝试 {attempt + 1}/{self.config.max_retries} 失败: {last_exception}. "
                               f"将在 {backoff_time:.1f} 秒后重试...")
                time.sleep(backoff_time)
                backoff_time *= 2  # Exponential backoff
            except (RateLimitError, APIStatusError) as e:
                last_exception = NonRetryableError(f"AI模型API错误 (不可重试): {e}")
                logger.error(f"{stage_info} 发生不可重试的API错误: {last_exception}")
                break  # Do not retry on these errors
            except Exception as e:
                last_exception = AIModelError(f"调用AI模型时发生未知错误: {e}")
                logger.error(f"{stage_info} 发生未知错误: {last_exception}\n{traceback.format_exc()}")
                break

        raise AIModelError(f"AI模型调用在 {self.config.max_retries} 次重试后失败") from last_exception

    async def _process_stream(self, response):
        """异步处理流式响应"""
        full_response = []
        async for chunk in response:
            content = chunk.choices[0].delta.content
            if content:
                full_response.append(content)
                # 如果需要在这里实现真正的流式处理，可以使用回调函数或其他方式
        
        full_text = "".join(full_response)
        logger.info(f"流式响应接收完成，总长度: {len(full_text)}")
        return full_text

    def count_tokens(self, text: str) -> int:
        """
        计算文本的token数量

        Args:
            text: 输入文本

        Returns:
            token数量
        """
        return len(text) // 1.5
        # return len(self.tokenizer.encode(text))

    @staticmethod
    def read_folder_content(folder_path: str) -> str:
        """
        读取指定文件夹下的所有文件内容并合并

        Args:
            folder_path: 文件夹路径

        Returns:
            合并后的文件内容
        """
        if not os.path.exists(folder_path):
            logger.warning(f"引用的文件夹不存在: {folder_path}")
            return ""
            
        context = ""
        try:
            for file_name in sorted(os.listdir(folder_path)):
                file_path = os.path.join(folder_path, file_name)
                if os.path.isfile(file_path):
                    try:
                        with open(file_path, "r", encoding="utf-8") as f:
                            context += f"--- 文件名: {file_name} ---\n"
                            context += f.read().strip()
                            context += "\n\n"
                    except Exception as read_err:
                        logger.error(f"读取文件失败 {file_path}: {read_err}")   
        except Exception as list_err:
            logger.error(f"列出目录失败 {folder_path}: {list_err}")                      
        return context.strip()
