#!/usr/bin/env python3
"""把 bge-m3 和 bge-reranker-v2-m3 拉起来，并写入 Xinference 自启动。

已经在跑的模型不会重复启动。登录口令是现网管理员口令。
"""
import json
import sqlite3
import urllib.request

BASE = "http://127.0.0.1:9997"
PASSWORD = "Xinf-bge-2026"
SPECS = [
    {
        "model_uid": "bge-m3",
        "model_name": "bge-m3",
        "model_type": "embedding",
        "model_engine": "sentence_transformers",
        "n_gpu": 1,
        "gpu_idx": [0],
        "download_hub": "modelscope",
        "replica": 1,
    },
    {
        "model_uid": "bge-reranker-v2-m3",
        "model_name": "bge-reranker-v2-m3",
        "model_type": "rerank",
        "model_engine": "sentence_transformers",
        "n_gpu": 1,
        "gpu_idx": [0],
        "download_hub": "modelscope",
        "replica": 1,
    },
]


def call(method, path, token=None, payload=None, timeout=900):
    headers = {}
    data = None
    if payload is not None:
        data = json.dumps(payload).encode()
        headers["Content-Type"] = "application/json"
    if token:
        headers["Authorization"] = "Bearer " + token
    req = urllib.request.Request(BASE + path, data=data, headers=headers, method=method)
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        body = resp.read().decode()
        return json.loads(body) if body else {}


def main():
    user = sqlite3.connect("/usr/cluster/dify/xinference/auth/auth.db").execute(
        "select username from users"
    ).fetchone()[0]
    token = call("POST", "/token", payload={"username": user, "password": PASSWORD})["access_token"]
    running = call("GET", "/v1/models", token=token).get("data") or []
    ids = {item.get("id") for item in running}
    print("before", sorted(ids), flush=True)
    for spec in SPECS:
        uid = spec["model_uid"]
        if uid not in ids:
            print("launch", uid, flush=True)
            print("launched", call("POST", "/v1/models?wait_ready=true", token=token, payload=spec), flush=True)
        else:
            print("already", uid, flush=True)
        call("POST", "/v1/autostart/models", token=token, payload={"enabled": True, "launch": spec})
        print("autostart", uid, flush=True)
    after = call("GET", "/v1/models", token=token).get("data") or []
    for item in after:
        print("model", item.get("id"), item.get("model_type") or item.get("type"), flush=True)


if __name__ == "__main__":
    main()