"""
LLM 服务
"""
import logging
from typing import AsyncGenerator, Dict, List, Optional

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from ai_platform.providers import (
    BaseLLMProvider,
    LLMConfig,
    LLMMessage,
    LLMResponse,
    ProviderRegistry,
)
from ai_platform.providers.base import LLMStreamChunk, ToolDefinition

logger = logging.getLogger(__name__)


class LLMService:
    """
    LLM 服务
    
    统一管理 LLM 调用，支持多种提供商
    """
    
    def __init__(self, db: Optional[AsyncSession] = None):
        self._provider_cache: Dict[str, BaseLLMProvider] = {}
        self._db = db
    
    async def _get_provider_async(self, model_id: str) -> tuple:
        """
        异步获取模型对应的提供商实例
        
        Args:
            model_id: 模型 ID
            
        Returns:
            (provider, model_name)
        """
        from ai_platform.models import LLMModel, LLMProvider
        
        if not self._db:
            raise ValueError("数据库会话未初始化")
        
        # 查询模型
        result = await self._db.execute(
            select(LLMModel).where(
                LLMModel.id == model_id,
                LLMModel.is_active == True,
                LLMModel.is_deleted == False
            )
        )
        model = result.scalar_one_or_none()
        if not model:
            raise ValueError(f'模型不存在或已禁用: {model_id}')
        
        # 查询提供商
        provider_result = await self._db.execute(
            select(LLMProvider).where(LLMProvider.id == model.provider_id)
        )
        provider = provider_result.scalar_one_or_none()
        if not provider or not provider.is_active:
            raise ValueError(f'提供商不存在或已禁用')
        
        # 缓存提供商实例
        cache_key = str(provider.id)
        if cache_key not in self._provider_cache:
            provider_instance = ProviderRegistry.create_instance(
                provider_type=provider.provider_type,
                api_key=provider.api_key,
                api_base=provider.api_base or "",
                ollama_host=provider.ollama_host or "",
            )
            if not provider_instance:
                raise ValueError(f'不支持的提供商类型: {provider.provider_type}')
            self._provider_cache[cache_key] = provider_instance
        
        return self._provider_cache[cache_key], model.model_name
    
    def _get_provider_sync(self, model_id: str, model_data: dict) -> tuple:
        """
        同步获取提供商实例（使用预加载的数据）
        
        Args:
            model_id: 模型 ID
            model_data: 预加载的模型和提供商数据
            
        Returns:
            (provider, model_name)
        """
        provider_type = model_data.get('provider_type')
        api_key = model_data.get('api_key', '')
        api_base = model_data.get('api_base', '')
        ollama_host = model_data.get('ollama_host', '')
        model_name = model_data.get('model_name', '')
        provider_id = model_data.get('provider_id', '')
        
        # 缓存提供商实例
        cache_key = str(provider_id)
        if cache_key not in self._provider_cache:
            provider_instance = ProviderRegistry.create_instance(
                provider_type=provider_type,
                api_key=api_key,
                api_base=api_base,
                ollama_host=ollama_host,
            )
            if not provider_instance:
                raise ValueError(f'不支持的提供商类型: {provider_type}')
            self._provider_cache[cache_key] = provider_instance
        
        return self._provider_cache[cache_key], model_name
    
    def chat_with_provider(
        self,
        provider: BaseLLMProvider,
        model_name: str,
        messages: List[Dict[str, str]],
        temperature: float = 0.7,
        max_tokens: int = 2048,
        tools: List[Dict] = None,
        tool_choice: str = 'auto',
        **kwargs
    ) -> LLMResponse:
        """
        使用指定提供商进行同步对话
        
        Args:
            provider: 提供商实例
            model_name: 模型名称
            messages: 消息列表 [{"role": "user", "content": "..."}]
            temperature: 温度参数
            max_tokens: 最大 Token
            tools: 工具定义列表 [{"name": "...", "description": "...", "parameters": {...}}]
            tool_choice: 工具选择策略 (auto, none, required, 或具体工具名)
            **kwargs: 其他参数
            
        Returns:
            LLMResponse
        """
        llm_messages = self._convert_messages(messages)
        
        # 转换工具定义
        tool_definitions = None
        if tools:
            tool_definitions = [
                ToolDefinition(
                    name=t['name'],
                    description=t.get('description', ''),
                    parameters=t.get('parameters', {}),
                )
                for t in tools
            ]
        
        config = LLMConfig(
            model=model_name,
            temperature=temperature,
            max_tokens=max_tokens,
            tools=tool_definitions,
            tool_choice=tool_choice,
            **kwargs
        )
        
        return provider.chat(llm_messages, config)
    
    def _convert_messages(self, messages: List[Dict]) -> List[LLMMessage]:
        """转换消息格式，支持 tool 消息"""
        llm_messages = []
        for m in messages:
            msg = LLMMessage(
                role=m['role'],
                content=m.get('content', ''),
                name=m.get('name'),
                tool_calls=m.get('tool_calls'),
                tool_call_id=m.get('tool_call_id'),
            )
            llm_messages.append(msg)
        return llm_messages
    
    async def chat_async(
        self,
        model_id: str,
        messages: List[Dict[str, str]],
        temperature: float = 0.7,
        max_tokens: int = 2048,
        tools: List[Dict] = None,
        tool_choice: str = 'auto',
        **kwargs
    ) -> LLMResponse:
        """
        异步对话
        
        Args:
            model_id: 模型 ID
            messages: 消息列表
            temperature: 温度参数
            max_tokens: 最大 Token
            tools: 工具定义列表
            tool_choice: 工具选择策略
            **kwargs: 其他参数
            
        Returns:
            LLMResponse
        """
        provider, model_name = await self._get_provider_async(model_id)
        
        llm_messages = self._convert_messages(messages)
        
        # 转换工具定义
        tool_definitions = None
        if tools:
            tool_definitions = [
                ToolDefinition(
                    name=t['name'],
                    description=t.get('description', ''),
                    parameters=t.get('parameters', {}),
                )
                for t in tools
            ]
        
        config = LLMConfig(
            model=model_name,
            temperature=temperature,
            max_tokens=max_tokens,
            tools=tool_definitions,
            tool_choice=tool_choice,
            **kwargs
        )
        
        return await provider.chat_async(llm_messages, config)
    
    async def chat_stream(
        self,
        model_id: str,
        messages: List[Dict[str, str]],
        temperature: float = 0.7,
        max_tokens: int = 2048,
        tools: List[Dict] = None,
        tool_choice: str = 'auto',
        **kwargs
    ) -> AsyncGenerator[LLMStreamChunk, None]:
        """
        异步流式对话
        
        Args:
            model_id: 模型 ID
            messages: 消息列表
            temperature: 温度参数
            max_tokens: 最大 Token
            tools: 工具定义列表
            tool_choice: 工具选择策略
            **kwargs: 其他参数
            
        Yields:
            LLMStreamChunk
        """
        provider, model_name = await self._get_provider_async(model_id)
        
        llm_messages = self._convert_messages(messages)
        
        # 转换工具定义
        tool_definitions = None
        if tools:
            tool_definitions = [
                ToolDefinition(
                    name=t['name'],
                    description=t.get('description', ''),
                    parameters=t.get('parameters', {}),
                )
                for t in tools
            ]
        
        config = LLMConfig(
            model=model_name,
            temperature=temperature,
            max_tokens=max_tokens,
            tools=tool_definitions,
            tool_choice=tool_choice,
            **kwargs
        )
        
        async for chunk in provider.chat_stream(llm_messages, config):
            yield chunk
    
    def chat_stream_sync(
        self,
        model_id: str,
        messages: List[Dict[str, str]],
        temperature: float = 0.7,
        max_tokens: int = 2048,
        tools: List[Dict] = None,
        tool_choice: str = 'auto',
        **kwargs
    ):
        """
        同步流式对话（生成器）
        
        注意：此方法会在内部运行异步代码来获取 provider，
        需要确保 LLMService 初始化时传入了 db_session
        
        Args:
            model_id: 模型 ID
            messages: 消息列表
            temperature: 温度参数
            max_tokens: 最大 Token
            tools: 工具定义列表
            tool_choice: 工具选择策略
            **kwargs: 其他参数
            
        Yields:
            LLMStreamChunk
        """
        import asyncio
        import concurrent.futures
        
        # 在同步上下文中运行异步代码获取 provider
        async def get_provider():
            return await self._get_provider_async(model_id)
        
        loop = asyncio.get_event_loop()
        if loop.is_running():
            # 如果已有事件循环在运行，使用线程池
            with concurrent.futures.ThreadPoolExecutor() as executor:
                future = executor.submit(asyncio.run, get_provider())
                provider, model_name = future.result()
        else:
            provider, model_name = loop.run_until_complete(get_provider())
        
        # 使用获取到的 provider 进行流式调用
        yield from self.chat_stream_sync_with_provider(
            provider=provider,
            model_name=model_name,
            messages=messages,
            temperature=temperature,
            max_tokens=max_tokens,
            tools=tools,
            tool_choice=tool_choice,
            **kwargs
        )
    
    def chat_stream_sync_with_provider(
        self,
        provider: BaseLLMProvider,
        model_name: str,
        messages: List[Dict[str, str]],
        temperature: float = 0.7,
        max_tokens: int = 2048,
        tools: List[Dict] = None,
        tool_choice: str = 'auto',
        **kwargs
    ):
        """
        使用指定提供商进行同步流式对话（生成器）
        
        Args:
            provider: 提供商实例
            model_name: 模型名称
            messages: 消息列表
            temperature: 温度参数
            max_tokens: 最大 Token
            tools: 工具定义列表
            tool_choice: 工具选择策略
            **kwargs: 其他参数
            
        Yields:
            LLMStreamChunk
        """
        llm_messages = self._convert_messages(messages)
        
        # 转换工具定义
        tool_definitions = None
        if tools:
            tool_definitions = [
                ToolDefinition(
                    name=t['name'],
                    description=t.get('description', ''),
                    parameters=t.get('parameters', {}),
                )
                for t in tools
            ]
        
        config = LLMConfig(
            model=model_name,
            temperature=temperature,
            max_tokens=max_tokens,
            tools=tool_definitions,
            tool_choice=tool_choice,
            **kwargs
        )
        
        for chunk in provider.chat_stream_sync(llm_messages, config):
            yield chunk
    
    @staticmethod
    def get_available_providers() -> List[Dict]:
        """获取可用的提供商列表"""
        return ProviderRegistry.get_all_types()
    
    @staticmethod
    def get_default_models(provider_type: str) -> List[Dict]:
        """获取提供商的默认模型列表"""
        return ProviderRegistry.get_default_models(provider_type)
