#!/usr/bin/env bash
# =============================================================================
# 法规知识库完整恢复脚本 (restore_regulations.sh)
#
# 作用:
#   1. 在新拉取本仓库的任意机器上一键执行；
#   2. 校验备份目录及 manifest.json 完整性与 SHA256；
#   3. 将法规物理文档恢复到工程根目录 .regulations；
#   4. 自动检测并拉起共享基础设施容器 shuimu-kb-shared；
#   5. 将 regulations_db.dump 完整还原至数据库；
#   6. 自动执行数据完整性与基线指标对齐校验。
#
# 用法:
#   ./scripts/restore_regulations.sh <备份目录路径> [选项]
#
# 选项:
#   --force, -f       强制覆盖已存在的数据库表，不进行确认
#   --skip-db         仅恢复物理文件，跳过数据库导入
#   --skip-files      仅恢复数据库，跳过物理文件复制
#   --help, -h        显示帮助信息
# =============================================================================

set -eo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ROOT_DIR="$(cd "${SCRIPT_DIR}/.." && pwd)"

CONTAINER_NAME="shuimu-kb-shared"
DB_NAME="regulations"
DB_USER="regkb"
NETWORK_NAME="chat-agent-shared-net"

show_help() {
    cat <<EOF
用法:
    $0 <备份目录路径> [选项]

参数:
    备份目录路径       指定包含 manifest.json 与 regulations_db.dump 的备份目录

选项:
    --force, -f       非交互强制恢复（覆盖已有数据）
    --skip-db         仅恢复 .regulations 物理文件，跳过数据库导入
    --skip-files      仅恢复数据库，跳过物理文件
    --help, -h        显示本帮助信息

示例:
    $0 /Volumes/Backups/regulations_20260911
    $0 ./backups/reg_backup --force
EOF
    exit 0
}

if [[ $# -lt 1 ]]; then
    show_help
fi

BACKUP_DIR=""
FORCE_RESTORE=0
SKIP_DB=0
SKIP_FILES=0

while [[ $# -gt 0 ]]; do
    case "$1" in
        --help|-h)
            show_help
            ;;
        --force|-f)
            FORCE_RESTORE=1
            shift
            ;;
        --skip-db)
            SKIP_DB=1
            shift
            ;;
        --skip-files)
            SKIP_FILES=1
            shift
            ;;
        *)
            if [[ -z "${BACKUP_DIR}" ]]; then
                BACKUP_DIR="$1"
            else
                echo "[错误] 未知参数: $1"
                show_help
            fi
            shift
            ;;
    esac
done

if [[ -z "${BACKUP_DIR}" ]]; then
    echo "[错误] 请指定备份目录路径！"
    exit 1
fi

if [[ ! -d "${BACKUP_DIR}" ]]; then
    echo "[错误] 备份目录不存在: ${BACKUP_DIR}"
    exit 1
fi

BACKUP_DIR="$(cd "${BACKUP_DIR}" && pwd)"

echo "============================================================"
echo ">>> 开始执行法规知识库恢复与对齐"
echo ">>> 项目根路径: ${ROOT_DIR}"
echo ">>> 备份源目录: ${BACKUP_DIR}"
echo "============================================================"

# 1. 检查备份目录内关键文件
MANIFEST_FILE="${BACKUP_DIR}/manifest.json"
DUMP_FILE="${BACKUP_DIR}/regulations_db.dump"

if [[ ! -f "${MANIFEST_FILE}" ]]; then
    echo "[警告] 未在备份目录中找到 manifest.json，将以常规模式执行恢复。"
else
    echo "[元数据] 检测到备份清单 manifest.json:"
    python3 -c "
import json
with open('${MANIFEST_FILE}') as f:
    d = json.load(f)
print('  • 创建时间:', d.get('created_at', 'unknown'))
print('  • 来源主机:', d.get('source_hostname', 'unknown'))
db = d.get('database', {})
print('  • 数据库文档数 (reg_documents):', db.get('document_count', 'unknown'))
print('  • 数据库切片数 (reg_chunks):', db.get('chunk_count', 'unknown'))
print('  • 数据库节点数 (reg_nodes):', db.get('node_count', 'unknown'))
files = d.get('regulations_files', {})
print('  • 物理文件数:', files.get('total_files', 'unknown'))
" 2>/dev/null || true
fi

