# -*- coding: utf-8 -*-
"""四家模型接口的协议整形：只算，不发。

为什么要单独一层：同一套协议现在有两个传输方在用——桌面端的同步 requests
（ai_chat.py）和网关的异步 httpx（model_gateway.py）。如果两边各写一份"怎么拼
Claude 的 payload"，迟早会漂：一边修了 system 的位置另一边没修，线上表现就是
"网关能用、桌面端 400"，而且极难查。

这里的函数全是纯函数：输入配置和消息，输出该发什么、怎么解析回来的东西。没有
网络、没有全局状态、没有副作用，所以两边都能安全复用，也好测。

四家的形状差异（照搬 OpenAI 的 payload 到别家一定报错）：

  OpenAI  /chat/completions   system 在 messages 里    Authorization: Bearer
  Claude  /v1/messages        system 是顶层参数        x-api-key + anthropic-version
                              max_tokens 必填
  Dify    /v1/chat-messages   system 在应用侧配置      Authorization: Bearer
                              图片要先 upload 再引用
  ComfyUI 文生图工作流引擎，不是聊天接口——不参与对话协议
"""

from __future__ import annotations

from urllib.parse import urlparse

# Anthropic Messages API 的必填版本头
ANTHROPIC_VERSION = "2023-06-01"

CHAT_KINDS = ("openai", "claude", "dify")
ALL_KINDS = CHAT_KINDS + ("comfyui",)

_BROWSER_UA = (
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
    "AppleWebKit/537.36 (KHTML, like Gecko) "
    "Chrome/124.0.0.0 Safari/537.36"
)


def detect_kind(kind: str, base_url: str) -> str:
    """声明了就用声明的；没声明就从地址猜。"""
    value = str(kind or "auto").strip().lower()
    if value in ALL_KINDS:
        return value
    parsed = urlparse((base_url or "").rstrip("/"))
    path = (parsed.path or "").lower()
    if "chat-messages" in path or "completion-messages" in path:
        return "dify"
    if "anthropic" in (parsed.hostname or ""):
        return "claude"
    return "openai"


ENDPOINT_MODES = ("auto", "exact")


def endpoint_url(kind: str, base_url: str, mode: str = "auto") -> str:
    """拼出实际请求地址。

    `mode="auto"`（默认）按接口类型补全路径：填 `https://api.openai.com/v1`
    就补成 `.../v1/chat/completions`。绝大多数服务商都是这个形状，所以它是默认。

    `mode="exact"` 原样使用，一个字符都不加。这是为那些路径不按套路来的服务准备
    的——比如 `https://api.example.com/custom/llm/invoke`，auto 会把它拼成
    `.../invoke/chat/completions`，请求发到一个不存在的地址，报 404，而排查的人
    会去怀疑密钥和网络。

    为什么用显式开关而不是"更聪明的猜测"：`https://api.example.com/openai` 到底
    是一个前缀（还要补 /chat/completions）还是完整端点，光看地址分不出来。猜错
    的两种方向都会静默地把请求发歪，不如让填的人说清楚。
    """
    base = (base_url or "").rstrip("/")
    if str(mode or "auto").strip().lower() == "exact":
        return base
    path = (urlparse(base).path or "").lower().rstrip("/")
    kind = detect_kind(kind, base_url)
    if kind == "claude":
        if path.endswith("/messages"):
            return base
        return f"{base}/messages" if path.endswith("/v1") else f"{base}/v1/messages"
    if kind == "dify":
        if path.endswith(("/chat-messages", "/completion-messages")):
            return base
        return f"{base}/chat-messages" if path.endswith("/v1") else f"{base}/v1/chat-messages"
    if path.endswith("/chat/completions"):
        return base
    if path.startswith("/v1/") and len(path) > len("/v1/"):
        return base
    return f"{base}/chat/completions"


def dify_api_root(base_url: str, mode: str = "auto") -> str:
    """Dify App API 根地址，用来拼 /files/upload。"""
    endpoint = endpoint_url("dify", base_url, mode).rstrip("/")
    for suffix in ("/chat-messages", "/completion-messages"):
        if endpoint.lower().endswith(suffix):
            return endpoint[: -len(suffix)]
    return endpoint


def auth_headers(kind: str, api_key: str, base_url: str = "") -> dict:
    """鉴权头。Claude 用 x-api-key，其余用 Bearer。"""
    key = str(api_key or "").strip()
    if detect_kind(kind, base_url) == "claude":
        auth = {"x-api-key": key, "anthropic-version": ANTHROPIC_VERSION}
    else:
        auth = {"Authorization": f"Bearer {key}"}
    return {
        **auth,
        "Content-Type": "application/json",
        "Accept": "application/json, text/plain, */*",
        "Accept-Language": "zh-CN,zh;q=0.9,en;q=0.8",
        "User-Agent": _BROWSER_UA,
        "Connection": "keep-alive",
    }


