#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""专题十三 · 步骤 3.2 路径 B：把 Dify 的 Xinference URL 改到 GPU 机（修改后）。"""
from __future__ import annotations

import sys

import _bootstrap  # noqa: F401
import _config as C
from dify_api import Dify, pp


def patch_credentials(d: Dify, model: str, model_type: str, server_url: str, api_key: str) -> None:
    body = {
        "model": model,
        "model_type": model_type,
        "credentials": {
            "server_url": server_url,
            "model_uid": model,
            "api_key": api_key,
            "invoke_timeout": "60",
            "max_retries": "3",
        },
        "name": model,
    }
    path = "/workspaces/current/model-providers/%s/models/credentials" % C.XINFERENCE_PROVIDER
    r = d.post(path, body)
    pp("POST %s %s" % (model_type, model), r)
    if r.get("code") not in ("200", "201"):
        r = d.req("PUT", path, body=body)
        pp("PUT %s %s" % (model_type, model), r)


def main() -> int:
    gpu_key = C.XF_KEY  # 新起 GPU 容器时换成 116 上生成的 xf- Key
    d = Dify(C.DIFY_EMAIL, C.DIFY_PASSWORD)
    try:
        d.login()
        print("=== 修改前 URL（207 CPU）===")
        print("embedding/rerank server_url:", C.XF_URL)
        print("\n=== 修改后 URL（116 GPU）===")
        print("embedding/rerank server_url:", C.XF_GPU_URL)
        patch_credentials(d, C.EMBED_MODEL, "text-embedding", C.XF_GPU_URL, gpu_key)
        patch_credentials(d, C.RERANK_MODEL, "rerank", C.XF_GPU_URL, gpu_key)
        return 0
    finally:
        d.close()


if __name__ == "__main__":
    sys.exit(main())