#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""专题一 · 步骤 3：设默认 LLM（API + tenant_default_models）。"""
from __future__ import annotations

import sys

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


def main() -> int:
    d = Dify(C.DIFY_EMAIL, C.DIFY_PASSWORD)
    try:
        d.login()
        r = d.post(
            "/workspaces/current/default-model",
            body={
                "model_settings": [{
                    "model_type": "llm",
                    "provider": C.OPENAI_PROVIDER,
                    "model": C.QWEN_MODEL,
                }]
            },
        )
        pp("set default llm (API)", r)

        sql = (
            "UPDATE tenant_default_models SET model='%s' WHERE model_type='llm';"
            % C.QWEN_MODEL
        )
        out, err = d.sh(
            "docker exec %s psql -U postgres -d dify -c \"%s\""
            % (C.PG_CONTAINER, sql)
        )
        print("tenant_default_models:\n", out or err)

        out2, err2 = d.sh(
            "cd %s && docker compose restart api plugin_daemon" % C.DIFY_DOCKER_DIR,
            t=120,
        )
        print(out2 or err2)
        return 0
    finally:
        d.close()


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