#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""专题十二 · 步骤 4 · 修改后：把 max_tokens 从 2048 改回 8192（快照+草稿，勿点发布）。"""
from __future__ import annotations

import sys

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


def main() -> int:
    d = Dify(C.DIFY_EMAIL, C.DIFY_PASSWORD)
    try:
        d.login()
        sql = "BEGIN;\n"
        for table, rid in (
            ("agent_config_snapshots", C.SNAP_ID),
            ("agent_config_drafts", C.DRAFT_ID),
        ):
            sql += (
                "UPDATE %s SET config_snapshot = jsonb_set("
                "config_snapshot::jsonb, '{model,model_settings,max_tokens}', '%s'::jsonb)::text "
                "WHERE id='%s';\n" % (table, C.MAX_TOKENS_DEFAULT, rid)
            )
        sql += "COMMIT;\n"
        print(A.psql(d, sql))
        A.verify_max_tokens(d, C.SNAP_ID, C.DRAFT_ID)
        print("\n[WARN] 改完不要点「构建/发布」；请新开对话验证")
        return 0
    finally:
        d.close()


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