def claude_messages(messages: list) -> tuple[str, list]:
    """OpenAI 形状的 messages → Anthropic 要的 (system, messages)。

    三处不兼容，照搬一定 400：system 是顶层参数、相邻同角色必须合并、
    第一条必须是 user。
    """
    system_parts: list[str] = []
    turns: list[dict] = []
    for item in messages or []:
        role = str((item or {}).get("role") or "user")
        content = (item or {}).get("content")
        if role == "system":
            system_parts.append(content if isinstance(content, str) else str(content))
            continue
        role = "assistant" if role == "assistant" else "user"
        if isinstance(content, str):
            blocks = [{"type": "text", "text": content}]
        elif isinstance(content, list):
            blocks = []
            for part in content:
                if not isinstance(part, dict):
                    blocks.append({"type": "text", "text": str(part)})
                elif part.get("type") == "text":
                    blocks.append({"type": "text", "text": str(part.get("text") or "")})
                elif part.get("type") == "image_url":
                    url = str((part.get("image_url") or {}).get("url") or "")
                    if url.startswith("data:"):
                        head, _, payload = url.partition(",")
                        media = head[5:].split(";")[0] or "image/png"
                        blocks.append({
                            "type": "image",
                            "source": {
                                "type": "base64",
                                "media_type": media,
                                "data": payload,
                            },
                        })
        else:
            blocks = [{"type": "text", "text": str(content)}]
        if turns and turns[-1]["role"] == role:
            turns[-1]["content"].extend(blocks)
        else:
            turns.append({"role": role, "content": blocks})
    while turns and turns[0]["role"] != "user":
        turns.pop(0)
    if not turns:
        turns = [{"role": "user", "content": [{"type": "text", "text": "请回复"}]}]
    return "\n\n".join(part for part in system_parts if part), turns


def _last_user_text(messages: list) -> str:
    for item in reversed(messages or []):
        if str((item or {}).get("role") or "") != "user":
            continue
        content = (item or {}).get("content")
        if isinstance(content, str):
            return content
        if isinstance(content, list):
            return "".join(
                str(part.get("text") or "")
                for part in content
                if isinstance(part, dict) and part.get("type") == "text"
            )
        return str(content)
    return ""


def chat_payload(
    kind: str,
    *,
    model: str,
    messages: list,
    max_tokens: int,
    temperature: float,
    base_url: str = "",
    tools: list | None = None,
    dify_user: str = "wechat-rpa",
    dify_files: list | None = None,
    dify_conversation_id: str = "",
) -> dict:
    """按各家的形状拼请求体。"""
    kind = detect_kind(kind, base_url)
    if kind == "claude":
        system_text, turns = claude_messages(messages)
        payload = {
            "model": model,
            "messages": turns,
            # Anthropic 的 max_tokens 是必填且必须 >= 1，给 0 会直接 400
            "max_tokens": max(1, int(max_tokens or 500)),
            "temperature": float(temperature),
        }
        if system_text:
            payload["system"] = system_text
        return payload
    if kind == "dify":
        payload = {
            "inputs": {},
            "query": _last_user_text(messages) or "请回复",
            "response_mode": "blocking",
            "user": dify_user or "wechat-rpa",
        }
        if dify_conversation_id:
            payload["conversation_id"] = dify_conversation_id
        if dify_files:
            payload["files"] = list(dify_files)
        return payload
    payload = {
        "model": model,
        "messages": messages,
        "max_tokens": int(max_tokens or 500),
        "temperature": float(temperature),
    }
    if tools:
        payload["tools"] = tools
        payload["tool_choice"] = "auto"
    return payload


def parse_chat(kind: str, data: dict, base_url: str = "") -> str:
    """从各家的响应里取出回复文本。取不到就抛，绝不返回空串蒙混过去。"""
    kind = detect_kind(kind, base_url)
    data = data if isinstance(data, dict) else {}
    if kind == "claude":
        parts = [
            str(block.get("text") or "")
            for block in (data.get("content") or [])
            if isinstance(block, dict) and block.get("type") == "text"
        ]
        text = "".join(parts).strip()
        if not text:
            raise ValueError("Claude 未返回文本内容")
        return text
    if kind == "dify":
        answer = data.get("answer")
        if answer is None:
            # 少数部署会包一层 data
            answer = (data.get("data") or {}).get("answer")
        if not answer:
            raise ValueError("Dify 未返回 answer")
        return str(answer)
    try:
        content = data["choices"][0]["message"]["content"]
    except (KeyError, IndexError, TypeError) as exc:
        raise ValueError("OpenAI 兼容响应里没有 choices[0].message.content") from exc
    if content is None:
        raise ValueError("OpenAI 兼容响应的 content 为空")
    return str(content)


def parse_message(kind: str, data: dict, base_url: str = "") -> dict:
    """取出完整的 assistant 消息，而不只是文本。

    `parse_chat` 只回文本，工具调用会被整个丢掉——模型说"我要查一下客户资料"，
    到调用方手里变成一句空回复。MCP 多轮必须看到 `tool_calls` 才能往下走。

    只有 OpenAI 兼容协议有工具调用。Dify 的工具在应用侧编排、对我们不可见，
    Claude 的 tool_use 是另一套形状——这两种都只回文本，让调用方按"没有工具
    调用"处理，比硬凑一个形状再在别处出错要好。
    """
    kind = detect_kind(kind, base_url)
    data = data if isinstance(data, dict) else {}
    if kind in ("claude", "dify"):
        return {"content": parse_chat(kind, data, base_url), "tool_calls": []}
    try:
        message = data["choices"][0]["message"]
    except (KeyError, IndexError, TypeError) as exc:
        raise ValueError("OpenAI 兼容响应里没有 choices[0].message") from exc
    if not isinstance(message, dict):
        raise ValueError("OpenAI 兼容响应的 message 不是对象")
    tool_calls = message.get("tool_calls") or []
    content = message.get("content")
    if content is None and not tool_calls:
        raise ValueError("OpenAI 兼容响应既没有 content 也没有 tool_calls")
    return {"content": "" if content is None else str(content), "tool_calls": list(tool_calls)}


def image_message(prompt: str, image_b64: str, mime: str = "image/png") -> dict:
    """一条带图的 user 消息（OpenAI 形状）。Claude 的转换由 claude_messages 负责。"""
    return {
        "role": "user",
        "content": [
            {"type": "text", "text": str(prompt or "")},
            {
                "type": "image_url",
                "image_url": {"url": f"data:{mime};base64,{image_b64}"},
            },
        ],
    }