import os
import json
from typing import Any, Dict, Iterable, List, Optional
from dotenv import load_dotenv

try:
    import httpx
except Exception:  # pragma: no cover
    httpx = None

load_dotenv()

_shared_async_client: Any = None

def _get_async_client() -> Any:
    global _shared_async_client
    if httpx is None:
        return None
    if _shared_async_client is None:
        _shared_async_client = httpx.AsyncClient(
            headers={},
            limits=httpx.Limits(max_keepalive_connections=20, max_connections=50),
            http2=True,
        )
    return _shared_async_client

_DEFAULT_ENDPOINTS: Dict[str, str] = {
    "openai": "https://api.openai.com/v1/chat/completions",
    "siliconflow": "https://api.siliconflow.cn/v1/chat/completions",
    "deepseek": "https://api.deepseek.com/v1/chat/completions",
}

_DEFAULT_MODELS: Dict[str, str] = {
    "openai": "gpt-4o-mini",
    "siliconflow": "deepseek-ai/DeepSeek-V3",
    "deepseek": "deepseek-v3",
    "r1": "Pro/deepseek-ai/DeepSeek-R1",
}

def _clean_str(s: str) -> str:
    if s is None:
        return ""
    s = s.strip()
    if (s.startswith("`") and s.endswith("`")) or (s.startswith('"') and s.endswith('"')) or (s.startswith("'") and s.endswith("'")):
        s = s[1:-1].strip()
    return s

def _normalize_endpoint(ep: str) -> str:
    if not ep:
        return ep
    s = _clean_str(ep).rstrip("/")
    if s.endswith("/v1"):
        return s + "/chat/completions"
    if s.endswith("/chat/completions"):
        return s
    return s


class LLMClient:
    def __init__(self):
        self.provider = os.getenv("LLM_PROVIDER", "openai").strip().lower()
        raw_endpoint = os.getenv("LLM_ENDPOINT", "") or _DEFAULT_ENDPOINTS.get(self.provider, _DEFAULT_ENDPOINTS["openai"])
        self.endpoint = _normalize_endpoint(raw_endpoint)
        self.model = _clean_str(os.getenv("LLM_MODEL", _DEFAULT_MODELS.get(self.provider, "gpt-4o-mini")))
        api_key = os.getenv("LLM_API_KEY") or os.getenv("OPENAI_API_KEY") or os.getenv("DEEPSEEK_API_KEY") or os.getenv("SILICONFLOW_API_KEY") or ""
        self.api_key = api_key
        self.simulate = os.getenv("LLM_SIMULATE", "false").lower() == "true"
        self.timeout = int(os.getenv("LLM_TIMEOUT", "300"))

    def _headers(self) -> Dict[str, str]:
        return {
            "Authorization": f"Bearer {self.api_key}" if self.api_key else "",
            "Content-Type": "application/json",
        }

    async def chat(self, messages: List[Dict[str, Any]], tools: Optional[List[Dict[str, Any]]] = None, stream: bool = False, model: Optional[str] = None) -> Any:
        if self.simulate or httpx is None:
            if stream:
                async def _sim_stream():
                    yield {"choices": [{"delta": {"content": "模拟流式输出：检测到错误日志，建议重启或kill相关进程"}, "index": 0}]}
                return _sim_stream()
            return {
                "choices": [
                    {
                        "message": {
                            "role": "assistant",
                            "content": "模拟输出：检测到错误日志，建议重启或kill相关进程",
                            "tool_calls": [],
                        }
                    }
                ]
            }
        
        target_model = model or self.model
        payload: Dict[str, Any] = {"model": target_model, "messages": messages, "stream": stream}
        if tools:
            payload["tools"] = tools
            payload["tool_choice"] = "auto"
        
        if stream:
            async def _stream_gen():
                client = _get_async_client()
                async with client.stream("POST", self.endpoint, headers=self._headers(), json=payload, timeout=self.timeout) as resp:
                        resp.raise_for_status()
                        async for line in resp.aiter_lines():
                             if not line or not line.startswith("data: "):
                                 continue
                             data_str = line[6:].strip()
                             if data_str == "[DONE]":
                                 break
                             try:
                                 yield json.loads(data_str)
                             except:
                                 continue
            return _stream_gen()

        client = _get_async_client()
        resp = await client.post(self.endpoint, headers=self._headers(), json=payload, timeout=self.timeout)
        resp.raise_for_status()
        return resp.json()