"""
能力(capability)注册中心

集中定义平台支持的全部模型能力，以及"每个能力对应的默认模型"的读写与解析。
- CAPABILITY_TYPES: 全部能力（id / 中文标签 / 测试类型 / 描述）
- get_capability_defaults / set_capability_default: 读写 site_settings 中的 model_capability_defaults
- resolve_auto_model: 根据能力解析平台默认模型（供聊天 / 技能 / 应用 / IM Bot 自动匹配）
"""
from typing import Optional
from sqlalchemy.orm import Session
from .models import AIModel, SiteSetting
import json


# ============================================================
# 平台支持的全部能力类型
# test_type 用于模型连接测试时选择对应的 API 路径
# ============================================================
CAPABILITY_TYPES = [
    {"id": "text",       "label": "文本对话", "test_type": "text",       "desc": "文本生成与多轮对话"},
    {"id": "thinking",   "label": "推理思考", "test_type": "text",       "desc": "深度推理 / 思维链"},
    {"id": "vision",     "label": "视觉理解", "test_type": "vision",     "desc": "图片 / 多模态理解"},
    {"id": "tool",       "label": "工具调用", "test_type": "text",       "desc": "Function Calling 函数调用"},
    {"id": "image",      "label": "图像生成", "test_type": "image",      "desc": "文生图 / 图像生成"},
    {"id": "audio",      "label": "语音合成", "test_type": "audio",      "desc": "TTS 文字转语音"},
    {"id": "asr",        "label": "语音识别", "test_type": "asr",        "desc": "ASR 语音转文字"},
    {"id": "ocr",        "label": "OCR识别",  "test_type": "ocr",        "desc": "图片文字识别提取"},
    {"id": "embedding",  "label": "嵌入向量", "test_type": "embedding",  "desc": "文本向量化 / 嵌入"},
    {"id": "rerank",     "label": "重排序",   "test_type": "rerank",     "desc": "语义重排 / Rerank"},
    {"id": "video",      "label": "视频生成", "test_type": "video",      "desc": "文生视频 / 视频生成"},
]

# 文本类能力（对话 / 视觉 / 推理 / 工具 / OCR）在缺少专属默认模型时，回退到全局默认模型
TEXT_FAMILY = {"text", "thinking", "vision", "tool", "ocr"}

# 仅用于前端展示的固定配色
CAPABILITY_COLORS = {
    "text": "#5B5FED", "thinking": "#8B5CF6", "vision": "#219B5E", "tool": "#3B82F6",
    "image": "#D9383A", "audio": "#E8922C", "asr": "#0EA5E9", "ocr": "#F59E0B",
    "embedding": "#10B981", "rerank": "#EC4899", "video": "#EF4444",
}

SETTING_KEY = "model_capability_defaults"


def get_capability_defaults(db: Session) -> dict:
    """读取 {capability: model_id} 映射"""
    row = db.query(SiteSetting).filter(SiteSetting.key == SETTING_KEY).first()
    if not row or not row.value:
        return {}
    try:
        data = json.loads(row.value)
        return data if isinstance(data, dict) else {}
    except Exception:
        return {}


def set_capability_default(db: Session, capability: str, model_id: Optional[str]) -> dict:
    """设置某个能力的默认模型（model_id 为空则清除）"""
    data = get_capability_defaults(db)
    if model_id:
        data[capability] = model_id
    else:
        data.pop(capability, None)
    row = db.query(SiteSetting).filter(SiteSetting.key == SETTING_KEY).first()
    if not row:
        from uuid import uuid4
        row = SiteSetting(id=str(uuid4()), key=SETTING_KEY, value="{}")
        db.add(row)
    row.value = json.dumps(data, ensure_ascii=False)
    db.commit()
    return data


def resolve_auto_model(capability: str, db: Session) -> Optional[AIModel]:
    """
    根据能力解析平台默认模型：
    1. 优先使用后台"按能力默认模型"中配置的值
    2. 文本类能力若未单独配置，回退到全局默认模型(is_default)
    3. 媒体类能力(image/video/audio)若未配置，自动发现第一个启用且具备该能力的模型
    4. 其余能力(embedding/rerank/asr 等)若未配置则返回 None
    """
    data = get_capability_defaults(db)
    mid = data.get(capability)
    if mid:
        m = db.query(AIModel).filter(
            AIModel.id == mid, AIModel.is_enabled == True
        ).first()
        if m:
            return m
    # 文本类能力回退到全局默认模型
    if capability in TEXT_FAMILY:
        m = db.query(AIModel).filter(
            AIModel.is_default == True, AIModel.is_enabled == True
        ).first()
        if m:
            return m
    # 媒体类能力(image/video/audio)：无显式配置时自动发现任一启用的模型
    # 避免用户未在后台配置"能力默认模型"时 Auto 回退到文本模型导致"无法生成图片/视频"
    if capability in ("image", "video", "audio"):
        from sqlalchemy import text as _text
        m = db.query(AIModel).filter(
            AIModel.is_enabled == True,
            AIModel.capabilities.op("JSON_CONTAINS")(
                AIModel.capabilities, f'"{capability}"'
            ),
        ).first()
        if m:
            return m
    return None