#!/usr/bin/env bash
#
# single-db-restore-enterprise.sh (v4 - detailed colors + logs + progress + import details)
#
# 用法:
#   ./single-db-restore-enterprise.sh <ABS_PATH> <USER> <PASS> [DEST_DB]
# 参数:
#   ABS_PATH : 绝对路径，指向 .archive[.gz] 文件，或包含此类文件的目录（目录时自动选最新）
#   USER/PASS: 认证凭证（认证库见 AUTH_DB）
#   DEST_DB  : 可选。不同名恢复时提供；同名恢复可省略
#
set -euo pipefail

# --- Configuration (env-overridable) ---
HOST="${HOST:-127.0.0.1}"
PORT="${PORT:-27017}"
AUTH_DB="${AUTH_DB:-admin}"
STOP_ON_ERROR="${STOP_ON_ERROR:-1}"          # 1=遇错停止
DROP_BEFORE_RESTORE="${DROP_BEFORE_RESTORE:-1}"   # 1=恢复前drop
JOBS="${JOBS:-$(nproc || echo 2)}"
RESTORE_TIMEOUT="${RESTORE_TIMEOUT:-0}"      # 秒。0=不超时
RESTORE_STALL_SECS="${RESTORE_STALL_SECS:-0}"   # 秒。0=不监测卡滞（无增长）
LOG_DIR="${LOG_DIR:-/var/log}"

# --- Style & Colors ---
CLR_RESET="\x1b[0m"; BOLD="\x1b[1m"
CLR_GRY="\x1b[90m"; CLR_SKY="\x1b[96m"
CLR_GRN="\x1b[32m"; CLR_YLW="\x1b[33m"
CLR_RED="\x1b[31m"

HOSTNAME_SHORT="$(hostname -s || echo 'localhost')"
ts(){ date +"%Y-%m-%d %H:%M:%S%z"; }
log_line(){ local lvl="$1";shift;local c="$1";shift;printf "%b%s [%s] [%s] %s%b\n" "$c" "$(ts)" "$HOSTNAME_SHORT" "$lvl" "$*" "$CLR_RESET"; }
info(){ log_line "INFO"  "$CLR_SKY" "$*"; }
ok(){   log_line "OK"    "$CLR_GRN" "$*"; }
warn(){ log_line "WARN"  "$CLR_YLW" "$*"; }
err(){  log_line "ERROR" "$CLR_RED" "$*"; }

print_header(){
  printf "\n%b╭──────────────────────────────────────────────────────────────────────────────╮%b\n" "$CLR_GRY" "$CLR_RESET"
  printf "%b│ %b%-75s%b │%b\n" "$CLR_GRY" "${BOLD}${CLR_SKY}" "$1" "$CLR_GRY" "$CLR_RESET"
  printf "%b╰──────────────────────────────────────────────────────────────────────────────╯%b\n" "$CLR_GRY" "$CLR_RESET"
}
print_kv(){ printf "  %b%-24s %b»%b %s\n" "$CLR_SKY" "$1" "$CLR_GRY" "$CLR_RESET" "$2"; }

need_tool(){
  local bin="$1" pkg="$2"
  if ! command -v "$bin" >/dev/null 2>&1; then
    warn "缺少 $bin，尝试安装 $pkg（需 root）。"
    if [[ $EUID -ne 0 ]]; then err "非 root 无法自动安装，请执行: apt-get update && apt-get install -y $pkg"; exit 2; fi
    apt-get update && apt-get install -y "$pkg"
  fi
}

format_bytes(){
  local b=$1
  if (( b<1024 )); then echo "${b} B"
  elif (( b<1024*1024 )); then printf "%.2f KB\n" "$(bc <<<"scale=2;$b/1024")"
  elif (( b<1024*1024*1024 )); then printf "%.2f MB\n" "$(bc <<<"scale=2;$b/1024/1024")"
  else printf "%.2f GB\n" "$(bc <<<"scale=2;$b/1024/1024/1024")"
  fi
}

