# -*- coding: utf-8 -*-
"""读写 agent_config_snapshots / agent_config_drafts 的共用函数。"""
from __future__ import annotations

import json

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 %s psql -U postgres -d dify -t -A -c "
        "\"SELECT config_snapshot FROM %s WHERE id='%s'\"" % (C.PG_CONTAINER, table, cid)
    )
    raw = out.strip()
    if not raw:
        raise RuntimeError(err or "empty config_snapshot for %s" % cid)
    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 %s psql -U postgres -d dify < %s" % (C.PG_CONTAINER, remote))
    print(table, (out or err).strip())


def set_max_tokens(d: Dify, snap_id: str, draft_id: str, value: int) -> None:
    sql = (
        "BEGIN; "
        "UPDATE agent_config_snapshots SET config_snapshot = jsonb_set("
        "config_snapshot::jsonb, '{model,model_settings,max_tokens}', '%d'::jsonb)::text "
        "WHERE id='%s'; "
        "UPDATE agent_config_drafts SET config_snapshot = jsonb_set("
        "config_snapshot::jsonb, '{model,model_settings,max_tokens}', '%d'::jsonb)::text "
        "WHERE id='%s'; "
        "COMMIT;"
    ) % (value, snap_id, value, draft_id)
    out, err = d.sh(
        'docker exec %s psql -U postgres -d dify -c "%s"' % (C.PG_CONTAINER, sql)
    )
    print(out or err)