from __future__ import annotations

from typing import Any

from niannian.llm.providers.registry import DEFAULT_LLM_BASE_URLS

from niannian.llm.model_registry import SILICONFLOW_FREE_MODELS, ZHIPU_FREE_MODELS
from niannian.llm.custom_providers import enrich_catalog_with_custom, load_providers_doc

# 国内 LLM 厂商目录（Console 默认展示；openai_compat 为硅基历史别名）
DOMESTIC_LLM_CATALOG: tuple[dict[str, Any], ...] = (
    {
        "id": "siliconflow",
        "label": "硅基流动",
        "aliases": ["openai_compat"],
        "base_url": "https://api.siliconflow.cn/v1",
        "docs_url": "https://cloud.siliconflow.cn",
        "models": list(SILICONFLOW_FREE_MODELS),
    },
    {
        "id": "deepseek",
        "label": "DeepSeek",
        "base_url": "https://api.deepseek.com/v1",
        "docs_url": "https://platform.deepseek.com",
        "models": ["deepseek-chat", "deepseek-reasoner"],
    },
    {
        "id": "custom",
        "label": "自定义 API",
        "base_url": "",
        "docs_url": None,
        "models": [],
        "builtin_custom": True,
    },
    {
        "id": "zhipu",
        "label": "智谱 AI",
        "base_url": "https://open.bigmodel.cn/api/paas/v4",
        "docs_url": "https://open.bigmodel.cn",
        "models": list(ZHIPU_FREE_MODELS),
        "exclusive": True,
    },
    {
        "id": "volcengine",
        "label": "火山引擎 / 豆包",
        "base_url": "https://ark.cn-beijing.volces.com/api/v3",
        "docs_url": "https://console.volcengine.com/ark",
        "models": [
            "doubao-seed-evolving",
            "doubao-seed-2-1-turbo-260628",
            "doubao-seed-2-0-mini-260215",
        ],
    },
    {
        "id": "dashscope",
        "label": "阿里百炼 / 通义",
        "base_url": "https://dashscope.aliyuncs.com/compatible-mode/v1",
        "docs_url": "https://bailian.console.aliyun.com",
        "models": ["qwen-plus", "qwen-turbo", "qwen-max", "text-embedding-v3"],
    },
    {
        "id": "minimax",
        "label": "MiniMax",
        "base_url": "https://api.minimax.chat/v1",
        "docs_url": "https://platform.minimaxi.com",
        "models": ["abab6.5s-chat", "abab6.5-chat"],
    },
    {
        "id": "baichuan",
        "label": "百川智能",
        "base_url": "https://api.baichuan-ai.com/v1",
        "docs_url": "https://platform.baichuan-ai.com",
        "models": ["Baichuan4-Turbo", "Baichuan3-Turbo"],
    },
    {
        "id": "moonshot",
        "label": "月之暗面 Kimi",
        "base_url": "https://api.moonshot.cn/v1",
        "docs_url": "https://platform.moonshot.cn",
        "models": ["moonshot-v1-8k", "moonshot-v1-32k", "moonshot-v1-128k"],
    },
    {
        "id": "qianfan",
        "label": "百度千帆",
        "base_url": "https://qianfan.baidubce.com/v2",
        "docs_url": "https://console.bce.baidu.com/qianfan",
        "models": ["ernie-4.0-8k", "ernie-3.5-8k"],
    },
    {
        "id": "spark",
        "label": "讯飞星火 / 星辰 MaaS",
        "base_url": "https://maas-api.cn-huabei-1.xf-yun.com/v2",
        "docs_url": "https://maas.xfyun.cn",
        "models": [
            "xopqwen36v35b",
            "xopqwen35v35b",
            "xop3qwen1b7",
            "xop35qwen2b",
            "xophunyuan7bmt",
        ],
    },
)

_BRAIN_MODULES: tuple[tuple[str, str, str | None], ...] = (
    ("right", "小思 · 表达/推理", None),
    ("left", "小忆 · 打标", None),
    ("hippocampus", "小海 · 巩固", None),
    ("cerebellum", "小织 · 向量编码", "embed"),
    ("retrieval_rerank", "小织 · 向量重排", "rerank"),
    ("ocr", "小路 · 视觉 OCR", "vision"),
    ("brainstem", "小枢 · 编排", None),
)

_SLOT_FIELDS = ("primary", "module_fallback", "global_fallback")


def list_catalog(providers_doc: dict | None = None) -> list[dict[str, Any]]:
    base: list[dict[str, Any]] = []
    for entry in DOMESTIC_LLM_CATALOG:
        pid = str(entry["id"])
        url = str(entry.get("base_url") or DEFAULT_LLM_BASE_URLS.get(pid, ""))
        base.append({**entry, "base_url": url.rstrip("/")})
    doc = providers_doc if providers_doc is not None else load_providers_doc()
    return enrich_catalog_with_custom(base, doc)


def catalog_by_id(providers_doc: dict | None = None) -> dict[str, dict[str, Any]]:
    m: dict[str, dict[str, Any]] = {}
    for item in list_catalog(providers_doc):
        m[str(item["id"])] = item
        for alias in item.get("aliases") or []:
            m[str(alias)] = item
    return m


def brain_modules() -> list[dict[str, str]]:
    out: list[dict[str, str]] = []
    for mid, label, capability in _BRAIN_MODULES:
        item: dict[str, str] = {"id": mid, "label": label}
        if capability:
            item["capability"] = capability
        if mid in ("cerebellum", "retrieval_rerank"):
            item["parent_brain"] = "cerebellum"
        elif mid == "ocr":
            item["parent_brain"] = "thalamus"
        out.append(item)
    return out


def all_provider_options() -> list[dict[str, str]]:
    opts: list[dict[str, str]] = []
    seen: set[str] = set()
    for item in list_catalog():
        pid = str(item["id"])
        if pid in seen:
            continue
        seen.add(pid)
        opts.append({"id": pid, "label": str(item["label"])})
        for alias in item.get("aliases") or []:
            opts.append({"id": str(alias), "label": f"{item['label']}（{alias}）"})
    return opts


def model_choices_for_provider(provider_id: str) -> list[str]:
    cat = catalog_by_id().get(provider_id) or {}
    return list(cat.get("models") or [])
