# -*- coding: utf-8 -*-
"""新 Agent 快照/草稿数据库读写（专题十一、十二）。"""
from __future__ import annotations

import json
import uuid
from typing import Any

import _config as C


def psql(d, sql: str) -> str:
    remote = "/tmp/agent_db_%s.sql" % uuid.uuid4().hex[:8]
    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)
    )
    return (out or err).strip()


def load_snapshot(d, table: str, row_id: str) -> dict[str, Any]:
    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, row_id)
    )
    raw = out.strip()
    if not raw:
        raise RuntimeError(err or "empty config_snapshot for %s.%s" % (table, row_id))
    return json.loads(raw)


def backup_snapshot(d, table: str, row_id: str, label: str) -> str:
    cfg = load_snapshot(d, table, row_id)
    remote = "/root/backup_%s_%s.json" % (label, uuid.uuid4().hex[:8])
    with d.sftp.open(remote, "w") as f:
        f.write(json.dumps(cfg, ensure_ascii=False, indent=2))
    print("backup %s -> %s" % (table, remote))
    return remote


def read_max_tokens(d, table: str, row_id: str) -> Any:
    out, err = d.sh(
        "docker exec %s psql -U postgres -d dify -t -A -c "
        "\"SELECT config_snapshot::json -> 'model' -> 'model_settings' -> 'max_tokens' "
        "FROM %s WHERE id='%s'\"" % (C.PG_CONTAINER, table, row_id)
    )
    val = out.strip()
    return None if val in ("", "null") else val


def set_max_tokens(d, table: str, row_id: str, value: int) -> None:
    sql = (
        "UPDATE %s SET config_snapshot = jsonb_set("
        "config_snapshot::jsonb, '{model,model_settings,max_tokens}', '%s'::jsonb)::text "
        "WHERE id='%s';" % (table, value, row_id)
    )
    result = psql(d, sql)
    print(table, result)


def verify_max_tokens(d, snap_id: str, draft_id: str) -> None:
    for which, table, rid in (
        ("snapshot", "agent_config_snapshots", snap_id),
        ("draft", "agent_config_drafts", draft_id),
    ):
        val = read_max_tokens(d, table, rid)
        print("%s max_tokens = %s" % (which, val))