"""LLM Providers Module"""
from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional
import json


class BaseLLMProvider(ABC):
    """LLM 提供商基类"""
    
    @abstractmethod
    async def chat(
        self, 
        messages: List[Dict[str, str]], 
        **kwargs
    ) -> str:
        """
        发送聊天请求
        
        Args:
            messages: 消息列表
            **kwargs: 其他参数
        
        Returns:
            回复内容
        """
        pass
    
    @abstractmethod
    async def structured_output(
        self, 
        messages: List[Dict[str, str]], 
        schema: Dict[str, Any]
    ) -> Any:
        """
        获取结构化输出
        
        Args:
            messages: 消息列表
            schema: 输出 schema
        
        Returns:
            结构化结果
        """
        pass
    
    @abstractmethod
    async def async_get_embedding(self, text: str) -> List[float]:
        """
        异步获取文本嵌入向量
        """
        pass
    
    def get_embedding(self, text: str) -> List[float]:
        """同步获取嵌入（默认实现，可覆盖）"""
        import asyncio
        return asyncio.run(self.async_get_embedding(text))


class OpenAIProvider(BaseLLMProvider):
    """OpenAI 提供商"""
    
    def __init__(
        self, 
        api_key: str, 
        model: str = "gpt-4o",
        embedding_model: str = "text-embedding-3-small"
    ):
        from openai import AsyncOpenAI
        self.client = AsyncOpenAI(api_key=api_key)
        self.model = model
        self.embedding_model = embedding_model
    
    async def chat(
        self, 
        messages: List[Dict[str, str]], 
        **kwargs
    ) -> str:
        response = await self.client.chat.completions.create(
            model=self.model,
            messages=messages,
            **kwargs
        )
        return response.choices[0].message.content
    
    async def structured_output(
        self, 
        messages: List[Dict[str, str]], 
        schema: Dict[str, Any]
    ) -> Any:
        response = await self.client.chat.completions.create(
            model=self.model,
            messages=messages,
            response_format={
                "type": "json_object",
                "schema": schema
            }
        )
        return json.loads(response.choices[0].message.content)
    
    def get_embedding(self, text: str) -> List[float]:
        import asyncio
        async def _create():
            return await self.client.embeddings.create(
                model=self.embedding_model,
                input=text
            )
        response = asyncio.run(_create())
        return response.data[0].embedding


class AnthropicProvider(BaseLLMProvider):
    """Anthropic 提供商"""
    
    def __init__(self, api_key: str, model: str = "claude-3-5-sonnet-20241022"):
        import anthropic
        self.client = anthropic.AsyncAnthropic(api_key=api_key)
        self.model = model
    
    async def chat(
        self, 
        messages: List[Dict[str, str]], 
        **kwargs
    ) -> str:
        # 转换消息格式
        system = None
        for msg in messages:
            if msg["role"] == "system":
                system = msg["content"]
        
        filtered = [m for m in messages if m["role"] != "system"]
        
        response = await self.client.messages.create(
            model=self.model,
            system=system,
            messages=filtered,
            **kwargs
        )
        return response.content[0].text
    
    async def structured_output(
        self, 
        messages: List[Dict[str, str]], 
        schema: Dict[str, Any]
    ) -> Any:
        # Anthropic 不直接支持 schema，使用提示词约束
        schema_str = json.dumps(schema)
        enhanced_messages = messages + [{
            "role": "user",
            "content": f"Respond with valid JSON matching this schema: {schema_str}"
        }]
        result = await self.chat(enhanced_messages)
        return json.loads(result)
    
    def get_embedding(self) -> List[float]:
        # Anthropic 没有 embedding API，使用 OpenAI 或本地模型
        raise NotImplementedError("Anthropic does not provide embeddings API")


class OllamaProvider(BaseLLMProvider):
    """Ollama 本地模型提供商"""
    
    def __init__(
        self,
        base_url: str = "http://localhost:11434",
        model: str = "llama3",
        embedding_model: str = "nomic-embed-text"
    ):
        self.base_url = base_url
        self.model = model
        self.embedding_model = embedding_model
    
    async def chat(
        self,
        messages: List[Dict[str, str]],
        **kwargs
    ) -> str:
        import aiohttp
        async with aiohttp.ClientSession() as session:
            payload = {
                "model": self.model,
                "messages": messages,
                **kwargs
            }
            async with session.post(
                f"{self.base_url}/api/chat",
                json=payload
            ) as resp:
                result = await resp.json()
                return result["message"]["content"]
    
    async def structured_output(
        self,
        messages: List[Dict[str, str]],
        schema: Dict[str, Any]
    ) -> Any:
        schema_str = json.dumps(schema)
        enhanced_messages = messages + [{
            "role": "user",
            "content": f"Respond with valid JSON matching this schema: {schema_str}"
        }]
        result = await self.chat(enhanced_messages)
        return json.loads(result)
    
    def get_embedding(self, text: str) -> List[float]:
        import aiohttp
        import asyncio
        
        async def _get_embedding():
            async with aiohttp.ClientSession() as session:
                payload = {
                    "model": self.embedding_model,
                    "input": text
                }
                async with session.post(
                    f"{self.base_url}/api/embeddings",
                    json=payload
                ) as resp:
                    result = await resp.json()
                    return result["embedding"]
        
        return asyncio.run(_get_embedding())


