#!/usr/bin/env bash
set -euo pipefail

usage() {
    cat <<'USAGE'
用法:
  scripts/run_rag_public_orphan_defaults_migration.sh [连接参数] [--migrate --yes]

连接方式二选一:
  1. 使用 DATABASE_URL:
     DATABASE_URL='postgresql://user:password@host:5432/dbname' scripts/run_rag_public_orphan_defaults_migration.sh

  2. 使用 psql 参数:
     scripts/run_rag_public_orphan_defaults_migration.sh -h <host> -U <user> -d <db_name> [-p <port>]

默认行为:
  只执行迁移前检查，不修改数据库。

执行迁移:
  先确认 precheck 输出符合预期，再追加 --migrate --yes:
     scripts/run_rag_public_orphan_defaults_migration.sh -h <host> -U <user> -d <db_name> --migrate --yes

环境变量:
  PGPASSWORD 可用于非交互输入数据库密码。
USAGE
}

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
PRECHECK_SQL="${REPO_ROOT}/scripts/创建sql/precheck_rag_public_orphan_defaults.sql"
MIGRATE_SQL="${REPO_ROOT}/scripts/创建sql/migrate_rag_public_orphan_defaults.sql"
LOG_DIR="${REPO_ROOT}/logs/rag-public-orphan-defaults"

PSQL_ARGS=()
RUN_MIGRATION=false
CONFIRMED=false

while [[ $# -gt 0 ]]; do
    case "$1" in
        -h|--host)
            [[ $# -ge 2 ]] || { echo "缺少 $1 的值" >&2; exit 2; }
            PSQL_ARGS+=("-h" "$2")
            shift 2
            ;;
        -U|--username)
            [[ $# -ge 2 ]] || { echo "缺少 $1 的值" >&2; exit 2; }
            PSQL_ARGS+=("-U" "$2")
            shift 2
            ;;
        -d|--dbname)
            [[ $# -ge 2 ]] || { echo "缺少 $1 的值" >&2; exit 2; }
            PSQL_ARGS+=("-d" "$2")
            shift 2
            ;;
        -p|--port)
            [[ $# -ge 2 ]] || { echo "缺少 $1 的值" >&2; exit 2; }
            PSQL_ARGS+=("-p" "$2")
            shift 2
            ;;
        --migrate)
            RUN_MIGRATION=true
            shift
            ;;
        --yes)
            CONFIRMED=true
            shift
            ;;
        --help)
            usage
            exit 0
            ;;
        *)
            echo "未知参数: $1" >&2
            usage
            exit 2
            ;;
    esac
done

if ! command -v psql >/dev/null 2>&1; then
    echo "未找到 psql，请先安装 PostgreSQL client。" >&2
    exit 1
fi

if [[ ! -f "${PRECHECK_SQL}" ]]; then
    echo "找不到预检查 SQL: ${PRECHECK_SQL}" >&2
    exit 1
fi

if [[ ! -f "${MIGRATE_SQL}" ]]; then
    echo "找不到迁移 SQL: ${MIGRATE_SQL}" >&2
    exit 1
fi

if [[ -n "${DATABASE_URL:-}" && ${#PSQL_ARGS[@]} -gt 0 ]]; then
    echo "请不要同时使用 DATABASE_URL 和 -h/-U/-d 参数。" >&2
    exit 2
fi

if [[ -z "${DATABASE_URL:-}" && ${#PSQL_ARGS[@]} -eq 0 ]]; then
    echo "缺少数据库连接参数。请设置 DATABASE_URL，或传入 -h/-U/-d。" >&2
    usage
    exit 2
fi

if [[ "${RUN_MIGRATION}" == true && "${CONFIRMED}" != true ]]; then
    echo "迁移会修改数据库。请先检查 precheck 输出，确认后追加 --yes。" >&2
    exit 2
fi

mkdir -p "${LOG_DIR}"
TIMESTAMP="$(date +%Y%m%d-%H%M%S)"
PRECHECK_LOG="${LOG_DIR}/precheck-${TIMESTAMP}.log"
MIGRATE_LOG="${LOG_DIR}/migrate-${TIMESTAMP}.log"

run_psql_file() {
    local sql_file="$1"
    local log_file="$2"

    if [[ -n "${DATABASE_URL:-}" ]]; then
        psql "${DATABASE_URL}" -v ON_ERROR_STOP=1 -f "${sql_file}" 2>&1 | tee "${log_file}"
    else
        psql "${PSQL_ARGS[@]}" -v ON_ERROR_STOP=1 -f "${sql_file}" 2>&1 | tee "${log_file}"
    fi
}

echo "开始执行 RAG 公共知识库未归属默认项预检查..."
echo "预检查输出: ${PRECHECK_LOG}"
run_psql_file "${PRECHECK_SQL}" "${PRECHECK_LOG}"

if [[ "${RUN_MIGRATION}" != true ]]; then
    cat <<EOF

预检查已完成，未执行迁移。
请确认上方第 1 段候选只包含需要从“未分配地区”迁到“公共”的历史公共知识库。
确认后执行同一命令并追加: --migrate --yes
EOF
    exit 0
fi

echo
echo "开始执行迁移..."
echo "迁移输出: ${MIGRATE_LOG}"
run_psql_file "${MIGRATE_SQL}" "${MIGRATE_LOG}"

cat <<EOF

迁移脚本已执行完成。
请检查迁移输出中的 public_default_count，应为 1。
然后刷新 /chat-with-llm/dataset-manager，原“未分配地区”的默认项应归入“公共”，并且非默认项可以删除。
EOF