"""Shared, standards-based LLM client for web chat + IM bot.

平台红线(07-24)：模型层以成熟/开源标准为设计基准。本模块使用官方 OpenAI SDK
(\"openai.AsyncOpenAI\")——它是 OpenAI 兼容端点的事实标准客户端。平台本身已在讲
OpenAI 协议(\"/v1/chat/completions\")，因此零协议改造成本，并天然获得：
  * 原生 SSE 流式(\"stream=True\")；
  * 结构化的 tool_calls / usage 解析；
  * 统一的代理(trust_env=False + 显式 proxy)与超时处理。

本模块取代 chat_service.py 与 im_llm.py 中两套手写 httpx 循环，使 web chat 与
IM bot 共用**同一模型层**、同一套 failover/代理/超时语义——这是「各端能力统一对齐」
红线在模型层的落地。

API:
  achat_complete(api_url, api_key, request_body, *, proxy=None, timeout=120, stream=False)
    - stream=False -> dict，形状对齐上游 \"data\" 响应：
        {"choices":[{"message":{"content","tool_calls","reasoning_content"}}], "usage":{...}}
    - stream=True  -> async generator，逐块产出 ("content"|"reasoning", text) 增量
"""
from __future__ import annotations

import httpx
import asyncio
from openai import AsyncOpenAI, OpenAI


def _base_url(api_url: str) -> str:
    """从完整 /v1/chat/completions 或 /v1/embeddings 地址推导 OpenAI SDK 所需的 base_url。"""
    u = (api_url or "").rstrip()
    for suf in ("/v1/chat/completions", "/chat/completions", "/v1/embeddings", "/embeddings"):
        if u.endswith(suf):
            return u[: -len(suf)]
    return u


def build_openai_client(api_url, api_key, *, proxy=None, timeout=120, auth_header=None) -> AsyncOpenAI:
    base = _base_url(api_url)
    # 始终 trust_env=False：避免容器 HTTP_PROXY 对国内端点(siliconflow 等)造成 502；
    # 代理经 proxy 显式传入（AGNES 等），未传则直连。
    http_client = httpx.AsyncClient(proxy=proxy, trust_env=False)
    # 故障快速失败（全栈统一，覆盖 web chat / IM bot / embedding）：
    #   connect 封顶 10s —— 坏端点不再让请求挂死事件循环；
    #   read/write 封顶 30s —— 长流式回复只要持续有分块即不受影响；
    #   保留数值型 timeout 接口，旧调用方无需改动。
    _num = float(timeout) if isinstance(timeout, (int, float)) else 120.0
    _to = httpx.Timeout(
        connect=min(10.0, _num),
        read=min(30.0, _num),
        write=min(30.0, _num),
        pool=10.0,
    )
    _hdrs = {}
    if auth_header and ":" in auth_header:
        _k, _v = auth_header.split(":", 1)
        _hdrs[_k.strip()] = _v.strip()
    elif auth_header:
        _hdrs[auth_header.strip()] = api_key

    return AsyncOpenAI(
        base_url=base,
        api_key=api_key or "EMPTY",
        timeout=_to,
        http_client=http_client,
        default_headers=_hdrs if _hdrs else None,
    )


def _create_kwargs(request_body: dict, stream: bool = False) -> dict:
    kw = {
        "model": request_body.get("model"),
        "messages": request_body.get("messages", []),
        "stream": stream,
    }
    if request_body.get("temperature") is not None:
        kw["temperature"] = request_body["temperature"]
    if request_body.get("top_p") is not None:
        kw["top_p"] = request_body["top_p"]
    if request_body.get("max_tokens") is not None:
        kw["max_tokens"] = request_body["max_tokens"]
    if request_body.get("stop") is not None:
        kw["stop"] = request_body["stop"]
    tools = request_body.get("tools")
    if tools:
        kw["tools"] = tools
        kw["tool_choice"] = request_body.get("tool_choice", "auto")
    # 2026-07-30: 将未在白名单的额外字段（如 enable_thinking）作为 extra_body 传递
    _whitelist = {"model", "messages", "stream", "temperature", "top_p", "max_tokens", "stop", "tools", "tool_choice"}
    _extras = {k: v for k, v in request_body.items() if k not in _whitelist}
    if _extras:
        kw["extra_body"] = _extras
    return kw


def _normalize_tool_calls(tcs):
    if not tcs:
        return None
    out = []
    for tc in tcs:
        out.append(
            {
                "id": getattr(tc, "id", None),
                "type": "function",
                "function": {
                    "name": getattr(tc.function, "name", None),
                    "arguments": getattr(tc.function, "arguments", "") or "",
                },
            }
        )
    return out


