#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""专题十七 · 步骤 2：把 OpenAI 插件钉住 —— 关掉 model 类自动升级（当前生产做法）。

这是本专题的「根治」动作。前面 §4 的批量补丁只解决「已经坏了」；
不改这张表，下次自动升级把 openai 装成新版本目录，补丁再次落空，模型再变「不可用」。

做两件事，互为双保险：
  1. strategy_setting = 'disabled'
     —— 整个 model 类不再做任何自动升级；
  2. upgrade_mode = 'exclude' + exclude_plugins = ['langgenius/openai']
     —— 即便日后有人把 setting 改回 fix_only / latest，openai 也仍在排除名单里。

安全设计：
  * 默认 **dry-run**（只打印改前改后，不写库）；确认无误后再 `--apply`；
  * 写库前先把现值备份成 JSON 落到 207 的 /opt/workbuddy/data/ 下；
  * 只动 category='model' 这一行，其余 5 行（tool / extension / agent-strategy /
    datasource / trigger）保持原样 —— 它们维持 fix_only 是合理的。

用法（Windows 开发机，项目根）：
    # 1) 先看现状、演练
    python scripts_cursor/Dify_专题十七_OpenAI插件升级后schema补丁丢失/专题十七_02_关闭OpenAI插件自动升级.py
    # 2) 确认输出无误后真正写库
    python scripts_cursor/Dify_专题十七_OpenAI插件升级后schema补丁丢失/专题十七_02_关闭OpenAI插件自动升级.py --apply

回滚：把备份 JSON 里的 strategy_setting / upgrade_mode / exclude_plugins 三个值
      用同样的 UPDATE 写回即可（脚本会把具体 SQL 一并打印出来）。
"""
from __future__ import annotations

import json
import sys
import time

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

PG = C.PG_CONTAINER
DB = "dify"
TARGET_PLUGIN = "langgenius/openai"
BACKUP_DIR = "/opt/workbuddy/data"

SELECT = (
    "select coalesce(json_agg(row_to_json(t))::text, '[]') from ("
    "select category, strategy_setting, upgrade_mode, exclude_plugins, include_plugins "
    "from tenant_plugin_auto_upgrade_strategies where category = 'model') t;"
)

UPDATE = (
    "update tenant_plugin_auto_upgrade_strategies "
    "set strategy_setting = 'disabled', "
    "    upgrade_mode = 'exclude', "
    "    exclude_plugins = '[\"%s\"]'::json, "
    "    updated_at = now() "
    "where category = 'model';"
) % TARGET_PLUGIN


def psql(d, sql: str):
    return d.sh('docker exec %s psql -U postgres -d %s -c "%s"' % (PG, DB, sql))


def main() -> int:
    apply_it = "--apply" in sys.argv
    d = Dify(C.DIFY_EMAIL, C.DIFY_PASSWORD)

    print("=" * 78)
    print("改前现状（category = model）")
    print("=" * 78)
    out, err = psql(d, SELECT)
    raw = (out or "").strip()
    # psql -c 输出会带表头与行数行，取第一段像 JSON 的行
    before = None
    for line in raw.splitlines():
        s = line.strip()
        if s.startswith("[") or s.startswith("{"):
            try:
                before = json.loads(s)
                break
            except Exception:
                continue
    print(raw or err)
    if before is None:
        print("[WARN] 没能解析出改前的 JSON，先人工核对上面的输出。")

    # ---- 备份到 207 ----
    if before is not None:
        ts = time.strftime("%Y%m%d_%H%M%S")
        path = "%s/plugin_autoupgrade_backup_%s.json" % (BACKUP_DIR, ts)
        payload = json.dumps(
            {"table": "tenant_plugin_auto_upgrade_strategies",
             "db": DB, "category": "model", "rows": before,
             "saved_at": ts}, ensure_ascii=False, indent=2)
        d.sh("mkdir -p %s" % BACKUP_DIR)
        d.sh("cat > %s <<'WBEOF'\n%s\nWBEOF" % (path, payload))
        print("\n[backup] 现值已存到 207：%s" % path)

    print("\n" + "=" * 78)
    print("将要执行的 SQL")
    print("=" * 78)
    print(UPDATE)

    if not apply_it:
        print("\n[dry-run] 未写库。确认无误后加 --apply 重跑。")
        print("回滚：把备份 JSON 里的三个字段用同样的 UPDATE 写回即可。")
        return 0

    print("\n" + "=" * 78)
    print("执行 UPDATE ...")
    print("=" * 78)
    out, err = psql(d, UPDATE)
    print(out or err)

    print("=" * 78)
    print("改后复核")
    print("=" * 78)
    out, err = psql(d, SELECT)
    print(out or err)

    print("\n预期结果：")
    print("  strategy_setting = disabled")
    print("  upgrade_mode     = exclude")
    print('  exclude_plugins  = ["%s"]' % TARGET_PLUGIN)
    print("  updated_at       = 刚刚")
    print("\n注意：关掉自动升级后，openai 插件要升级必须**手动**做，")
    print("      升完记得立刻重跑 专题一_04 给新版本目录补 llm.py，再刷新模型列表。")
    return 0


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