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

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)
        elif self.provider:
            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.")
        else:
            self.provider_instance = litellm.LiteLLMProvider(**kwargs)

    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:
                # Jugde whether the model starts with self.provider, if not, add it
                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:
        # Judge whether the provider is supported by LiteLLM, and if provider name should be added to the model name
        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']:
            # provider name should be added to the model name
            return True, 1
        elif provider in ['openai', 'claude', 'cohere', 'ai21', 'luminous']:
            # provider name should not be added to the model name
            return True, 2
        else:
            return False, 0
        
# This is the new function to simplify calling
def unionchat(model: str, messages: List[dict], **kwargs) -> Any:
    client = UnionLLM(**kwargs)
    return client.completion(model=model, messages=messages, **kwargs)