#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Qwen 手册 · 步骤 6：对齐 Qwen context_size=64000（勿写 max_tokens 进 model_properties）。"""
from __future__ import annotations

import sys
import time

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


def main() -> int:
    d = Dify(C.DIFY_EMAIL, C.DIFY_PASSWORD)
    try:
        ok, _, _ = d.login()
        print("login ok:", ok)
        if not ok:
            return 1

        r = d.get(
            "/workspaces/current/model-providers/%s/models/credentials?model_type=llm&model=%s"
            % (C.OPENAI_PROVIDER, C.QWEN_MODEL)
        )
        js = r.get("json") or {}
        cur_cred = js.get("credentials") or {}
        cid = js.get("current_credential_id")
        print("current_credential_id:", cid)

        cred = dict(cur_cred)
        cred["model_properties"] = {"context_size": C.QWEN_CTX}
        name = js.get("current_credential_name") or C.QWEN_MODEL

        r = d.req(
            "PUT",
            "/workspaces/current/model-providers/%s/models/credentials" % C.OPENAI_PROVIDER,
            body={
                "model": C.QWEN_MODEL,
                "model_type": "llm",
                "credential_id": cid,
                "credentials": cred,
                "name": name,
            },
        )
        print("PUT cred HTTP", r.get("code"), (r.get("text") or "")[:200])

        time.sleep(2)
        r = d.get("/workspaces/current/models/model-types/llm")
        for p in (r.get("json") or {}).get("data") or []:
            if p.get("provider") == C.OPENAI_PROVIDER:
                for m in p.get("models") or []:
                    if m.get("model") == C.QWEN_MODEL:
                        mp = m.get("model_properties") or {}
                        print("VERIFY context_size=", mp.get("context_size"))
        return 0 if r.get("code") == "200" else 1
    finally:
        d.close()


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