import logging
from typing import Any, List
from .providers import base_provider
from .providers import zhipu, moonshot, minimax, qwen, tiangong, baichuan, wenxin, xunfei, dify, fastgpt, coze, litellm
from .exceptions import ProviderError

logger = logging.getLogger(__name__)

class UnionLLM:
    def __init__(self, provider: str, **kwargs):
        self.provider = provider.lower()
        self.litellm_call_type = None
        if self.provider == "zhipuai":
            self.provider_instance = zhipu.ZhipuAIProvider(**kwargs)
        elif self.provider == "moonshot":
            self.provider_instance = moonshot.MoonshotAIProvider(**kwargs)
        elif self.provider == "minimax":
            self.provider_instance = minimax.MinimaxAIProvider(**kwargs)
        elif self.provider == "qwen":
            self.provider_instance = qwen.QwenAIProvider(**kwargs)
        elif self.provider == "tiangong":
            self.provider_instance = tiangong.TianGongAIProvider(**kwargs)
        elif self.provider == "baichuan":
            self.provider_instance = baichuan.BaiChuanAIProvider(**kwargs)
        elif self.provider == "wenxin":
            self.provider_instance = wenxin.WenXinAIProvider(**kwargs)
        elif self.provider == "xunfei":
            self.provider_instance = xunfei.XunfeiAIProvider(**kwargs)
        elif self.provider == "dify":
            self.provider_instance = dify.DifyAIProvider(**kwargs)
        elif self.provider == "fastgpt":
            self.provider_instance = fastgpt.FastGPTProvider(**kwargs)
        elif self.provider == "coze":
            self.provider_instance = coze.CozeAIProvider(**kwargs)
        else:
            if_litellm_support, support_type = self.check_litellm_providers(provider=self.provider)
            if if_litellm_support:
                self.provider_instance = litellm.LiteLLMProvider(**kwargs)
                if support_type == 1:
                    self.litellm_call_type = 1
                elif support_type == 2:
                    self.litellm_call_type = 2
            else:
                raise ProviderError(f"Provider '{self.provider}' is not supported.")

    def completion(self, model: str, messages: List[str], **kwargs) -> Any:
        if not self.provider_instance:
            raise ProviderError(f"Provider '{self.provider}' is not initialized.")
        
        if self.litellm_call_type:
            if self.litellm_call_type == 1:
                # 判断model是否以self.provider开头，如果不是则加上
                if not model.startswith(self.provider):
                    model = f"{self.provider}/{model}"
                return self.provider_instance.completion(model, messages, **kwargs)
            elif self.litellm_call_type == 2:
                return self.provider_instance.completion(model, messages, **kwargs)
        else:
            return self.provider_instance.completion(model, messages, **kwargs)
        
    def check_litellm_providers(self, provider: str) -> bool:
        if provider in ['azure', 'sagemaker', 'bedrock', 'vertex_ai', 'palm', 'gemini', 'mistral', 'cloudflare', 'huggingface', 'replicate', 'together_ai', 'openrouter', 'baseten', 'nlp_cloud', 'petals', 'ollama', 'perplexity', 'groq', 'anyscale', 'watsonx', 'voyage', 'xinference']:
            return True, 1
        elif provider in ['openai', 'claude', 'cohere', 'ai21', 'luminous']:
            return True, 2
        else:
            return False, 0