class GoogleGeminiProvider(BaseLLMProvider):
    """Google Gemini 提供商"""
    
    def __init__(
        self,
        api_key: str,
        model: str = "gemini-1.5-pro"
    ):
        import google.genai as genai
        genai.configure(api_key=api_key)
        self.client = genai
        self.model = model
    
    async def chat(
        self,
        messages: List[Dict[str, str]],
        **kwargs
    ) -> str:
        import asyncio
        
        system = ""
        filtered = []
        for msg in messages:
            if msg["role"] == "system":
                system = msg["content"]
            else:
                filtered.append(msg)
        
        async def _chat():
            model = self.client.GenerativeModel(self.model)
            contents = [m["content"] for m in filtered]
            response = await model.generate_content_async(
                contents,
                **kwargs
            )
            return response.text
        
        return await _chat()
    
    async def structured_output(
        self,
        messages: List[Dict[str, str]],
        schema: Dict[str, Any]
    ) -> Any:
        schema_str = json.dumps(schema)
        enhanced_messages = messages + [{
            "role": "user",
            "content": f"Respond with valid JSON matching this schema: {schema_str}"
        }]
        result = await self.chat(enhanced_messages)
        return json.loads(result)
    
    def get_embedding(self, text: str) -> List[float]:
        import asyncio
        async def _create():
            return await self.client.embeddings.create(
                model="text-embedding-004",
                input=text
            )
        result = asyncio.run(_create())
        return result.data[0].embedding_values


class AzureOpenAIProvider(BaseLLMProvider):
    """Azure OpenAI 提供商"""
    
    def __init__(
        self,
        api_key: str,
        endpoint: str,
        api_version: str = "2024-02-01",
        model: str = "gpt-4o",
        embedding_model: str = "text-embedding-3-small"
    ):
        from openai import AsyncAzureOpenAI
        self.client = AsyncAzureOpenAI(
            api_key=api_key,
            azure_endpoint=endpoint,
            api_version=api_version
        )
        self.model = model
        self.embedding_model = embedding_model
    
    async def chat(
        self,
        messages: List[Dict[str, str]],
        **kwargs
    ) -> str:
        response = await self.client.chat.completions.create(
            model=self.model,
            messages=messages,
            **kwargs
        )
        return response.choices[0].message.content
    
    async def structured_output(
        self,
        messages: List[Dict[str, str]],
        schema: Dict[str, Any]
    ) -> Any:
        response = await self.client.chat.completions.create(
            model=self.model,
            messages=messages,
            response_format={
                "type": "json_object",
                "schema": schema
            }
        )
        return json.loads(response.choices[0].message.content)
    
    def get_embedding(self, text: str) -> List[float]:
        import asyncio
        async def _create():
            return await self.client.embeddings.create(
                model=self.embedding_model,
                input=text
            )
        response = asyncio.run(_create())
        return response.data[0].embedding


class CohereProvider(BaseLLMProvider):
    """Cohere 提供商 - 擅长 embedding"""
    
    def __init__(
        self,
        api_key: str,
        model: str = "command-r-plus",
        embedding_model: str = "embed-english-v3.0"
    ):
        import cohere
        self.client = cohere.AsyncClient(api_key=api_key)
        self.model = model
        self.embedding_model = embedding_model
    
    async def chat(
        self,
        messages: List[Dict[str, str]],
        **kwargs
    ) -> str:
        system = ""
        filtered = []
        for msg in messages:
            if msg["role"] == "system":
                system = msg["content"]
            else:
                filtered.append(msg)
        
        response = await self.client.chat(
            message=filtered[-1]["content"],
            chat_history=filtered[:-1],
            preamble=system,
            **kwargs
        )
        return response.text
    
    async def structured_output(
        self,
        messages: List[Dict[str, str]],
        schema: Dict[str, Any]
    ) -> Any:
        schema_str = json.dumps(schema)
        enhanced_messages = messages + [{
            "role": "user",
            "content": f"Respond with valid JSON matching this schema: {schema_str}"
        }]
        result = await self.chat(enhanced_messages)
        return json.loads(result)
    
    def get_embedding(self, text: str) -> List[float]:
        import asyncio
        async def _get():
            response = await self.client.embed(
                texts=[text],
                model=self.embedding_model,
                input_type="search_query"
            )
            return response.embeddings[0]
        return asyncio.run(_get())


