"""
lambdagent.providers.openai_compat_provider — OpenAI-compatible API provider.

Covers: OpenAI, Ollama, DashScope/Qwen, DeepSeek, Moonshot, Zhipu, and any
OpenAI-compatible endpoint.

L02: Added chat_typed() for ChatMessage/ChatResponse support,
overrides provider_name and default_model properties. Added PRESETS
class attribute for easy provider lookup.

L03: Token-usage accumulation in chat() so get_usage() / reset_session()
work the same way as ClaudeCodeProvider, enabling cost tracking in agentpaas.
"qwen" added as a friendly alias for "dashscope".
"""
from __future__ import annotations

import http.client
import json
import os
import urllib.request
from typing import Dict, List

from .base import LLMProvider, ProviderConfig, ProviderError, ChatMessage, ChatResponse


# 瞬时网络断连(dashscope/任意 OpenAI 兼容端)应可重试,而不是把一条长 run 杀掉。
# "Remote end closed connection without response" = http.client.RemoteDisconnected,
# 既非 HTTPError 也非 URLError,会落到兜底 except → 原默认 retryable=False → 一断就死。
_TRANSIENT_NET_MARKERS = (
    "remote end closed", "connection reset", "connection aborted",
    "broken pipe", "timed out", "timeout", "incompleteread",
    "connection refused", "temporarily unavailable", "eof occurred",
)


def _parse_tool_calls(raw) -> list:
    """把 OpenAI 响应的 message.tool_calls 规整成 [{"id","name","arguments"(dict)}]。
    arguments 端点上可能是 dict 或 JSON 字符串(甚至空/坏串)——都兼容。"""
    out = []
    for tc in (raw or []):
        if not isinstance(tc, dict):
            continue
        fn = tc.get("function", {}) or {}
        args = fn.get("arguments")
        if isinstance(args, str):
            try:
                args = json.loads(args) if args.strip() else {}
            except (json.JSONDecodeError, ValueError):
                args = {}
        if not isinstance(args, dict):
            args = {}
        out.append({"id": tc.get("id", ""), "name": fn.get("name", ""), "arguments": args})
    return out


def _is_transient_net_error(exc: Exception) -> bool:
    if isinstance(exc, (ConnectionError, http.client.RemoteDisconnected,
                        http.client.IncompleteRead, TimeoutError)):
        return True
    msg = str(exc).lower()
    return any(m in msg for m in _TRANSIENT_NET_MARKERS)


# Default base URLs and env keys per provider.
# "qwen" is a friendly alias for DashScope (Alibaba Cloud).
_PROVIDER_DEFAULTS = {
    "openai":     ("https://api.openai.com/v1",                              "OPENAI_API_KEY"),
    "ollama":     ("http://127.0.0.1:11434/v1",                              ""),
    "dashscope":  ("https://dashscope.aliyuncs.com/compatible-mode/v1",      "DASHSCOPE_API_KEY"),
    "qwen":       ("https://dashscope.aliyuncs.com/compatible-mode/v1",      "DASHSCOPE_API_KEY"),
    "deepseek":   ("https://api.deepseek.com/v1",                            "DEEPSEEK_API_KEY"),
    "moonshot":   ("https://api.moonshot.cn/v1",                             "MOONSHOT_API_KEY"),
    "zhipu":      ("https://open.bigmodel.cn/api/paas/v4",                   "ZHIPU_API_KEY"),
}