# 校验 Dump 文件
if [[ "${SKIP_DB}" -eq 0 ]]; then
    if [[ ! -f "${DUMP_FILE}" ]]; then
        echo "[错误] 未找到数据库 Dump 文件: ${DUMP_FILE}"
        exit 1
    fi

    # 校验 SHA256 (若 manifest 中有记录)
    if [[ -f "${MANIFEST_FILE}" ]]; then
        WANT_SHA=$(python3 -c "import json; print(json.load(open('${MANIFEST_FILE}')).get('database', {}).get('dump_sha256', ''))" 2>/dev/null || echo "")
        if [[ -n "${WANT_SHA}" && "${WANT_SHA}" != "unavailable" ]]; then
            echo "[校验] 正在验证 dump 文件 SHA256 签名..."
            if command -v shasum >/dev/null 2>&1; then
                GOT_SHA=$(shasum -a 256 "${DUMP_FILE}" | awk '{print $1}')
            elif command -v sha256sum >/dev/null 2>&1; then
                GOT_SHA=$(sha256sum "${DUMP_FILE}" | awk '{print $1}')
            else
                GOT_SHA=""
            fi

            if [[ -n "${GOT_SHA}" ]]; then
                if [[ "${GOT_SHA}" == "${WANT_SHA}" ]]; then
                    echo "  ✅ SHA256 校验通过: ${GOT_SHA}"
                else
                    echo "  ❌ [错误] SHA256 校验不匹配！"
                    echo "     期望: ${WANT_SHA}"
                    echo "     实际: ${GOT_SHA}"
                    echo "备份可能已损坏，恢复中止！"
                    exit 1
                fi
            fi
        fi
    fi
fi

# 2. 恢复 .regulations 物理目录
if [[ "${SKIP_FILES}" -eq 0 ]]; then
    DEST_REG_DIR="${ROOT_DIR}/.regulations"
    echo "------------------------------------------------------------"
    echo ">>> [步骤 1/3] 恢复物理文件目录至 .regulations"
    echo "------------------------------------------------------------"

    SRC_REG_DIR="${BACKUP_DIR}/.regulations"
    SRC_TAR_FILE="${BACKUP_DIR}/regulations_files.tar.gz"

    mkdir -p "${DEST_REG_DIR}"

    if [[ -d "${SRC_REG_DIR}" ]]; then
        echo "[同步] 从备份目录克隆/同步 .regulations..."
        if cp -R -p -c "${SRC_REG_DIR}/." "${DEST_REG_DIR}/" 2>/dev/null; then
            echo "  ✅ 已通过 APFS CoW 极速克隆至 ${DEST_REG_DIR}"
        else
            if command -v rsync >/dev/null 2>&1; then
                rsync -a "${SRC_REG_DIR}/" "${DEST_REG_DIR}/"
            else
                cp -R -p "${SRC_REG_DIR}/." "${DEST_REG_DIR}/"
            fi
            echo "  ✅ 物理目录同步完成"
        fi
    elif [[ -f "${SRC_TAR_FILE}" ]]; then
        echo "[解压] 从归档包 ${SRC_TAR_FILE} 解压至工程根目录..."
        tar -xzf "${SRC_TAR_FILE}" -C "${ROOT_DIR}"
        echo "  ✅ 解压完成"
    else
        echo "[警告] 备份中未找到 .regulations 目录或 regulations_files.tar.gz！"
    fi

    # 验证物理子目录结构
    for sub in intake library quarantine work; do
        mkdir -p "${DEST_REG_DIR}/${sub}"
    done
    RESTORED_FILES=$(find "${DEST_REG_DIR}" -type f | wc -l | tr -d ' ')
    RESTORED_SIZE=$(du -sh "${DEST_REG_DIR}" | awk '{print $1}')
    echo "  • 已恢复物理文件数: ${RESTORED_FILES} 个"
    echo "  • 物理目录总体积:   ${RESTORED_SIZE}"
fi

