#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""专题十五 · 步骤 5.2：改 Agent 快照+草稿 knowledge.sets 绑库与 Rerank。"""
from __future__ import annotations

import json
import sys

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


def bind_and_rerank(cfg: dict) -> dict:
    know = cfg.setdefault("knowledge", {})
    sets = know.setdefault("sets", [{}])
    s0 = sets[0]
    if "dataset_id" not in s0 and "id" not in s0:
        s0["dataset_id"] = C.DS_ID
    elif "dataset_id" in s0:
        s0["dataset_id"] = C.DS_ID
    retr = s0.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 = AC.load_cfg(d, "agent_config_snapshots", C.SNAP_ID)
        draft = AC.load_cfg(d, "agent_config_drafts", C.DRAFT_ID)
        AC.save_cfg(d, "agent_config_snapshots", C.SNAP_ID, bind_and_rerank(snap), "snap15")
        AC.save_cfg(d, "agent_config_drafts", C.DRAFT_ID, bind_and_rerank(draft), "draft15")
        out, _ = d.sh(
            "docker exec %s psql -U postgres -d dify -t -A -c "
            "\"SELECT config_snapshot::json -> 'knowledge' FROM agent_config_snapshots WHERE id='%s';\""
            % (C.PG_CONTAINER, C.SNAP_ID)
        )
        print("knowledge:", out.strip())
        print("\n[WARN] 改完不要点「构建/发布」")
        return 0
    finally:
        d.close()


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