#!/usr/bin/env python3
"""
批量安装模型供应商插件脚本
解决设置页面只显示 openai-api-compatible 的问题

使用方法:
    cd dify-main/api
    python install_model_plugins.py
"""

import sys
import os

# 确保在正确的目录中运行
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

# 常用模型供应商插件列表
# 格式: langgenius/{provider_name}
MODEL_PROVIDER_PLUGINS = [
    "langgenius/openai",           # OpenAI (GPT-4o, GPT-4, etc.)
    "langgenius/anthropic",        # Anthropic (Claude)
    "langgenius/deepseek",         # DeepSeek
    "langgenius/google",           # Google (Gemini)
    "langgenius/azure_openai",     # Azure OpenAI
    "langgenius/ollama",           # Ollama (本地模型)
    "langgenius/zhipuai",          # 智谱 AI (GLM)
    "langgenius/tongyi",           # 通义千问
    "langgenius/spark",            # 讯飞星火
    "langgenius/minimax",          # MiniMax
    "langgenius/moonshot",         # Moonshot (Kimi)
    "langgenius/baichuan",         # 百川
    "langgenius/yi",               # 零一万物 (Yi)
    "langgenius/mistralai",        # Mistral AI
    "langgenius/groq",             # Groq
    "langgenius/cohere",           # Cohere
    "langgenius/replicate",        # Replicate
    "langgenius/togetherai",       # Together AI
    "langgenius/openrouter",       # OpenRouter
    "langgenius/huggingface_hub",  # Hugging Face
    "langgenius/siliconflow",      # SiliconFlow
    "langgenius/fireworks",        # Fireworks AI
    "langgenius/volcengine_maas",  # 火山引擎
    "langgenius/tencent",          # 腾讯混元
    "langgenius/wenxin",           # 文心一言
    "langgenius/nvidia_nim",       # NVIDIA NIM
]


def get_tenant_id():
    """从数据库获取第一个 tenant_id"""
    from extensions.ext_database import db
    from models.account import Tenant

    tenant = db.session.query(Tenant).first()
    if not tenant:
        print("错误: 数据库中没有找到租户信息")
        sys.exit(1)
    return tenant.id


def fetch_plugin_identifiers(plugin_ids: list[str]) -> dict[str, str]:
    """从 marketplace 获取插件的唯一标识符"""
    from core.helper.marketplace import batch_fetch_plugin_manifests

    result = {}
    try:
        manifests = batch_fetch_plugin_manifests(plugin_ids)
        for manifest in manifests:
            plugin_id = manifest.plugin_id
            identifier = manifest.latest_package_identifier
            if identifier:
                result[plugin_id] = identifier
                print(f"  ✓ {plugin_id} -> {identifier}")
            else:
                print(f"  ✗ {plugin_id} -> 无法获取标识符")
    except Exception as e:
        print(f"  获取插件信息失败: {e}")

    return result


def install_plugins(tenant_id: str, plugin_identifiers: dict[str, str]):
    """安装插件到 Plugin Daemon"""
    from core.plugin.impl.plugin import PluginInstaller
    from core.plugin.entities.plugin_daemon import PluginInstallationSource
    from core.plugin.plugin_service import PluginService

    installer = PluginInstaller()

    # 检查已安装的插件
    try:
        installed = installer.list_plugins(tenant_id)
        installed_ids = {p.plugin_id for p in installed}
        print(f"\n已安装 {len(installed_ids)} 个插件")
    except Exception as e:
        print(f"获取已安装插件列表失败: {e}")
        installed_ids = set()

    # 过滤出需要安装的插件
    to_install = {
        pid: identifier
        for pid, identifier in plugin_identifiers.items()
        if pid not in installed_ids
    }

    if not to_install:
        print("\n所有插件均已安装，无需操作")
        return

    print(f"\n需要安装 {len(to_install)} 个插件...")

    # 分批安装（每批最多 64 个）
    identifiers_list = list(to_install.values())
    metas = [{"plugin_unique_identifier": uid} for uid in identifiers_list]

    for i in range(0, len(identifiers_list), 64):
        batch = identifiers_list[i:i + 64]
        batch_metas = metas[i:i + 64]

        try:
            response = installer.install_from_identifiers(
                tenant_id,
                batch,
                PluginInstallationSource.Marketplace,
                batch_metas,
            )
            print(f"  批次 {i // 64 + 1}: 已提交 {len(batch)} 个插件安装任务")
        except Exception as e:
            print(f"  批次 {i // 64 + 1}: 安装失败 - {e}")

    # 清除缓存
    try:
        PluginService.invalidate_plugin_model_providers_cache(tenant_id)
        print("\n已清除插件缓存")
    except Exception as e:
        print(f"\n清除缓存失败: {e}")

    print("\n安装完成！请刷新设置页面查看模型供应商列表")


def main():
    print("=" * 60)
    print("  Dify 模型供应商插件批量安装工具")
    print("=" * 60)

    # 初始化 Flask 应用
    from app import create_app
    app = create_app()

    with app.app_context():
        # 获取 tenant_id
        print("\n[1/3] 获取租户信息...")
        tenant_id = get_tenant_id()
        print(f"  租户 ID: {tenant_id}")

        # 获取插件标识符
        print("\n[2/3] 从 Marketplace 获取插件信息...")
        plugin_identifiers = fetch_plugin_identifiers(MODEL_PROVIDER_PLUGINS)

        if not plugin_identifiers:
            print("\n错误: 无法获取任何插件信息，请检查网络连接和 Marketplace 配置")
            sys.exit(1)

        # 安装插件
        print(f"\n[3/3] 安装插件 (共 {len(plugin_identifiers)} 个)...")
        install_plugins(tenant_id, plugin_identifiers)

        print("\n" + "=" * 60)
        print("  完成！")
        print("=" * 60)


if __name__ == "__main__":
    main()