from __future__ import annotations

import copy
import os
import re
from pathlib import Path
from typing import Any

import yaml

from niannian.core.config import _load_yaml, resolve_config_path

# 硅基永久免费模型（与 config/docker.yaml 脑区 blueprint 对齐）
SILICONFLOW_FREE_MODELS: tuple[str, ...] = (
    "Qwen/Qwen3-8B",
    "Qwen/Qwen3.5-4B",
    "Qwen/Qwen2.5-7B-Instruct",
    "deepseek-ai/DeepSeek-R1-0528-Qwen3-8B",
    "BAAI/bge-m3",
    "BAAI/bge-large-zh-v1.5",
    "BAAI/bge-large-en-v1.5",
    "BAAI/bge-reranker-v2-m3",
    "deepseek-ai/DeepSeek-OCR",
    "PaddlePaddle/PaddleOCR-VL-1.5",
)

ZHIPU_FREE_MODELS: tuple[str, ...] = (
    "glm-4.7-flash",
    "glm-4-flash",
)

DEFAULT_FREE_MODELS: dict[str, tuple[str, ...]] = {
    "siliconflow": SILICONFLOW_FREE_MODELS,
    "zhipu": ZHIPU_FREE_MODELS,
}

_MODEL_ID_RE = re.compile(r"^[\w./+-]+$")


class ModelRegistryError(Exception):
    pass


def resolve_llm_models_path(*, for_write: bool = False) -> Path:
    env = os.environ.get("NIANNIAN_LLM_MODELS_CONFIG")
    if env:
        return Path(env).expanduser().resolve()
    cfg_dir = resolve_config_path().parent
    path = (cfg_dir / "llm_models.local.yaml").resolve()
    if for_write or path.is_file():
        return path
    pkg = Path(__file__).resolve().parents[2] / "config" / "llm_models.local.yaml"
    return pkg if pkg.is_file() else path


def _empty_doc() -> dict[str, Any]:
    return {"provider_models": {}}


def load_model_registry() -> dict[str, Any]:
    path = resolve_llm_models_path()
    if path.is_file():
        raw = _load_yaml(path)
        if isinstance(raw, dict):
            raw.setdefault("provider_models", {})
            return raw
    return _empty_doc()


def save_model_registry(doc: dict[str, Any]) -> str:
    path = resolve_llm_models_path(for_write=True)
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as f:
        yaml.safe_dump(doc, f, allow_unicode=True, default_flow_style=False, sort_keys=False)
    return str(path)


def _validate_model_id(model_id: str) -> str:
    mid = model_id.strip()
    if not mid or not _MODEL_ID_RE.match(mid):
        raise ModelRegistryError(f"invalid model id: {model_id!r}")
    return mid


def list_provider_model_entries(provider_id: str) -> list[dict[str, Any]]:
    doc = load_model_registry()
    stored = (doc.get("provider_models") or {}).get(provider_id)
    free_defaults = set(DEFAULT_FREE_MODELS.get(provider_id, ()))
    if not stored:
        return [
            {"id": mid, "free": mid in free_defaults, "builtin": True}
            for mid in DEFAULT_FREE_MODELS.get(provider_id, ())
        ]
    out: list[dict[str, Any]] = []
    for item in stored:
        if isinstance(item, str):
            mid = item.strip()
            if mid:
                out.append({"id": mid, "free": mid in free_defaults, "builtin": mid in free_defaults})
        elif isinstance(item, dict) and item.get("id"):
            mid = str(item["id"])
            out.append(
                {
                    "id": mid,
                    "free": bool(item.get("free")) or mid in free_defaults,
                    "builtin": mid in free_defaults,
                }
            )
    return out


def model_ids_for_provider(provider_id: str) -> list[str]:
    return [e["id"] for e in list_provider_model_entries(provider_id)]


def init_free_models(*provider_ids: str) -> dict[str, list[str]]:
    doc = load_model_registry()
    pm = doc.setdefault("provider_models", {})
    out: dict[str, list[str]] = {}
    for pid in provider_ids:
        defaults = list(DEFAULT_FREE_MODELS.get(pid, ()))
        if defaults:
            pm[pid] = list(defaults)
            out[pid] = list(defaults)
    save_model_registry(doc)
    return out


def add_provider_model(provider_id: str, model_id: str) -> list[dict[str, Any]]:
    mid = _validate_model_id(model_id)
    doc = load_model_registry()
    pm = doc.setdefault("provider_models", {})
    current = model_ids_for_provider(provider_id)
    if mid in current:
        raise ModelRegistryError(f"model already exists: {mid}")
    if provider_id not in pm or not pm[provider_id]:
        pm[provider_id] = list(DEFAULT_FREE_MODELS.get(provider_id, ()))
        current = list(pm[provider_id])
    pm[provider_id] = current + [mid]
    save_model_registry(doc)
    return list_provider_model_entries(provider_id)


def remove_provider_model(provider_id: str, model_id: str) -> list[dict[str, Any]]:
    mid = _validate_model_id(model_id)
    doc = load_model_registry()
    pm = doc.setdefault("provider_models", {})
    current = model_ids_for_provider(provider_id)
    if mid not in current:
        raise ModelRegistryError(f"model not found: {mid}")
    pm[provider_id] = [m for m in current if m != mid]
    save_model_registry(doc)
    return list_provider_model_entries(provider_id)


def all_provider_models_view() -> dict[str, list[dict[str, Any]]]:
    from niannian.llm.provider_catalog import list_catalog

    out: dict[str, list[dict[str, Any]]] = {}
    for item in list_catalog():
        pid = str(item["id"])
        entries = list_provider_model_entries(pid)
        if entries:
            out[pid] = entries
        elif item.get("models"):
            out[pid] = [{"id": m, "free": False, "builtin": False} for m in item["models"]]
    return out
