#!/usr/bin/env bash
# =============================================================================
# 共享法规知识库冷迁移脚本 (migrate_shared_kb.sh)
#
# 作用:
#   将参考项目容器 agent-regkb (端口 55432) 的 2,907 篇法规、31 万条切片
#   通过流式无损方式迁移至本项目专有基础设施容器 shuimu-kb-shared (端口 5432)
#
# 原则:
#   1. 源容器 agent-regkb 纯只读零侵入，严禁写操作；
#   2. 目标容器使用本项目专属数据卷 chat-agent-shared-kb-pgdata；
#   3. 迁移完成后执行数据基线校验。
# =============================================================================

set -eo pipefail

SOURCE_CONTAINER="agent-regkb"
TARGET_CONTAINER="shuimu-kb-shared"
DB_NAME="regulations"
DB_USER="regkb"

echo "============================================================"
echo ">>> 开始法规知识库流式数据迁移 (解耦自立)"
echo ">>> 源容器: ${SOURCE_CONTAINER} (只读)"
echo ">>> 目标容器: ${TARGET_CONTAINER} (本项目专有)"
echo "============================================================"

# 1. 检查源容器是否运行
if ! docker ps --format '{{.Names}}' | grep -q "^${SOURCE_CONTAINER}$"; then
    echo "[错误] 源容器 ${SOURCE_CONTAINER} 未运行，无法执行数据读取！"
    exit 1
fi

# 2. 检查目标容器是否运行
if ! docker ps --format '{{.Names}}' | grep -q "^${TARGET_CONTAINER}$"; then
    echo "[提示] 目标容器 ${TARGET_CONTAINER} 未启动，尝试启动..."
    docker compose -f docker-compose.shared.yml up -d
fi

# 3. 等待目标容器就绪
echo "[等待] 检查 ${TARGET_CONTAINER} 数据库就绪状态..."
for i in {1..30}; do
    if docker exec -i "${TARGET_CONTAINER}" pg_isready -U "${DB_USER}" -d "${DB_NAME}" >/dev/null 2>&1; then
        echo "[就绪] ${TARGET_CONTAINER} 已接受连接。"
        break
    fi
    sleep 1
    if [ "$i" -eq 30 ]; then
        echo "[错误] 目标数据库连接超时！"
        exit 1
    fi
done

# 4. 执行流式无损导出与导入
echo "[迁移] 正在从 ${SOURCE_CONTAINER} 导出并流式导入到 ${TARGET_CONTAINER}..."
START_TIME=$(date +%s)

docker exec -i "${SOURCE_CONTAINER}" pg_dump -U "${DB_USER}" -d "${DB_NAME}" --clean --if-exists \
    | sed -e '/DROP EXTENSION IF EXISTS vector/d' \
    | docker exec -i "${TARGET_CONTAINER}" psql -U "${DB_USER}" -d "${DB_NAME}" >/dev/null

END_TIME=$(date +%s)
DURATION=$((END_TIME - START_TIME))
echo "[完成] 数据流式导入完成，耗时: ${DURATION} 秒。"

# 5. 数据完整性比对校验
echo "============================================================"
echo ">>> 执行数据完整性基线校验"
echo "============================================================"

SRC_DOC_COUNT=$(docker exec -i "${SOURCE_CONTAINER}" psql -U "${DB_USER}" -d "${DB_NAME}" -t -A -c "SELECT count(*) FROM reg_documents;")
TGT_DOC_COUNT=$(docker exec -i "${TARGET_CONTAINER}" psql -U "${DB_USER}" -d "${DB_NAME}" -t -A -c "SELECT count(*) FROM reg_documents;")

SRC_CHUNK_COUNT=$(docker exec -i "${SOURCE_CONTAINER}" psql -U "${DB_USER}" -d "${DB_NAME}" -t -A -c "SELECT count(*) FROM reg_chunks;")
TGT_CHUNK_COUNT=$(docker exec -i "${TARGET_CONTAINER}" psql -U "${DB_USER}" -d "${DB_NAME}" -t -A -c "SELECT count(*) FROM reg_chunks;")

SRC_NODE_COUNT=$(docker exec -i "${SOURCE_CONTAINER}" psql -U "${DB_USER}" -d "${DB_NAME}" -t -A -c "SELECT count(*) FROM reg_nodes;")
TGT_NODE_COUNT=$(docker exec -i "${TARGET_CONTAINER}" psql -U "${DB_USER}" -d "${DB_NAME}" -t -A -c "SELECT count(*) FROM reg_nodes;")

TGT_DB_SIZE=$(docker exec -i "${TARGET_CONTAINER}" psql -U "${DB_USER}" -d "${DB_NAME}" -t -A -c "SELECT pg_size_pretty(pg_database_size('regulations'));")

echo "法规文档数 (reg_documents): 源=${SRC_DOC_COUNT} | 目标=${TGT_DOC_COUNT}"
echo "法规切片数 (reg_chunks):   源=${SRC_CHUNK_COUNT} | 目标=${TGT_CHUNK_COUNT}"
echo "法规节点数 (reg_nodes):    源=${SRC_NODE_COUNT} | 目标=${TGT_NODE_COUNT}"
echo "目标数据库大小:            ${TGT_DB_SIZE}"

if [[ "${SRC_DOC_COUNT}" == "${TGT_DOC_COUNT}" && "${SRC_CHUNK_COUNT}" == "${TGT_CHUNK_COUNT}" && "${SRC_NODE_COUNT}" == "${TGT_NODE_COUNT}" ]]; then
    echo "============================================================"
    echo ">>> [PASS] 知识库数据无损迁移成功！数据 100% 一致！"
    echo "============================================================"
else
    echo "============================================================"
    echo ">>> [FAIL] 数据校验不一致，请排查导入日志！"
    echo "============================================================"
    exit 1
fi