async def achat_complete(api_url, api_key, request_body, *, proxy=None, timeout=120, stream=False, auth_header=None):
    """统一的 OpenAI 兼容聊天补全。详见模块 docstring。"""
    if stream:

        async def _gen():
            client = build_openai_client(api_url, api_key, proxy=proxy, timeout=timeout, auth_header=auth_header)
            _content_parts = []
            _reasoning_parts = []
            _tool_acc = {}  # index -> {id, type, name, arguments}
            try:
                _to = min(30.0, float(timeout) if isinstance(timeout, (int, float)) else 120.0)
                gen = await asyncio.wait_for(
                    client.chat.completions.create(
                        **_create_kwargs(request_body, stream=True)
                    ),
                    timeout=_to,
                )
                _it = gen.__aiter__()
                while True:
                    try:
                        chunk = await asyncio.wait_for(_it.__anext__(), timeout=_to)
                    except StopAsyncIteration:
                        break
                    except asyncio.TimeoutError:
                        raise TimeoutError("stream chunk read timeout")
                    if not chunk.choices:
                        continue
                    delta = chunk.choices[0].delta
                    rc = getattr(delta, "reasoning_content", None)
                    if rc:
                        _reasoning_parts.append(rc)
                        yield ("reasoning", rc)
                    ct = getattr(delta, "content", None)
                    if ct:
                        _content_parts.append(ct)
                        yield ("content", ct)
                    # 累积 OpenAI 流式 tool_calls（供后续工具循环执行）
                    tcs = getattr(delta, "tool_calls", None)
                    if tcs:
                        for tc in tcs:
                            _idx = getattr(tc, "index", 0)
                            _slot = _tool_acc.setdefault(
                                _idx, {"id": None, "type": "function", "name": "", "arguments": ""}
                            )
                            _tcid = getattr(tc, "id", None)
                            if _tcid:
                                _slot["id"] = _tcid
                            _fn = getattr(tc, "function", None)
                            if _fn is not None:
                                _fn_name = getattr(_fn, "name", None)
                                if _fn_name:
                                    _slot["name"] += _fn_name
                                _fn_args = getattr(_fn, "arguments", None)
                                if _fn_args:
                                    _slot["arguments"] += _fn_args
            finally:
                try:
                    await asyncio.wait_for(client.close(), timeout=5)
                except Exception:
                    pass
            # 流式结束时吐出已组装的完整 message（含 tool_calls），供调用方无缝衔接工具循环
            _tool_calls = None
            if _tool_acc:
                _tool_calls = []
                for _i in sorted(_tool_acc.keys()):
                    _s = _tool_acc[_i]
                    _tool_calls.append({
                        "id": _s["id"],
                        "type": "function",
                        "function": {"name": _s["name"], "arguments": _s["arguments"]},
                    })
            yield ("message", {
                "content": "".join(_content_parts),
                "tool_calls": _tool_calls,
                "reasoning_content": "".join(_reasoning_parts) or None,
            })

        return _gen()

    client = build_openai_client(api_url, api_key, proxy=proxy, timeout=timeout, auth_header=auth_header)
    try:
        _to = min(30.0, float(timeout) if isinstance(timeout, (int, float)) else 120.0)
        resp = await asyncio.wait_for(
            client.chat.completions.create(**_create_kwargs(request_body)),
            timeout=_to,
        )
        msg = resp.choices[0].message
        return {
            "choices": [
                {
                    "message": {
                        "content": msg.content or "",
                        "tool_calls": _normalize_tool_calls(getattr(msg, "tool_calls", None)),
                        "reasoning_content": getattr(msg, "reasoning_content", None),
                    }
                }
            ],
            "usage": (resp.usage.model_dump() if resp.usage else {}),
        }
    finally:
        try:
            await asyncio.wait_for(client.close(), timeout=5)
        except Exception:
            pass


def build_openai_client_sync(api_url, api_key, *, proxy=None, timeout=120, auth_header=None) -> OpenAI:
    """同步客户端工厂（供 IM bot 等同步上下文复用同一标准模型层）。"""
    base = _base_url(api_url)
    http_client = httpx.Client(proxy=proxy, trust_env=False)
    _hdrs = {}
    if auth_header and ":" in auth_header:
        _k, _v = auth_header.split(":", 1)
        _hdrs[_k.strip()] = _v.strip()
    elif auth_header:
        _hdrs[auth_header.strip()] = api_key
    return OpenAI(base_url=base, api_key=api_key or "EMPTY", timeout=timeout, http_client=http_client, default_headers=_hdrs if _hdrs else None)


