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

"""
AI Service
AI服务 - 重构后的AI交互服务，移除硬编码依赖
"""

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

from ..config import AIModelConfig
from ..exceptions import (
    AIServiceError, NetworkError, RateLimitError as CustomRateLimitError,
    QuotaExceededError, ModelNotFoundError, TimeoutError as CustomTimeoutError
)

logger = logging.getLogger(__name__)


class AIService:
    """
    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
        )
        
        # 初始化tokenizer (可选)
        self.tokenizer = None
        try:
            if hasattr(tiktoken, 'encoding_for_model'):
                self.tokenizer = tiktoken.encoding_for_model(self.config.model)
        except (KeyError, Exception) as e:
            logger.warning(f"无法为模型 '{self.config.model}' 初始化tokenizer: {e}")
            try:
                self.tokenizer = tiktoken.get_encoding("cl100k_base")
            except Exception:
                logger.warning("使用默认tokenizer失败，token计数功能将不可用")

    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 = "",
        **kwargs
    ) -> Tuple[str, int, int, float]:
        """
        生成文本内容

        Args:
            system_prompt: 系统提示词
            user_prompt: 用户提示词
            use_stream: 是否使用流式输出
            temperature: 温度参数（覆盖配置）
            top_p: top_p参数（覆盖配置）
            presence_penalty: presence_penalty参数（覆盖配置）
            stage: 生成阶段标识（用于日志）
            **kwargs: 其他参数

        Returns:
            Tuple[生成的文本, 输入token数, 输出token数, 耗时秒数]

        Raises:
            AIServiceError: AI服务相关异常
        """
        start_time = time.time()
        
        # 使用传入参数或配置默认值
        params = {
            "temperature": temperature if temperature is not None else self.config.temperature,
            "top_p": top_p if top_p is not None else self.config.top_p,
            "presence_penalty": presence_penalty if presence_penalty is not None else self.config.presence_penalty,
        }
        
        # 过滤None值
        params = {k: v for k, v in params.items() if v is not None}
        
        # 计算输入token数
        input_tokens = self._count_tokens(system_prompt + user_prompt)
        
        try:
            if use_stream and self.config.enable_stream:
                content = await self._generate_with_stream(system_prompt, user_prompt, stage, **params)
            else:
                content = await self._generate_without_stream(system_prompt, user_prompt, stage, **params)
            
            # 计算输出token数和耗时
            output_tokens = self._count_tokens(content)
            elapsed_time = time.time() - start_time
            
            logger.info(f"{stage}生成完成 - 输入tokens: {input_tokens}, 输出tokens: {output_tokens}, 耗时: {elapsed_time:.2f}s")
            
            return content, input_tokens, output_tokens, elapsed_time
            
        except Exception as e:
            elapsed_time = time.time() - start_time
            error_msg = f"{stage}生成失败: {str(e)}"
            logger.error(f"{error_msg} (耗时: {elapsed_time:.2f}s)")
            
            # 转换为自定义异常
            raise self._convert_exception(e, error_msg)

    async def _generate_without_stream(
        self, 
        system_prompt: str, 
        user_prompt: str, 
        stage: str,
        **params
    ) -> str:
        """非流式生成"""
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ]
        
        response = await self.client.chat.completions.create(
            model=self.config.model,
            messages=messages,
            **params
        )
        
        return response.choices[0].message.content

    async def _generate_with_stream(
        self, 
        system_prompt: str, 
        user_prompt: str, 
        stage: str,
        **params
    ) -> str:
        """流式生成"""
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ]
        
        content_parts = []
        
        async for chunk in await self.client.chat.completions.create(
            model=self.config.model,
            messages=messages,
            stream=True,
            **params
        ):
            if chunk.choices[0].delta.content:
                content_parts.append(chunk.choices[0].delta.content)
                
                # 可以在这里添加流式输出的回调
                if len(content_parts) % 10 == 0:  # 每10个chunk记录一次进度
                    logger.debug(f"{stage}流式生成进度: {len(content_parts)} chunks")
        
        return "".join(content_parts)

    async def generate_with_stream_callback(
        self,
        system_prompt: str,
        user_prompt: str,
        callback=None,
        **kwargs
    ) -> AsyncGenerator[str, None]:
        """
        带回调的流式生成

        Args:
            system_prompt: 系统提示词
            user_prompt: 用户提示词
            callback: 流式回调函数
            **kwargs: 其他参数

        Yields:
            生成的文本片段
        """
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ]
        
        params = {
            "temperature": kwargs.get("temperature", self.config.temperature),
            "top_p": kwargs.get("top_p", self.config.top_p),
            "presence_penalty": kwargs.get("presence_penalty", self.config.presence_penalty),
        }
        params = {k: v for k, v in params.items() if v is not None}
        
        try:
            async for chunk in await self.client.chat.completions.create(
                model=self.config.model,
                messages=messages,
                stream=True,
                **params
            ):
                if chunk.choices[0].delta.content:
                    content = chunk.choices[0].delta.content
                    
                    if callback:
                        await callback(content)
                    
                    yield content
                    
        except Exception as e:
            raise self._convert_exception(e, "流式生成失败")

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

        Args:
            text: 要计算的文本

        Returns:
            token数量
        """
        if not self.tokenizer or not text:
            return 0
        
        try:
            return len(self.tokenizer.encode(text))
        except Exception as e:
            logger.warning(f"Token计数失败: {e}")
            return 0

    def _convert_exception(self, original_exception: Exception, message: str) -> AIServiceError:
        """
        转换原始异常为自定义异常

        Args:
            original_exception: 原始异常
            message: 错误消息

        Returns:
            自定义异常实例
        """
        if isinstance(original_exception, APITimeoutError):
            return CustomTimeoutError(f"请求超时: {message}")
        elif isinstance(original_exception, APIConnectionError):
            return NetworkError(f"网络连接失败: {message}")
        elif isinstance(original_exception, RateLimitError):
            return CustomRateLimitError(f"API频率限制: {message}")
        elif isinstance(original_exception, APIStatusError):
            if original_exception.status_code == 429:
                return CustomRateLimitError(f"API频率限制: {message}")
            elif original_exception.status_code == 402:
                return QuotaExceededError(f"配额不足: {message}")
            elif original_exception.status_code == 404:
                return ModelNotFoundError(self.config.model)
            else:
                return AIServiceError(f"API错误 ({original_exception.status_code}): {message}")
        else:
            return AIServiceError(f"未知错误: {message}")

    async def test_connection(self) -> bool:
        """
        测试AI服务连接

        Returns:
            连接是否成功
        """
        try:
            await self.generate_text(
                system_prompt="你是一个测试助手。",
                user_prompt="请回复'连接成功'",
                stage="连接测试"
            )
            return True
        except Exception as e:
            logger.error(f"AI服务连接测试失败: {e}")
            return False

    def get_model_info(self) -> dict:
        """
        获取模型信息

        Returns:
            模型信息字典
        """
        return {
            "model": self.config.model,
            "api_url": self.config.api_url,
            "timeout": self.config.timeout,
            "max_retries": self.config.max_retries,
            "temperature": self.config.temperature,
            "top_p": self.config.top_p,
            "presence_penalty": self.config.presence_penalty
        } 