"""
统一 LLM 客户端 - 支持多种提供商

支持的提供商:
- ollama: 本地运行 (免费，但占用 CPU/内存)
- deepseek: 云端 API (便宜，¥1/百万 tokens)
- openai: 云端 API (GPT-4)
- qwen: 阿里云通义千问 (便宜，国内稳定)
"""

import httpx
from typing import Dict, List, Optional, Any
import os
import json


class LLMClient:
    """统一 LLM 客户端"""

    def __init__(self):
        self.provider = os.getenv("LLM_PROVIDER", "deepseek").lower()
        self.client = self._create_client()

    def _create_client(self) -> 'LLMProvider':
        """根据配置创建对应的提供商客户端"""
        if self.provider == "ollama":
            return OllamaClient()
        elif self.provider == "deepseek":
            return OpenAICompatibleClient(
                api_key=os.getenv("DEEPSEEK_API_KEY"),
                base_url=os.getenv("DEEPSEEK_BASE_URL", "https://api.deepseek.com/v1"),
                model=os.getenv("DEEPSEEK_MODEL", "deepseek-chat")
            )
        elif self.provider == "openai":
            return OpenAICompatibleClient(
                api_key=os.getenv("OPENAI_API_KEY"),
                base_url=os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1"),
                model=os.getenv("OPENAI_MODEL", "gpt-4o-mini")
            )
        elif self.provider == "qwen":
            return OpenAICompatibleClient(
                api_key=os.getenv("QWEN_API_KEY"),
                base_url=os.getenv("QWEN_BASE_URL", "https://dashscope.aliyuncs.com/compatible-mode/v1"),
                model=os.getenv("QWEN_MODEL", "qwen-plus")
            )
        else:
            raise ValueError(f"不支持的 LLM 提供商: {self.provider}")

    def generate(self, prompt: str, **kwargs) -> str:
        """生成文本"""
        return self.client.generate(prompt, **kwargs)

    def generate_json(self, prompt: str, **kwargs) -> Dict:
        """生成 JSON 格式响应"""
        return self.client.generate_json(prompt, **kwargs)

    def get_embedding(self, text: str) -> List[float]:
        """获取向量嵌入"""
        return self.client.get_embedding(text)


class LLMProvider:
    """LLM 提供商基类"""
    def generate(self, prompt: str, **kwargs) -> str:
        raise NotImplementedError

    def generate_json(self, prompt: str, **kwargs) -> Dict:
        raise NotImplementedError

    def get_embedding(self, text: str) -> List[float]:
        raise NotImplementedError


class OllamaClient(LLMProvider):
    """Ollama 本地客户端"""

    def __init__(self):
        self.host = os.getenv("OLLAMA_HOST", "http://127.0.0.1:11434").rstrip('/')
        self.chat_model = os.getenv("OLLAMA_CHAT_MODEL", "llama3:8b")
        self.embedding_model = os.getenv("OLLAMA_EMBEDDING_MODEL", "nomic-embed-text")
        self.client = httpx.Client(base_url=self.host, timeout=300.0, trust_env=False)

    def generate(self, prompt: str, **kwargs) -> str:
        response = self.client.post("/api/generate", json={
            "model": self.chat_model,
            "prompt": prompt,
            "stream": False,
            **kwargs
        })
        response.raise_for_status()
        return response.json()["response"]

    def generate_json(self, prompt: str, **kwargs) -> Dict:
        response = self.client.post("/api/generate", json={
            "model": self.chat_model,
            "prompt": prompt,
            "stream": False,
            "format": "json",
            **kwargs
        })
        response.raise_for_status()
        content = response.json()["response"]
        # 清理 markdown 代码块
        if content.startswith("```json"):
            content = content[7:]
        if content.startswith("```"):
            content = content[3:]
        if content.endswith("```"):
            content = content[:-3]
        content = content.strip()
        return json.loads(content)

    def get_embedding(self, text: str) -> List[float]:
        response = self.client.post("/api/embeddings", json={
            "model": self.embedding_model,
            "prompt": text
        })
        response.raise_for_status()
        return response.json().get("embedding", [])


class OpenAICompatibleClient(LLMProvider):
    """OpenAI 兼容客户端 (适用于 DeepSeek, Qwen 等)"""

    def __init__(self, api_key: str, base_url: str, model: str):
        self.model = model
        self.base_url = base_url
        self.api_key = api_key
        # 延迟初始化客户端，在实际调用时才检查 API_KEY
        self._client = None

    @property
    def client(self):
        """延迟初始化客户端"""
        if self._client is None:
            if not self.api_key:
                raise ValueError(f"请设置 {self.base_url} 的 API_KEY")
            self._client = httpx.Client(
                base_url=self.base_url,
                timeout=120.0,
                headers={"Authorization": f"Bearer {self.api_key}"}
            )
        return self._client

    def generate(self, prompt: str, **kwargs) -> str:
        response = self.client.post("/chat/completions", json={
            "model": self.model,
            "messages": [{"role": "user", "content": prompt}],
            **kwargs
        })
        response.raise_for_status()
        return response.json()["choices"][0]["message"]["content"]

    def generate_json(self, prompt: str, **kwargs) -> Dict:
        response = self.client.post("/chat/completions", json={
            "model": self.model,
            "messages": [{"role": "user", "content": prompt}],
            "response_format": {"type": "json_object"},
            **kwargs
        })
        response.raise_for_status()
        content = response.json()["choices"][0]["message"]["content"]
        return json.loads(content)

    def get_embedding(self, text: str) -> List[float]:
        # DeepSeek 不支持 embedding，使用替代方案或返回空
        if "deepseek" in self.base_url:
            print("[WARNING] DeepSeek 不支持 embedding API，使用简化向量")
            # 简化方案：基于文本内容生成伪向量（实际使用时应该用其他服务）
            import hashlib
            hash_obj = hashlib.md5(text.encode())
            # 生成 768 维伪向量（与 nomic-embed-text 维度一致）
            return [float((hash_obj.digest()[i % 16] / 255.0) * 2 - 1) for _ in range(768)]

        response = self.client.post("/embeddings", json={
            "model": self.model,
            "input": text
        })
        response.raise_for_status()
        return response.json()["data"][0]["embedding"]


# 单例实例
llm_client = LLMClient()