def chat_complete(api_url, api_key, request_body, *, proxy=None, timeout=120, auth_header=None):
    """同步版（与 achat_complete 行为/返回形状完全一致），供非 async 上下文使用。

    2026-07-30 重写：直接使用 requests 替代 OpenAI SDK/httpx。
    OpenAI SDK 的 httpx 连接池在慢长轮询（如 iLink QR 状态 12-19s）后会出现
    Connection error，已实测复现。requests 库没有这个问题。
    """
    import requests
    import base64
    import random
    body = dict(request_body)
    body["stream"] = False
    # enable_thinking 等额外字段已在 request_body 中

    # Headers: Bearer token + X-WECHAT-UIN (兼容 iLink)
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}" if api_key else "Bearer empty",
    }
    if auth_header and ":" in auth_header:
        _k, _v = auth_header.split(":", 1)
        headers[_k.strip()] = _v.strip()
    else:
        # X-WECHAT-UIN 防重放
        headers["X-WECHAT-UIN"] = base64.b64encode(str(random.randint(0, 2**32-1)).encode()).decode()

    # 2026-07-31 修复根因: proxy 为 None 时传 proxies={} 而非 None。
    # requests 在 proxies=None 时会继续读取环境 HTTP_PROXY(trust_env 默认 True),
    # 而 uvicorn 进程(run.sh)带有 HTTP_PROXY=172.18.0.1:7890, 该代理无法解析
    # Docker 内部域名(如 qihu-aiflow-litellm), 导致所有内部模型调用静默失败、返回空。
    # 显式 proxies={} + trust_env=False 可完全绕过环境代理, 使 _proxy_for_url 的
    # 决策(返回 None=不代理)真正生效。
    _proxies = {"http": proxy, "https": proxy} if proxy else {}
    _session = requests.Session()
    _session.trust_env = False
    last_body = {}
    for attempt in range(2):
        try:
            resp = _session.post(
                api_url.rstrip("/"),
                headers=headers,
                json=body,
                proxies=_proxies,
                timeout=timeout,
            )
            try:
                last_body = resp.json() if resp.content else {}
            except Exception:
                last_body = {}
            if resp.ok:
                import logging as _l6
                _l6.warning(f'[chat_complete] body keys={list(last_body.keys()) if isinstance(last_body, dict) else type(last_body)}, content_len={len(resp.content)}, content_preview={str(resp.content)[:200]}')
                return _shape_response(last_body)
        except Exception:
            pass
    return {"choices": [{"message": {"content": "", "tool_calls": None, "reasoning_content": None}}]}


def _shape_response(body: dict) -> dict:
    """Normalize OpenAI-compatible response to standard shape."""
    choices = body.get("choices", [])
    import logging as _l7
    if choices:
        msg = choices[0].get("message", {})
        result = {
            "choices": [{
                "message": {
                    "content": msg.get("content") or "",
                    "tool_calls": _normalize_tool_calls(msg.get("tool_calls")),
                    "reasoning_content": msg.get("reasoning_content") or "",
                }
            }]
        }
        _l7.warning(f'[_shape] content_len={len(result["choices"][0]["message"]["content"])}, keys={list(result["choices"][0]["message"].keys())}')
        return result
    _l7.warning(f'[_shape] NO choices in body, keys={list(body.keys())[:5]}')
    return body


def _embed_kwargs(input, model):
    return {"model": model, "input": input if isinstance(input, list) else [input]}


async def aembed(api_url, api_key, input, model, *, proxy=None, timeout=60) -> list:
    """统一的 OpenAI 兼容文本嵌入（异步）。返回 list[list[float]]。

    供 pgvector_search.get_embedding 等 embeddings 路径复用，把散落的裸 httpx
    嵌入调用收敛到同一标准模型层（与聊天共用代理/超时/failover 语义）。
    """
    client = build_openai_client(api_url, api_key, proxy=proxy, timeout=timeout, auth_header=auth_header)
    try:
        _to = min(60.0, float(timeout) if isinstance(timeout, (int, float)) else 60.0)
        resp = await asyncio.wait_for(
            client.embeddings.create(**_embed_kwargs(input, model)),
            timeout=_to,
        )
        return [d.embedding for d in resp.data]
    finally:
        try:
            await asyncio.wait_for(client.close(), timeout=5)
        except Exception:
            pass


def embed(api_url, api_key, input, model, *, proxy=None, timeout=60) -> list:
    """同步版嵌入（与非 async 上下文复用同一标准模型层）。"""
    client = build_openai_client_sync(api_url, api_key, proxy=proxy, timeout=timeout, auth_header=auth_header)
    with client:
        resp = client.embeddings.create(**_embed_kwargs(input, model))
        return [d.embedding for d in resp.data]