# --- Args ---
if [[ $# -lt 3 || $# -gt 4 ]]; then err "用法: $0 <ABS_PATH> <USER> <PASS> [DEST_DB]"; exit 1; fi
ABS_PATH="$1"; USER="$2"; PASS="$3"; DEST_DB="${4:-}"

# --- Preflight ---
[[ -w "$LOG_DIR" ]] || LOG_DIR="/tmp"
print_header "环境检查 / PRE-FLIGHT CHECKS"
if [[ -z "$ABS_PATH" || -z "$USER" || -z "$PASS" ]]; then err "参数不能为空。"; exit 1; fi
if [[ "${ABS_PATH}" != /* ]]; then err "参数1必须为绝对路径。"; exit 1; fi
need_tool "mongorestore" "mongodb-database-tools"
need_tool "mongosh"      "mongodb-mongosh"
need_tool "sha256sum"    "coreutils"
need_tool "bc"           "bc"
ok "工具就绪。"

# --- 定位归档 ---
print_header "定位备份文件 / LOCATE BACKUP ARTIFACT"
resolve_archive(){
  local p="$1"
  if [[ -f "$p" ]]; then echo "$p"
  elif [[ -d "$p" ]]; then
    local f
    f="$(find "$p" -maxdepth 1 -type f \( -name '*.archive.gz' -o -name '*.archive' \) -printf '%T@ %p\n' | sort -nr | awk 'NR==1{print $2}')"
    [[ -n "$f" ]] && echo "$f" || return 1
  else return 1
  fi
}
ARCHIVE="$(resolve_archive "$ABS_PATH" || true)"
[[ -n "${ARCHIVE:-}" ]] || { err "未找到有效的备份文件: $ABS_PATH"; exit 1; }

FILE_NAME="$(basename -- "$ARCHIVE")"
FILE_SIZE="$(stat -c '%s' "$ARCHIVE")"
FILE_MODE="$(stat -c '%A' "$ARCHIVE")"
SHA256="$(sha256sum "$ARCHIVE" | awk '{print $1}')"
GZIP_OPT=""; [[ "$ARCHIVE" == *.gz ]] && GZIP_OPT="--gzip"

print_kv "备份文件" "$ARCHIVE"
print_kv "文件大小" "$(format_bytes "$FILE_SIZE") (${FILE_SIZE} B)"
print_kv "文件权限" "$FILE_MODE"
print_kv "SHA256"   "$SHA256"

# 从文件名推断源库名: <DB>-YYYYmmdd-HHMMSS.archive[.gz]
SRC_DB=""; BN="${FILE_NAME%.archive.gz}"; BN="${BN%.archive}"
[[ "$BN" =~ ^(.+)-([0-9]{8})-([0-9]{6})$ ]] && SRC_DB="${BASH_REMATCH[1]}"
[[ -n "$SRC_DB" ]] && print_kv "推断的源库名" "$SRC_DB" || warn "无法从文件名推断源库名，将按归档内原名恢复。"

# --- 连接性 ---
print_header "连接性检查 / CONNECTIVITY CHECK"
if ! mongosh --quiet --host "$HOST" --port "$PORT" --username "$USER" --password "$PASS" --authenticationDatabase "$AUTH_DB" --eval 'db.runCommand({ ping: 1 })' >/dev/null 2>&1; then
  err "无法连接或认证失败。"; exit 3
fi
ok "数据库连接成功。"

# --- 恢复计划 ---
print_header "恢复计划 / RESTORE PLAN"
RESTORE_USERS_AND_ROLES=0
MAP_ARGS=()
TARGET_DB_MSG="原库名（归档内）"

if [[ -n "$DEST_DB" ]]; then
  if [[ -n "$SRC_DB" && "$DEST_DB" != "$SRC_DB" ]]; then
    MAP_ARGS+=(--nsFrom "${SRC_DB}.*" --nsTo "${DEST_DB}.*")
    TARGET_DB_MSG="$DEST_DB（命名空间映射）"
    RESTORE_USERS_AND_ROLES=0  # 重命名恢复默认不恢复用户与角色
    warn "重命名恢复：${SRC_DB} → ${DEST_DB}。默认不恢复用户与角色。"
  else
    TARGET_DB_MSG="$DEST_DB"; RESTORE_USERS_AND_ROLES=1
  fi
else
  TARGET_DB_MSG="${SRC_DB:-'归档内原名'}"; RESTORE_USERS_AND_ROLES=1
fi

print_kv "目标数据库" "$TARGET_DB_MSG"
print_kv "Drop 既有数据" "$([[ "$DROP_BEFORE_RESTORE" -eq 1 ]] && echo yes || echo no)"
print_kv "StopOnError" "$([[ "$STOP_ON_ERROR" -eq 1 ]] && echo yes || echo no)"
print_kv "并发线程"   "$JOBS"
print_kv "恢复用户与角色" "$([[ "$RESTORE_USERS_AND_ROLES" -eq 1 ]] && echo yes || echo no)"

# --- 快照（恢复前总文档数，用于进度与增量） ---
VERIFY_DB="$DEST_DB"; [[ -z "$VERIFY_DB" ]] && VERIFY_DB="$SRC_DB"
SNAP_JS=$(cat <<'JS'
(function(){
  const d=db.getSiblingDB("@@DB@@");
  let total=0;
  try{
    let names=[];
    try{ names=d.getCollectionInfos({ name:{ $not:/^system\./ }}).filter(ci=>ci.type==="collection").map(ci=>ci.name); }
    catch(e){ names=d.getCollectionNames().filter(n=>!n.startsWith("system.")); }
    names.forEach(n=>{ try{ total+=d.getCollection(n).estimatedDocumentCount(); }catch(e){} });
  }catch(e){}
  print(total);
})();
JS
)

PRE_DOCS=""
if [[ -n "$VERIFY_DB" ]]; then
  print_header "采集恢复前快照 / SNAPSHOT BEFORE RESTORE"
  PRE_DOCS="$(mongosh --quiet --host "$HOST" --port "$PORT" --username "$USER" --password "$PASS" --authenticationDatabase "$AUTH_DB" --eval "${SNAP_JS//@@DB@@/${VERIFY_DB}}" 2>/dev/null || true)"
  [[ -n "$PRE_DOCS" ]] && ok "前快照：总文档数=$PRE_DOCS" || warn "前快照为空（新库或权限不足）。"
else
  warn "无法确定目标库名，进度与增量统计将不可用。"
fi

# --- 进度监控 ---
progress_monitor(){
  local pid="$1"; local db="$2"; local stall="${3:-0}"
  local last_total=""; local last_ts="$(date +%s)"
  while kill -0 "$pid" >/dev/null 2>&1; do
    if [[ -n "$db" ]]; then
      local t
      t="$(mongosh --quiet --host "$HOST" --port "$PORT" --username "$USER" --password "$PASS" --authenticationDatabase "$AUTH_DB" --eval "${SNAP_JS//@@DB@@/${db}}" 2>/dev/null || true)"
      if [[ -n "$t" ]]; then
        printf "%b%s [%s] [INFO] 进度: %s 文档%b\n" "$CLR_SKY" "$(ts)" "$HOSTNAME_SHORT" "$t" "$CLR_RESET"
        if [[ "$t" != "$last_total" ]]; then last_total="$t"; last_ts="$(date +%s)"; fi
      else
        printf "%b%s [%s] [INFO] 进度: N/A（权限不足或库尚未创建）%b\n" "$CLR_SKY" "$(ts)" "$HOSTNAME_SHORT" "$CLR_RESET"
      fi
    else
      printf "%b%s [%s] [INFO] 心跳: mongorestore 运行中%b\n" "$CLR_SKY" "$(ts)" "$HOSTNAME_SHORT" "$CLR_RESET"
    fi
    if (( stall>0 )); then
      local now_ts="$(date +%s)"
      if (( now_ts - last_ts > stall )); then
        printf "%b%s [%s] [ERROR] 长时间无进度，判定卡滞（>%ss）。%b\n" "$CLR_RED" "$(ts)" "$HOSTNAME_SHORT" "$stall" "$CLR_RESET"
        kill "$pid" 2>/dev/null || true
        return 124
      fi
    fi
    sleep 15
  done
  return 0
}

# --- 执行恢复 ---
print_header "执行恢复 / EXECUTION"
STAMP="$(date +%Y%m%d-%H%M%S)"
LOG_FILE="${LOG_DIR}/mongorestore_${HOSTNAME_SHORT}_${STAMP}.log"

RESTORE_CMD=(mongorestore
  --host "$HOST" --port "$PORT"
  --username "$USER" --password "$PASS"
  --authenticationDatabase "$AUTH_DB"
  --archive="$ARCHIVE" "$GZIP_OPT"
  --numInsertionWorkersPerCollection "$JOBS"
  --verbose
)

# 命名空间映射（重命名恢复）
[[ ${#MAP_ARGS[@]} -gt 0 ]] && RESTORE_CMD+=("${MAP_ARGS[@]}")
# 常规开关
[[ "$DROP_BEFORE_RESTORE" -eq 1 ]] && RESTORE_CMD+=(--drop)
[[ "$STOP_ON_ERROR" -eq 1 ]] && RESTORE_CMD+=(--stopOnError)

# 关键修复：恢复用户与角色时必须加 --db
if [[ "$RESTORE_USERS_AND_ROLES" -eq 1 ]]; then
  if [[ -z "$VERIFY_DB" ]]; then
    err "使用 --restoreDbUsersAndRoles 需要明确数据库名。请提供参数4或使用可解析源库名的归档文件名。"
    exit 1
  fi
  RESTORE_CMD+=(--db "$VERIFY_DB" --restoreDbUsersAndRoles)
fi

# 展示安全命令（隐藏密码）
SAFE_CMD=("${RESTORE_CMD[@]}"); for i in "${!SAFE_CMD[@]}"; do [[ "${SAFE_CMD[$i]}" == "$PASS" ]] && SAFE_CMD[$i]="******"; done
print_kv "命令" "${SAFE_CMD[*]}"
print_kv "日志文件" "$LOG_FILE"
(( RESTORE_TIMEOUT>0 )) && print_kv "超时" "${RESTORE_TIMEOUT}s"

# 可选超时包装
RUN_WRAP=()
if (( RESTORE_TIMEOUT>0 )); then RUN_WRAP=(timeout -k 10s "$RESTORE_TIMEOUT"); fi

set +e
("${RUN_WRAP[@]}" "${RESTORE_CMD[@]}" >"$LOG_FILE" 2>&1) & RESTORE_PID=$!
progress_monitor "$RESTORE_PID" "$VERIFY_DB" "$RESTORE_STALL_SECS" & MON_PID=$!
wait "$RESTORE_PID"; RET=$?
kill "$MON_PID" >/dev/null 2>&1 || true
set -e

if [[ $RET -ne 0 ]]; then
  err "恢复失败，退出码: $RET"
  printf "%b--- LOG TAIL (last 100 lines) ---%b\n" "$CLR_RED" "$CLR_RESET"
  tail -n 100 "$LOG_FILE" | sed "s/^/  /"
  exit $RET
fi
ok "恢复完成。详见日志：$LOG_FILE"

# --- 导入集合详情 ---
print_header "导入集合详情 / IMPORTED COLLECTION DETAILS"
DETAIL_JS=$(cat <<'JS'
(function(){
  const d=db.getSiblingDB("@@DB@@");
  function safe(fn, def){ try{ return fn(); }catch(e){ return def; } }
  let names=[];
  try{ names=d.getCollectionInfos({ name:{ $not:/^system\./ }}).filter(ci=>ci.type==="collection").map(ci=>ci.name); }
  catch(e){ names=safe(()=>d.getCollectionNames().filter(n=>!n.startsWith("system.")), []); }
  names.forEach(n=>{
    const cnt=safe(()=>d.getCollection(n).estimatedDocumentCount(), 0);
    const idx=safe(()=>d.getCollection(n).getIndexes().length, 0);
    const cs =safe(()=>d.runCommand({collStats:n,scale:1}), {ok:0});
    const stor=(cs&&cs.ok===1)?(cs.storageSize||0):0;
    const size=(cs&&cs.ok===1)?(cs.size||0):0;
    print(n+"\t"+cnt+"\t"+stor+"\t"+idx+"\t"+size);
  });
})();
JS
)

TOTAL_DOCS=0; TOTAL_STOR=0
if [[ -n "$VERIFY_DB" ]]; then
  POST_OUT="$(mongosh --quiet --host "$HOST" --port "$PORT" --username "$USER" --password "$PASS" --authenticationDatabase "$AUTH_DB" --eval "${DETAIL_JS//@@DB@@/${VERIFY_DB}}" 2>/dev/null || true)"
  printf "  %b%-42s %-14s %-12s %-14s %-8s%b\n" "$CLR_YLW" "Collection" "Docs(After)" "ΔDocs" "Storage(After)" "Indexes" "$CLR_RESET"
  printf "  %b%s%b\n" "$CLR_GRY" "-------------------------------------------------------------------------------------------------" "$CLR_RESET"
  if [[ -n "$POST_OUT" ]]; then
    while IFS=$'\t' read -r n pc ps pi _sz; do
      [[ -z "${n:-}" ]] && continue
      delta="N/A"; [[ -n "$PRE_DOCS" && "$PRE_DOCS" =~ ^[0-9]+$ ]] && delta="~$(($pc))" # 前快照为总量，逐集合增量无法逐一精确，展示近似
      printf "  %-42s %-14s %-12s %-14s %-8s\n" "$n" "$pc" "$delta" "$(format_bytes "$ps")" "$pi"
      TOTAL_DOCS=$((TOTAL_DOCS + pc)); TOTAL_STOR=$((TOTAL_STOR + ps))
    done < <(printf "%s" "$POST_OUT" | awk -F'\t' '{print $1"\t"$2"\t"$3"\t"$4"\t"$5}' | sort -t$'\t' -k3,3nr)
  else
    warn "无法获取集合详情（权限不足或库为空）。"
  fi
else
  warn "未确定目标库名，跳过集合详情。"
fi

printf "\n  %b📊 汇总%b\n" "${BOLD}${CLR_GRN}" "$CLR_RESET"
print_kv "目标数据库" "${VERIFY_DB:-未知}"
[[ -n "$PRE_DOCS" && "$PRE_DOCS" =~ ^[0-9]+$ ]] && print_kv "文档增量(≈导入量)" "$((TOTAL_DOCS - PRE_DOCS))"
print_kv "总文档数(恢复后)" "$TOTAL_DOCS"
print_kv "总存储(恢复后)" "$(format_bytes "$TOTAL_STOR")"

print_header "结果验证 / VERIFICATION"
[[ $TOTAL_DOCS -ge 0 ]] && ok "验证完成：已统计集合、文档、索引与存储。" || warn "验证信息不足，请人工复核。"

printf "\n%b╭──────────────────────────────────────────────────────────────────────────────╮%b\n" "$CLR_GRN" "$CLR_RESET"
printf "%b│ %b%-75s%b │%b\n" "$CLR_GRN" "$BOLD" "SUCCESS: Restore completed with logs, progress and import details." "$CLR_GRN" "$CLR_RESET"
printf "%b╰──────────────────────────────────────────────────────────────────────────────╯%b\n" "$CLR_GRN" "$CLR_RESET"
exit 0