#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""专题十三 · 步骤 7 · 修改后：同步改 Agent 快照与草稿的 knowledge.sets[0].retrieval。"""
from __future__ import annotations

import json
import sys

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


def load_cfg(d: Dify, table: str, cid: str) -> dict:
    out, err = d.sh(
        "docker exec docker-db_postgres-1 psql -U postgres -d dify -t -A -c "
        "\"SELECT config_snapshot FROM %s WHERE id='%s'\"" % (table, cid)
    )
    raw = out.strip()
    if not raw:
        raise RuntimeError(err or "empty config_snapshot")
    return json.loads(raw)


def save_cfg(d: Dify, table: str, cid: str, cfg: dict, tag: str) -> None:
    sql = "UPDATE %s SET config_snapshot = $cfg$%s$cfg$ WHERE id='%s';" % (
        table,
        json.dumps(cfg, ensure_ascii=False),
        cid,
    )
    remote = "/tmp/upd_%s.sql" % tag
    with d.sftp.open(remote, "w") as f:
        f.write(sql)
    out, err = d.sh("docker exec -i docker-db_postgres-1 psql -U postgres -d dify < %s" % remote)
    print(table, (out or err).strip())


def enable_agent_rerank(cfg: dict) -> dict:
    retr = cfg["knowledge"]["sets"][0].setdefault("retrieval", {})
    retr["top_k"] = 5
    retr["reranking_enable"] = True
    retr["reranking_mode"] = "reranking_model"
    retr["reranking_model"] = {
        "provider": C.XINFERENCE_PROVIDER,
        "model": C.RERANK_MODEL,
    }
    return cfg


def main() -> int:
    d = Dify(C.DIFY_EMAIL, C.DIFY_PASSWORD)
    try:
        d.login()
        snap = load_cfg(d, "agent_config_snapshots", C.SNAP_ID)
        draft = load_cfg(d, "agent_config_drafts", C.DRAFT_ID)
        save_cfg(d, "agent_config_snapshots", C.SNAP_ID, enable_agent_rerank(snap), "snap_rerank")
        save_cfg(d, "agent_config_drafts", C.DRAFT_ID, enable_agent_rerank(draft), "draft_rerank")
        out, _ = d.sh(
            "docker exec docker-db_postgres-1 psql -U postgres -d dify -t -A -c "
            "\"SELECT config_snapshot::json -> 'knowledge' -> 'sets' -> 0 -> 'retrieval' "
            "FROM agent_config_snapshots WHERE id='%s';\"" % C.SNAP_ID
        )
        print("验证 retrieval:", out.strip())
        print("\n[WARN] 改完不要点「构建/发布」")
        return 0
    finally:
        d.close()


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