class LocalEmbeddingProvider(BaseLLMProvider):
    """本地 Embedding 提供商 (sentence-transformers)"""
    
    def __init__(
        self,
        model_name: str = "sentence-transformers/all-MiniLM-L6-v2"
    ):
        from sentence_transformers import SentenceTransformer
        self.model = SentenceTransformer(model_name)
        self.embedding_dim = self.model.get_sentence_embedding_dimension()
    
    async def chat(self, messages: List[Dict[str, str]], **kwargs) -> str:
        raise NotImplementedError("Local embedding provider does not support chat")
    
    async def structured_output(
        self,
        messages: List[Dict[str, str]],
        schema: Dict[str, Any]
    ) -> Any:
        raise NotImplementedError("Local embedding provider does not support structured output")
    
    def get_embedding(self, text: str) -> List[float]:
        return self.model.encode(text).tolist()


class QwenProvider(BaseLLMProvider):
    """阿里 Qwen (通义千问) 提供商 - DashScope"""
    
    def __init__(
        self,
        api_key: str,
        model: str = "qwen-turbo",
        base_url: str = "https://dashscope.aliyuncs.com/compatible-mode/v1",
        embedding_model: str = "text-embedding-v3"
    ):
        from openai import AsyncOpenAI
        self.client = AsyncOpenAI(
            api_key=api_key,
            base_url=base_url
        )
        self.model = model
        self.embedding_model = embedding_model
    
    async def chat(
        self,
        messages: List[Dict[str, str]],
        **kwargs
    ) -> str:
        response = await self.client.chat.completions.create(
            model=self.model,
            messages=messages,
            **kwargs
        )
        return response.choices[0].message.content
    
    async def structured_output(
        self,
        messages: List[Dict[str, str]],
        schema: Dict[str, Any]
    ) -> Any:
        response = await self.client.chat.completions.create(
            model=self.model,
            messages=messages,
            response_format={
                "type": "json_object",
                "schema": schema
            }
        )
        return json.loads(response.choices[0].message.content)
    
    def get_embedding(self, text: str) -> List[float]:
        import asyncio
        async def _create():
            return await self.client.embeddings.create(
                model=self.embedding_model,
                input=text
            )
        response = asyncio.run(_create())
        return response.data[0].embedding


class ZhipuAIProvider(BaseLLMProvider):
    """智谱 AI (ChatGLM) 提供商"""
    
    def __init__(
        self,
        api_key: str,
        model: str = "glm-4",
        embedding_model: str = "embedding-3"
    ):
        import httpx
        self.api_key = api_key
        self.model = model
        self.embedding_model = embedding_model
        self.base_url = "https://open.bigmodel.cn/api/paas/v4"
    
    async def chat(
        self,
        messages: List[Dict[str, str]],
        **kwargs
    ) -> str:
        import httpx
        headers = {
            "Authorization": f"Bearer {self.api_key}"
        }
        payload = {
            "model": self.model,
            "messages": messages,
            **kwargs
        }
        async with httpx.AsyncClient() as client:
            resp = await client.post(
                f"{self.base_url}/chat/completions",
                headers=headers,
                json=payload
            )
            result = resp.json()
            return result["choices"][0]["message"]["content"]
    
    async def structured_output(
        self,
        messages: List[Dict[str, str]],
        schema: Dict[str, Any]
    ) -> Any:
        schema_str = json.dumps(schema)
        enhanced_messages = messages + [{
            "role": "user",
            "content": f"Respond with valid JSON matching this schema: {schema_str}"
        }]
        result = await self.chat(enhanced_messages)
        return json.loads(result)
    
    def get_embedding(self, text: str) -> List[float]:
        import httpx
        import asyncio
        
        async def _get():
            headers = {"Authorization": f"Bearer {self.api_key}"}
            payload = {
                "model": self.embedding_model,
                "input": text
            }
            async with httpx.AsyncClient() as client:
                resp = await client.post(
                    f"{self.base_url}/embeddings",
                    headers=headers,
                    json=payload
                )
                result = resp.json()
                return result["data"][0]["embedding"]
        
        return asyncio.run(_get())