class OpenAICompatProvider(LLMProvider):
    """OpenAI-compatible chat completion API."""

    # L02: Provider configs: name -> (base_url, default_model, api_key_env)
    PRESETS = {
        "openai":    ("https://api.openai.com/v1",                         "gpt-4o",            "OPENAI_API_KEY"),
        "dashscope": ("https://dashscope.aliyuncs.com/compatible-mode/v1", "qwen-max",          "DASHSCOPE_API_KEY"),
        "qwen":      ("https://dashscope.aliyuncs.com/compatible-mode/v1", "qwen-max",          "DASHSCOPE_API_KEY"),
        "deepseek":  ("https://api.deepseek.com/v1",                       "deepseek-chat",     "DEEPSEEK_API_KEY"),
        "moonshot":  ("https://api.moonshot.cn/v1",                        "moonshot-v1-128k",  "MOONSHOT_API_KEY"),
        "ollama":    ("http://127.0.0.1:11434/v1",                         "qwen2.5:7b",        ""),
        "zhipu":     ("https://open.bigmodel.cn/api/paas/v4",              "glm-4",             "ZHIPU_API_KEY"),
    }

    # 各模型真实输入上限(token)。base.py 默认 200000 对 qwen-max 是错的——qwen-max
    # 实际只有 30720(dashscope 报 "Range of input length should be [1, 30720]")。
    # 不设对的话 ConversationLam 的 max_history(min(cw//2,80000))会放行 80000 token
    # 历史 → 累积超限 → 400。按模型名前缀匹配,未知则保持 config 原值。
    _MODEL_CONTEXT_WINDOWS = {
        "qwen-max": 30720, "qwen-plus": 131072, "qwen-turbo": 1000000,
        "qwen-long": 10000000, "qwen2.5": 32768, "qwen2": 32768, "qwen3": 32768,
        "deepseek": 65536, "moonshot-v1-8k": 8192, "moonshot-v1-32k": 32768,
        "moonshot-v1-128k": 128000, "gpt-4o": 128000, "gpt-4-turbo": 128000,
        "gpt-4": 8192, "glm-4": 128000,
    }

    @classmethod
    def _resolve_context_window(cls, model: str):
        m = (model or "").lower()
        # 最长前缀优先(qwen-max-longcontext 命中 qwen-max;但更具体的先匹配)
        best = None
        for key, cw in cls._MODEL_CONTEXT_WINDOWS.items():
            if m.startswith(key) and (best is None or len(key) > len(best[0])):
                best = (key, cw)
        return best[1] if best else None

    def __init__(self, config: ProviderConfig, provider_name: str = "openai",
                 api_key: str = "", base_url: str = "", model: str = ""):
        super().__init__(config)
        self._provider_name = provider_name

        # Resolve base URL and API key
        defaults = _PROVIDER_DEFAULTS.get(provider_name, ("", ""))
        self.base_url = base_url or config.extra.get("base_url", defaults[0])
        env_key = defaults[1]
        self.api_key = api_key or config.extra.get("api_key", "")
        if not self.api_key and env_key:
            self.api_key = os.environ.get(env_key, "")
        if provider_name == "ollama":
            self.api_key = self.api_key or "ollama"

        # L02: Override model from explicit param if provided
        if model:
            self.config.model = model

        # 按真实模型设 context_window(治 qwen-max 被当 200000 → 历史超 30720 → 400)。
        _cw = self._resolve_context_window(self.config.model)
        if _cw:
            self.config.context_window = _cw

        # L03: Token-usage accumulator — mirrors ClaudeCodeProvider.get_usage()
        # so agentpaas cost tracking works for any provider.
        self._usage_input  = 0
        self._usage_output = 0

    def chat(self, messages: List[Dict[str, str]]) -> str:
        url = f"{self.base_url.rstrip('/')}/chat/completions"

        body = json.dumps({
            "model": self.config.model,
            "messages": messages,
            "temperature": self.config.temperature,
            "max_tokens": self.config.max_tokens,
        }, ensure_ascii=False).encode("utf-8")

        headers = {"Content-Type": "application/json"}
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"

        req = urllib.request.Request(url, data=body, headers=headers, method="POST")

        try:
            with urllib.request.urlopen(req, timeout=self.config.timeout) as resp:
                data = json.loads(resp.read())
                # L03: accumulate usage so get_usage() returns session totals
                usage = data.get("usage", {})
                self._usage_input  += usage.get("prompt_tokens", 0)
                self._usage_output += usage.get("completion_tokens", 0)
                return data["choices"][0]["message"]["content"].strip()
        except urllib.error.HTTPError as e:
            error_body = e.read().decode()[:500]
            raise ProviderError(f"{self._provider_name} API error {e.code}: {error_body}",
                                self._provider_name, retryable=e.code >= 500)
        except urllib.error.URLError as e:
            raise ProviderError(f"{self._provider_name} connection error: {e}",
                                self._provider_name, retryable=True)
        except Exception as e:
            # 瞬时网络断连(RemoteDisconnected / reset / timeout 等)判为可重试,
            # 否则一次抖动就杀掉整条 run(实测 dashscope "Remote end closed
            # connection without response" 中断了一条 59 分钟的 run、写残了 paper.tex)。
            raise ProviderError(f"{self._provider_name} error: {e}",
                                self._provider_name,
                                retryable=_is_transient_net_error(e))

    def supports_function_calling(self) -> bool:
        """OpenAI 兼容端默认支持原生 tool_calls。不支持的端点会忽略 tools 字段、
        返回 content,FC loop 有 content 兜底,所以保守返回 True 即可。"""
        return True

    def chat_with_tools(self, messages, tools=None, tool_choice="auto", model="",
                        temperature=None, max_tokens=None) -> dict:
        """原生 function-calling。发送 tools(OpenAI 规格),返回结构化结果:
        {"content": str|None, "tool_calls": [{"id","name","arguments"(dict)}],
         "usage": {...}, "finish_reason": str}。无需解析文本 JSON。"""
        url = f"{self.base_url.rstrip('/')}/chat/completions"
        payload = {
            "model": model or self.config.model,
            "messages": messages,
            "temperature": self.config.temperature if temperature is None else temperature,
            "max_tokens": self.config.max_tokens if max_tokens is None else max_tokens,
        }
        if tools:
            payload["tools"] = tools
            if tool_choice:
                payload["tool_choice"] = tool_choice
        body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
        headers = {"Content-Type": "application/json"}
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"
        req = urllib.request.Request(url, data=body, headers=headers, method="POST")
        try:
            with urllib.request.urlopen(req, timeout=self.config.timeout) as resp:
                data = json.loads(resp.read())
        except urllib.error.HTTPError as e:
            error_body = e.read().decode()[:500]
            raise ProviderError(f"{self._provider_name} API error {e.code}: {error_body}",
                                self._provider_name, retryable=e.code >= 500)
        except urllib.error.URLError as e:
            raise ProviderError(f"{self._provider_name} connection error: {e}",
                                self._provider_name, retryable=True)
        except Exception as e:
            raise ProviderError(f"{self._provider_name} error: {e}",
                                self._provider_name, retryable=_is_transient_net_error(e))
        choice = (data.get("choices") or [{}])[0]
        msg = choice.get("message", {}) or {}
        usage = data.get("usage", {}) or {}
        self._usage_input  += usage.get("prompt_tokens", 0)
        self._usage_output += usage.get("completion_tokens", 0)
        return {
            "content": (msg.get("content") or None),
            "tool_calls": _parse_tool_calls(msg.get("tool_calls")),
            "usage": {"input_tokens": usage.get("prompt_tokens", 0),
                      "output_tokens": usage.get("completion_tokens", 0)},
            "finish_reason": choice.get("finish_reason", ""),
        }

    # ── L03: Session-level usage API (matches ClaudeCodeProvider interface) ──

    def get_usage(self) -> dict:
        """Return accumulated token usage for this session.

        Keys match ClaudeCodeProvider.get_usage() so _execute_agent's cost
        tracking works without any changes. Cache fields are zero (not
        applicable to OpenAI-compatible endpoints).
        """
        return {
            "input_tokens":                 self._usage_input,
            "output_tokens":                self._usage_output,
            "cache_creation_input_tokens":  0,
            "cache_read_input_tokens":      0,
            # Provider tag lets _compute_cost_usd apply the right pricing table
            "provider": self._provider_name,
        }

    def reset_session(self) -> None:
        """Reset per-session usage counters (called between runs)."""
        self._usage_input  = 0
        self._usage_output = 0

    # ── L02: Unified interface overrides ──

    def chat_typed(self, messages: List[ChatMessage],
                   model: str = "", temperature: float = 0.0,
                   max_tokens: int = 4096) -> ChatResponse:
        """Typed chat with response metadata via urllib (no openai SDK required)."""
        url = f"{self.base_url.rstrip('/')}/chat/completions"
        api_msgs = [{"role": m.role, "content": m.content} for m in messages]

        body = json.dumps({
            "model": model or self.default_model,
            "messages": api_msgs,
            "temperature": temperature,
            "max_tokens": max_tokens,
        }, ensure_ascii=False).encode("utf-8")

        headers = {"Content-Type": "application/json"}
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"

        req = urllib.request.Request(url, data=body, headers=headers, method="POST")

        try:
            with urllib.request.urlopen(req, timeout=self.config.timeout) as resp:
                data = json.loads(resp.read())
                choice = data["choices"][0]
                usage = data.get("usage", {})
                # Accumulate session usage so get_usage() is consistent whether
                # the caller used chat() or chat_typed() (cost tracking parity).
                self._usage_input  += usage.get("prompt_tokens", 0)
                self._usage_output += usage.get("completion_tokens", 0)
                return ChatResponse(
                    text=(choice.get("message", {}).get("content") or "").strip(),
                    input_tokens=usage.get("prompt_tokens", 0),
                    output_tokens=usage.get("completion_tokens", 0),
                    model=data.get("model", ""),
                    finish_reason=choice.get("finish_reason", ""),
                )
        except urllib.error.HTTPError as e:
            error_body = e.read().decode()[:500]
            raise ProviderError(f"{self._provider_name} API error {e.code}: {error_body}",
                                self._provider_name, retryable=e.code >= 500)
        except urllib.error.URLError as e:
            raise ProviderError(f"{self._provider_name} connection error: {e}",
                                self._provider_name, retryable=True)
        except Exception as e:
            # 瞬时网络断连(RemoteDisconnected / reset / timeout 等)判为可重试,
            # 否则一次抖动就杀掉整条 run(实测 dashscope "Remote end closed
            # connection without response" 中断了一条 59 分钟的 run、写残了 paper.tex)。
            raise ProviderError(f"{self._provider_name} error: {e}",
                                self._provider_name,
                                retryable=_is_transient_net_error(e))

    async def achat(self, messages: List[ChatMessage],
                    model: str = "", temperature: float = 0.0,
                    max_tokens: int = 4096) -> ChatResponse:
        import asyncio
        return await asyncio.get_event_loop().run_in_executor(
            None, lambda: self.chat_typed(messages, model, temperature, max_tokens)
        )

    @property
    def provider_name(self) -> str:
        return self._provider_name

    @property
    def default_model(self) -> str:
        return self.config.model
