import json

from dify_api import Dify



SNAP_ID = "<snapshot_id>"

DRAFT_ID = "<draft_id>"

DS_ID = "<dataset_id>"

PROVIDER = "langgenius/xinference/xinference"

RERANK = "bge-reranker-v2-m3"



d = Dify("38205036@qq.com", "Difyai123456")



def load_cfg(table, cid):

    out, _ = 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)

    )

    return json.loads(out.strip())



def save_cfg(table, cid, cfg, tag):

    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 bind_and_rerank(cfg):

    know = cfg.setdefault("knowledge", {})

    sets = know.setdefault("sets", [{}])

    sets[0]["dataset_id"] = DS_ID   # 若现有结构用 id / dataset 字段，以 load 出来的为准

    retr = sets[0].setdefault("retrieval", {})

    retr["top_k"] = 5

    retr["reranking_enable"] = True

    retr["reranking_mode"] = "reranking_model"

    retr["reranking_model"] = {"provider": PROVIDER, "model": RERANK}

    return cfg



snap = load_cfg("agent_config_snapshots", SNAP_ID)

draft = load_cfg("agent_config_drafts", DRAFT_ID)

# 先把 snap/draft 备份成本地 json

save_cfg("agent_config_snapshots", SNAP_ID, bind_and_rerank(snap), "snap")

save_cfg("agent_config_drafts", DRAFT_ID, bind_and_rerank(draft), "draft")