class DoubaoProvider(BaseLLMProvider):
    """字节跳动 Doubao (火山引擎方舟) 提供商"""
    
    def __init__(
        self,
        api_key: str,
        model: str = "doubao-latest",
        base_url: str = "https://ark.cn-beijing.volces.com/api/v3",
        embedding_model: str = "doubao-embedding"
    ):
        from openai import AsyncOpenAI
        self.client = AsyncOpenAI(
            api_key=api_key,
            base_url=base_url
        )
        self.model = model
        self.embedding_model = embedding_model
    
    async def chat(
        self,
        messages: List[Dict[str, str]],
        **kwargs
    ) -> str:
        response = await self.client.chat.completions.create(
            model=self.model,
            messages=messages,
            **kwargs
        )
        return response.choices[0].message.content
    
    async def structured_output(
        self,
        messages: List[Dict[str, str]],
        schema: Dict[str, Any]
    ) -> Any:
        response = await self.client.chat.completions.create(
            model=self.model,
            messages=messages,
            response_format={
                "type": "json_object",
                "schema": schema
            }
        )
        return json.loads(response.choices[0].message.content)
    
    async def async_get_embedding(self, text: str) -> List[float]:
        response = await self.client.embeddings.create(
            model=self.embedding_model,
            input=text
        )
        return response.data[0].embedding
    
    def get_embedding(self, text: str) -> List[float]:
        import asyncio
        return asyncio.run(self.async_get_embedding(text))


class MiniMaxProvider(BaseLLMProvider):
    """MiniMax 提供商"""
    
    def __init__(
        self,
        api_key: str,
        model: str = "abab6.5s-chat",
        embedding_model: str = "embo-01"
    ):
        import httpx
        self.api_key = api_key
        self.model = model
        self.embedding_model = embedding_model
        self.base_url = "https://api.minimax.chat/v1"
    
    async def chat(
        self,
        messages: List[Dict[str, str]],
        **kwargs
    ) -> str:
        import httpx
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        payload = {
            "model": self.model,
            "messages": messages,
            **kwargs
        }
        async with httpx.AsyncClient() as client:
            resp = await client.post(
                f"{self.base_url}/text/chatcompletion_v2",
                headers=headers,
                json=payload
            )
            result = resp.json()
            return result["choices"][0]["message"]["content"]
    
    async def structured_output(
        self,
        messages: List[Dict[str, str]],
        schema: Dict[str, Any]
    ) -> Any:
        schema_str = json.dumps(schema)
        enhanced_messages = messages + [{
            "role": "user",
            "content": f"Respond with valid JSON matching this schema: {schema_str}"
        }]
        result = await self.chat(enhanced_messages)
        return json.loads(result)
    
    def get_embedding(self, text: str) -> List[float]:
        import httpx
        import asyncio
        
        async def _get():
            headers = {
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json"
            }
            payload = {
                "model": self.embedding_model,
                "input": text
            }
            async with httpx.AsyncClient() as client:
                resp = await client.post(
                    f"{self.base_url}/text/embeddings",
                    headers=headers,
                    json=payload
                )
                result = resp.json()
                return result["data"][0]["embedding"]
        
        return asyncio.run(_get())


class DeepSeekProvider(BaseLLMProvider):
    """DeepSeek 提供商 - OpenAI 兼容"""
    
    def __init__(
        self,
        api_key: str,
        model: str = "deepseek-chat",
        base_url: str = "https://api.deepseek.com/v1",
        embedding_model: str = "deepseek-embedding"
    ):
        from openai import AsyncOpenAI
        self.client = AsyncOpenAI(
            api_key=api_key,
            base_url=base_url
        )
        self.model = model
        self.embedding_model = embedding_model
    
    async def chat(
        self,
        messages: List[Dict[str, str]],
        **kwargs
    ) -> str:
        response = await self.client.chat.completions.create(
            model=self.model,
            messages=messages,
            **kwargs
        )
        return response.choices[0].message.content
    
    async def structured_output(
        self,
        messages: List[Dict[str, str]],
        schema: Dict[str, Any]
    ) -> Any:
        response = await self.client.chat.completions.create(
            model=self.model,
            messages=messages,
            response_format={
                "type": "json_object",
                "schema": schema
            }
        )
        return json.loads(response.choices[0].message.content)
    
    def get_embedding(self, text: str) -> List[float]:
        import asyncio
        async def _create():
            return await self.client.embeddings.create(
                model=self.embedding_model,
                input=text
            )
        response = asyncio.run(_create())
        return response.data[0].embedding


