from __future__ import annotations
from typing import Any, Optional, Dict, Tuple
import json


def oa_messages_to_gemini(messages: list[dict]) -> tuple[Optional[dict], list[dict]]:
    system_texts: list[str] = []
    contents: list[dict] = []

    for m in messages:
        role = m.get("role")
        content = m.get("content", "")

        # 兼容 content 为 list（OpenAI 多模态结构），这里只取 text
        if isinstance(content, list):
            texts = []
            for part in content:
                if isinstance(part, dict) and part.get("type") == "text":
                    texts.append(part.get("text", ""))
            content = "\n".join([t for t in texts if t])

        content = "" if content is None else str(content)

        if role == "system":
            if content.strip():
                system_texts.append(content)
            continue

        g_role = "model" if role == "assistant" else "user"
        contents.append({"role": g_role, "parts": [{"text": content}]})

    system_instruction = None
    if system_texts:
        system_instruction = {"parts": [{"text": "\n".join(system_texts)}]}
    return system_instruction, contents


def gemini_text(resp_json: dict) -> str:
    cands = resp_json.get("candidates") or []
    if not cands:
        return ""
    content = (cands[0].get("content") or {})
    parts = content.get("parts") or []
    return "".join([p.get("text", "") for p in parts if isinstance(p, dict)])


def _usage_norm_from_gemini(resp_json: dict) -> dict:
    """
    Gemini usage 主要在 usageMetadata:
      {
        "usageMetadata": {
          "promptTokenCount": ...,
          "candidatesTokenCount": ...,
          "totalTokenCount": ...,
          "cachedContentTokenCount": ... (若命中 cachedContent/缓存)
        }
      }
    """
    um = resp_json.get("usageMetadata") or resp_json.get("usage_metadata") or {}
    prompt = um.get("promptTokenCount")
    completion = um.get("candidatesTokenCount")
    cached = um.get("cachedContentTokenCount") or um.get("cachedTokenCount")  # 兼容字段差异
    total = um.get("totalTokenCount")

    norm = {
        "prompt_tokens": int(prompt) if isinstance(prompt, int) else None,
        "completion_tokens": int(completion) if isinstance(completion, int) else None,
        "cached_prompt_tokens": int(cached) if isinstance(cached, int) else 0,
        "total_tokens": int(total) if isinstance(total, int) else None,
        "vendor": "gemini",
    }
    return norm


def gemini_to_openai_like(resp_json: dict) -> dict:
    text = gemini_text(resp_json)
    return {
        "object": "chat.completion",
        "choices": [{
            "index": 0,
            "message": {"role": "assistant", "content": text},
            "finish_reason": "stop",
        }],
        "usage_norm": _usage_norm_from_gemini(resp_json),
        "raw": resp_json,
    }


def extract_cached_content_ref(body: dict) -> Optional[str]:
    """
    兼容多种入口：
    - body["cached_content"]
    - body["extra_body"]["cached_content"]
    - body["vendor_extra"]["cached_content"]  (flatten 后也会提升)
    - body["cachedContent"] / body["cached_content_name"]（兼容业务命名）
    """
    for k in ("cached_content", "cachedContent", "cached_content_name", "cachedContentName"):
        v = body.get(k)
        if isinstance(v, str) and v.strip():
            return v.strip()

    extra = body.get("extra_body")
    if isinstance(extra, dict):
        v = extra.get("cached_content") or extra.get("cachedContent")
        if isinstance(v, str) and v.strip():
            return v.strip()

    ve = body.get("vendor_extra")
    if isinstance(ve, dict):
        v = ve.get("cached_content") or ve.get("cachedContent")
        if isinstance(v, str) and v.strip():
            return v.strip()

    return None


def build_gemini_payload(body: dict) -> dict:
    """
    构建 Gemini generateContent payload。

    缓存策略约定（重要）：
    - 若请求体带 cached_content（cachedContents/xxx），则 payload 顶层带 cachedContent 引用。
    - 此时 **仍可发送 contents（变化部分）**，但建议不要再重复超长 systemInstruction；
      你应当在创建 cachedContents 时把 systemInstruction + 稳定前缀 contents 一并放入缓存。
    - 这里的实现：如果命中 cachedContent，则默认不重复 systemInstruction（除非调用方显式传 force_system_instruction=True）。
    """
    sys_inst, contents = oa_messages_to_gemini(body.get("messages") or [])
    payload: dict[str, Any] = {"contents": contents}

    cached_ref = extract_cached_content_ref(body)
    if cached_ref:
        payload["cachedContent"] = cached_ref

    force_sys = bool(body.get("force_system_instruction"))
    if sys_inst and (force_sys or not cached_ref):
        payload["systemInstruction"] = sys_inst

    gen_cfg = {}
    if "temperature" in body:
        gen_cfg["temperature"] = body["temperature"]
    if "max_tokens" in body:
        gen_cfg["maxOutputTokens"] = body["max_tokens"]

    if body.get("response_mime_type"):
        gen_cfg["responseMimeType"] = body["response_mime_type"]

    schema = body.get("response_schema_json") or body.get("response_schema")
    if schema:
        gen_cfg["responseSchema"] = schema

    if gen_cfg:
        payload["generationConfig"] = gen_cfg

    # 工具/函数等：保持扩展口（透传）
    if isinstance(body.get("tools"), list):
        payload["tools"] = body["tools"]
    if isinstance(body.get("tool_config"), dict):
        payload["toolConfig"] = body["tool_config"]

    return payload