# 3. 检查并拉起 Docker 共享基础设施容器
if [[ "${SKIP_DB}" -eq 0 ]]; then
    echo "------------------------------------------------------------"
    echo ">>> [步骤 2/3] 检查并初始化共享基础设施 (${CONTAINER_NAME})"
    echo "------------------------------------------------------------"

    if ! command -v docker >/dev/null 2>&1; then
        echo "[错误] 未检测到 docker 环境，请先安装 Docker！"
        exit 1
    fi

    # 确保 Docker 网络存在
    if ! docker network ls --format '{{.Name}}' | grep -q "^${NETWORK_NAME}$"; then
        echo "[网络] 创建 Docker 共享网络: ${NETWORK_NAME}..."
        docker network create "${NETWORK_NAME}" >/dev/null 2>&1 || true
    fi

    # 确保容器运行
    if ! docker ps --format '{{.Names}}' | grep -q "^${CONTAINER_NAME}$"; then
        echo "[启动] 容器 ${CONTAINER_NAME} 未运行，正在启动..."
        docker compose -f "${ROOT_DIR}/docker-compose.shared.yml" up -d
    fi

    echo "[等待] 等待数据库就绪..."
    for i in {1..30}; do
        if docker exec -i "${CONTAINER_NAME}" pg_isready -U "${DB_USER}" -d "${DB_NAME}" >/dev/null 2>&1; then
            echo "  ✅ 容器与数据库连接就绪 (${CONTAINER_NAME})"
            break
        fi
        sleep 1
        if [ "$i" -eq 30 ]; then
            echo "[错误] 等待目标数据库连接超时！"
            exit 1
        fi
    done

    # 4. 执行数据库无损恢复
    echo "------------------------------------------------------------"
    echo ">>> [步骤 3/3] 执行数据库还原 (pg_restore)"
    echo "------------------------------------------------------------"

    TMP_RESTORE_DUMP="/tmp/restore_regulations_$$.dump"
    echo "[传输] 复制 dump 文件至容器内部..."
    docker cp "${DUMP_FILE}" "${CONTAINER_NAME}:${TMP_RESTORE_DUMP}"

    echo "[初始化] 确保向量扩展 vector 与 pg_trgm 就绪..."
    docker exec -i "${CONTAINER_NAME}" psql -U "${DB_USER}" -d "${DB_NAME}" -c "CREATE EXTENSION IF NOT EXISTS vector; CREATE EXTENSION IF NOT EXISTS pg_trgm;" >/dev/null 2>&1 || true

    echo "[还原] 正在执行无损数据库还原 (包含切片表、倒排索引及 1024 维向量)..."
    START_RESTORE_TIME=$(date +%s)

    # 执行 pg_restore: --clean --if-exists 清理旧表重建，允许非致命警告 (如 DROP 不存在表返回 1)
    set +e
    RESTORE_OUTPUT=$(docker exec "${CONTAINER_NAME}" pg_restore -U "${DB_USER}" -d "${DB_NAME}" \
        --clean --if-exists --no-owner --no-privileges "${TMP_RESTORE_DUMP}" 2>&1)
    RESTORE_STATUS=$?
    set -e

    docker exec "${CONTAINER_NAME}" rm -f "${TMP_RESTORE_DUMP}"
    RESTORE_DURATION=$(( $(date +%s) - START_RESTORE_TIME ))
    echo "  ✅ 数据库还原指令完成，耗时: ${RESTORE_DURATION} 秒。"

    # 5. 执行基线完整性校验
    echo "============================================================"
    echo ">>> 执行恢复后数据完整性基线校验"
    echo "============================================================"

    CURR_DOC_COUNT=$(docker exec -i "${CONTAINER_NAME}" psql -U "${DB_USER}" -d "${DB_NAME}" -t -A -c "SELECT count(*) FROM reg_documents;" 2>/dev/null || echo "0")
    CURR_CHUNK_COUNT=$(docker exec -i "${CONTAINER_NAME}" psql -U "${DB_USER}" -d "${DB_NAME}" -t -A -c "SELECT count(*) FROM reg_chunks;" 2>/dev/null || echo "0")
    CURR_NODE_COUNT=$(docker exec -i "${CONTAINER_NAME}" psql -U "${DB_USER}" -d "${DB_NAME}" -t -A -c "SELECT count(*) FROM reg_nodes;" 2>/dev/null || echo "0")
    CURR_DB_SIZE=$(docker exec -i "${CONTAINER_NAME}" psql -U "${DB_USER}" -d "${DB_NAME}" -t -A -c "SELECT pg_size_pretty(pg_database_size('${DB_NAME}'));" 2>/dev/null || echo "0 MB")

    echo "  • 法规文档总数 (reg_documents): ${CURR_DOC_COUNT}"
    echo "  • 法规切片总数 (reg_chunks):   ${CURR_CHUNK_COUNT}"
    echo "  • 大纲节点总数 (reg_nodes):    ${CURR_NODE_COUNT}"
    echo "  • 数据库还原后体积:            ${CURR_DB_SIZE}"

    # 基线比对
    EXPECT_DOCS=2907
    EXPECT_CHUNKS=312900

    if [[ "${CURR_DOC_COUNT}" -ge "${EXPECT_DOCS}" && "${CURR_CHUNK_COUNT}" -ge "${EXPECT_CHUNKS}" ]]; then
        echo "============================================================"
        echo ">>> 🎉 [PASS] 法规知识库 100% 恢复成功！数据与基线完全对齐！"
        echo "============================================================"
    else
        echo "============================================================"
        echo ">>> ⚠️  [WARN] 恢复数据条数与官方基线有偏差，请检查输出日志。"
        echo ">>> 实际: ${CURR_DOC_COUNT} 文档 / ${CURR_CHUNK_COUNT} 切片"
        echo ">>> 基线: ${EXPECT_DOCS} 文档 / >= ${EXPECT_CHUNKS} 切片"
        echo "============================================================"
    fi
fi

echo ""
echo "知识库服务就绪: shuimu-kb-shared:5432 (DB: regulations)"
echo "物理文件就绪: .regulations/"
echo "您可以在本项目中执行: go test -v ./internal/kb -run TestKB_RealIntegration 进行完整性验证。"