class LLMProviderFactory:
    """LLM 提供商工厂"""
    
    _providers: Dict[str, BaseLLMProvider] = {}
    _default: str = "openai"
    
    @classmethod
    def register(
        cls, 
        name: str, 
        provider: BaseLLMProvider
    ):
        cls._providers[name] = provider
    
    @classmethod
    def get(cls, name: str = None) -> BaseLLMProvider:
        name = name or cls._default
        return cls._providers.get(name)
    
    @classmethod
    def configure(cls, config: Dict[str, Any]):
        """配置提供商"""
        if "openai" in config and config["openai"].get("api_key"):
            cls.register(
                "openai",
                OpenAIProvider(
                    api_key=config["openai"]["api_key"],
                    model=config["openai"].get("model", "gpt-4o"),
                    embedding_model=config["openai"].get("embedding_model", "text-embedding-3-small")
                )
            )
        
        if "anthropic" in config and config["anthropic"].get("api_key"):
            cls.register(
                "anthropic",
                AnthropicProvider(
                    api_key=config["anthropic"]["api_key"],
                    model=config["anthropic"].get("model", "claude-3-5-sonnet-20241022")
                )
            )
        
        if "ollama" in config:
            cls.register(
                "ollama",
                OllamaProvider(
                    base_url=config["ollama"].get("base_url", "http://localhost:11434"),
                    model=config["ollama"].get("model", "llama3"),
                    embedding_model=config["ollama"].get("embedding_model", "nomic-embed-text")
                )
            )
        
        if "gemini" in config and config["gemini"].get("api_key"):
            cls.register(
                "gemini",
                GoogleGeminiProvider(
                    api_key=config["gemini"]["api_key"],
                    model=config["gemini"].get("model", "gemini-1.5-pro")
                )
            )
        
        if "azure" in config and config["azure"].get("api_key"):
            cls.register(
                "azure",
                AzureOpenAIProvider(
                    api_key=config["azure"]["api_key"],
                    endpoint=config["azure"]["endpoint"],
                    api_version=config["azure"].get("api_version", "2024-02-01"),
                    model=config["azure"].get("model", "gpt-4o"),
                    embedding_model=config["azure"].get("embedding_model", "text-embedding-3-small")
                )
            )
        
        if "cohere" in config and config["cohere"].get("api_key"):
            cls.register(
                "cohere",
                CohereProvider(
                    api_key=config["cohere"]["api_key"],
                    model=config["cohere"].get("model", "command-r-plus"),
                    embedding_model=config["cohere"].get("embedding_model", "embed-english-v3.0")
                )
            )
        
        if "local" in config:
            cls.register(
                "local",
                LocalEmbeddingProvider(
                    model_name=config["local"].get("model_name", "sentence-transformers/all-MiniLM-L6-v2")
                )
            )
        
        if "embedding" in config:
            cls.register(
                "embedding",
                LocalEmbeddingProvider(
                    model_name=config["embedding"].get("model_name", "sentence-transformers/all-MiniLM-L6-v2")
                )
            )
        
        if "qwen" in config and config["qwen"].get("api_key"):
            cls.register(
                "qwen",
                QwenProvider(
                    api_key=config["qwen"]["api_key"],
                    model=config["qwen"].get("model", "qwen-turbo"),
                    base_url=config["qwen"].get("base_url", "https://dashscope.aliyuncs.com/compatible-mode/v1")
                )
            )
        
        if "zhipu" in config and config["zhipu"].get("api_key"):
            cls.register(
                "zhipu",
                ZhipuAIProvider(
                    api_key=config["zhipu"]["api_key"],
                    model=config["zhipu"].get("model", "glm-4")
                )
            )
        
        if "doubao" in config and config["doubao"].get("api_key"):
            cls.register(
                "doubao",
                DoubaoProvider(
                    api_key=config["doubao"]["api_key"],
                    model=config["doubao"].get("model", "doubao-latest"),
                    base_url=config["doubao"].get("base_url", "https://ark.cn-beijing.volces.com/api/v3"),
                    embedding_model=config["doubao"].get(
                        "embedding_model", "doubao-embedding"
                    ),
                )
            )
        
        if "minimax" in config and config["minimax"].get("api_key"):
            cls.register(
                "minimax",
                MiniMaxProvider(
                    api_key=config["minimax"]["api_key"],
                    model=config["minimax"].get("model", "abab6.5s-chat")
                )
            )
        
        if "deepseek" in config and config["deepseek"].get("api_key"):
            cls.register(
                "deepseek",
                DeepSeekProvider(
                    api_key=config["deepseek"]["api_key"],
                    model=config["deepseek"].get("model", "deepseek-chat"),
                    base_url=config["deepseek"].get("base_url", "https://api.deepseek.com/v1")
                )
            )
