#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""专题十四 · 步骤 4：快照+草稿 max_tokens 改为 32768（不要点发布）。"""
from __future__ import annotations

import sys

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


def main() -> int:
    d = Dify(C.DIFY_EMAIL, C.DIFY_PASSWORD)
    try:
        d.login()
        out, err = d.sh(
            """docker exec %s psql -U postgres -d dify -c "
SELECT 'snapshot' AS which,
       config_snapshot::json -> 'model' -> 'model_settings' -> 'max_tokens' AS max_tokens
FROM agent_config_snapshots WHERE id='%s'
UNION ALL
SELECT 'draft',
       config_snapshot::json -> 'model' -> 'model_settings' -> 'max_tokens'
FROM agent_config_drafts WHERE id='%s';
"
"""
            % (C.PG_CONTAINER, C.SNAP_ID, C.DRAFT_ID)
        )
        print("before:\n", out or err)

        AC.set_max_tokens(d, C.SNAP_ID, C.DRAFT_ID, C.MAX_TOKENS)

        out2, err2 = d.sh(
            """docker exec %s psql -U postgres -d dify -c "
SELECT 'snapshot' AS which,
       config_snapshot::json -> 'model' -> 'model_settings' -> 'max_tokens' AS max_tokens
FROM agent_config_snapshots WHERE id='%s'
UNION ALL
SELECT 'draft',
       config_snapshot::json -> 'model' -> 'model_settings' -> 'max_tokens'
FROM agent_config_drafts WHERE id='%s';
"
"""
            % (C.PG_CONTAINER, C.SNAP_ID, C.DRAFT_ID)
        )
        print("after:\n", out2 or err2)
        print("\n[WARN] 改完不要点「构建/发布」；新开对话验证")
        return 0
    finally:
        d.close()


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