#!/bin/bash
# 统一恢复脚本 — 支持远程/本地模式
# 用法：
#   bash restore.sh                           # 交互式站点恢复（默认）
#   bash restore.sh 20260708_120000           # 指定时间戳
#   bash restore.sh --db issue issue.db       # 数据库恢复
#   bash restore.sh --full backup.tar.gz      # 灾难恢复
#   bash restore.sh --local                   # 本地模式（直接在服务器上运行）
set -e

BACKUP_DIR="/Users/niuniu/programs/docker/db-backups"
REMOTE_BASE="/root/docker"
MODE="remote"

log()   { echo "[$(date '+%H:%M:%S')] $1"; }
warn()  { echo "[$(date '+%H:%M:%S')] [警告] $1"; }
error() { echo "[$(date '+%H:%M:%S')] [错误] $1"; exit 1; }

# 参数解析
ARGS=(); ACTION="default"
for arg in "$@"; do
  case "$arg" in
    --local)  MODE="local" ;;
    --db)     ACTION="db" ;;
    --full)   ACTION="full" ;;
    --help|-h) ACTION="help" ;;
    *)        ARGS+=("$arg") ;;
  esac
done

# 远程/本地适配
runcmd()   { [[ "$MODE" == "local" ]] && bash -c "$*" || ssh txjp "$*"; }
fetch()    { [[ "$MODE" == "local" ]] && cp "$1" "$2" || scp "txjp:$1" "$2"; }
upload()   { [[ "$MODE" == "local" ]] && cp "$1" "$2" || scp "$1" "txjp:$2"; }

RESTORE_DIR="$(mktemp -d)"

# === 自动发现可恢复站点 ===
PROJECT_ROOT="$(cd "$(dirname "$0")/.." && pwd)"
ALL_SITES="shared"
for dir in "$PROJECT_ROOT"/*-ai/; do
  s=$(basename "$dir"); [[ "$s" == "nginx-proxy-ai" ]] && continue
  [ -f "$dir/docker-compose.yml" ] || continue
  s="${s%-ai}"; ALL_SITES="$ALL_SITES $s"
done

# === 工具函数 ===
check_backup() {
  local site="$1" ts="$2"
  [[ -f "${BACKUP_DIR}/${site}-${ts}.tar.gz" ]] && echo "local:${BACKUP_DIR}/${site}-${ts}.tar.gz" && return
  runcmd "test -f /root/docker/db-backups/${site}-${ts}.tar.gz" 2>/dev/null && echo "cloud:${site}-${ts}.tar.gz" && return
  echo ""
}

download_and_extract() {
  local site="$1" ts="$2"
  local bk=$(check_backup "$site" "$ts")
  [[ -z "$bk" ]] && error "${site} 备份文件不存在: ${ts}"

  rm -rf "$RESTORE_DIR/backup-${site}"; mkdir -p "$RESTORE_DIR/backup-${site}"

  if [[ "$bk" == local:* ]]; then
    log "从本地解压: ${BACKUP_DIR}/${site}-${ts}.tar.gz"
    tar xzf "${BACKUP_DIR}/${site}-${ts}.tar.gz" -C "$RESTORE_DIR/backup-${site}"
  else
    log "从服务器下载: ${ts}"
    fetch "/root/docker/db-backups/${site}-${ts}.tar.gz" "$RESTORE_DIR/"
    tar xzf "$RESTORE_DIR/${site}-${ts}.tar.gz" -C "$RESTORE_DIR/backup-${site}"
  fi
}

restore_src() {
  local src="${RESTORE_DIR}/backup-$1/backup-$1"
  [[ -d "$src" ]] && echo "$src" || echo "${RESTORE_DIR}/backup-$1"
}

# === 站点恢复函数 ===
restore_shared() {
  local src=$(restore_src "shared")
  [[ -f "$src/nginx"/*.conf ]] && { mkdir -p "$src/nginx"; upload "$src/nginx/"* "$REMOTE_BASE/nginx-proxy-ai/conf.d/"; }
  runcmd "docker exec nginx-ai nginx -t && docker exec nginx-ai nginx -s reload" 2>/dev/null || warn "nginx 重载失败"
  log "共享配置恢复完成"
}

restore_assets() {
  log "恢复 assets..."; download_and_extract "assets" "$TIMESTAMP"
  local src=$(restore_src "assets")
  for f in docker-compose.yml .env Dockerfile; do [[ -f "$src/$f" ]] && upload "$src/$f" "$REMOTE_BASE/assets-ai/"; done
  [[ -d "$src/data/uploads" ]] && runcmd "cp -a $src/data/uploads /root/docker/assets-ai/data/"
  [[ -f "$src/data/assets.db" ]] && runcmd "docker cp $src/data/assets.db assets-ai:/app/data/assets.db"
  runcmd "cd $REMOTE_BASE/assets-ai && docker compose down 2>/dev/null || true && docker compose up -d"
  log "assets 恢复完成"
}

restore_issue() {
  log "恢复 issue..."; download_and_extract "issue" "$TIMESTAMP"
  local src=$(restore_src "issue")
  for f in docker-compose.yml .env Dockerfile; do [[ -f "$src/$f" ]] && upload "$src/$f" "$REMOTE_BASE/issue-ai/"; done
  [[ -d "$src/data/uploads" ]] && runcmd "cp -a $src/data/uploads /root/docker/issue-ai/data/"
  [[ -d "$src/data/reports" ]] && runcmd "cp -a $src/data/reports /root/docker/issue-ai/data/"
  [[ -f "$src/data/issue.db" ]] && runcmd "docker cp $src/data/issue.db issue-ai:/app/data/issue.db"
  runcmd "cd $REMOTE_BASE/issue-ai && docker compose down 2>/dev/null || true && docker compose up -d"
  log "issue 恢复完成"
}

restore_ldap() {
  log "恢复 ldap..."; download_and_extract "ldap" "$TIMESTAMP"
  local src=$(restore_src "ldap")
  for f in docker-compose.yml .env Dockerfile; do [[ -f "$src/$f" ]] && upload "$src/$f" "$REMOTE_BASE/ldap-ai/"; done
  [[ -f "$src/data/users.db" ]] && runcmd "docker cp $src/data/users.db lldap:/data/users.db"
  runcmd "cd $REMOTE_BASE/ldap-ai && docker compose down 2>/dev/null || true && docker compose up -d"
  log "ldap 恢复完成"
}

restore_gitea() {
  log "恢复 gitea..."; download_and_extract "gitea" "$TIMESTAMP"
  local src=$(restore_src "gitea")
  [[ -f "$src/docker-compose.yml" ]] && upload "$src/docker-compose.yml" "$REMOTE_BASE/gitea-ai/"
  [[ -f "$src/data/gitea.db" ]] && runcmd "docker cp $src/data/gitea.db gitea-ai:/data/gitea.db"
  runcmd "cd $REMOTE_BASE/gitea-ai && docker compose down && docker compose up -d"
  log "gitea 恢复完成"
}

restore_oa() {
  log "恢复 oa..."; download_and_extract "oa" "$TIMESTAMP"
  local src=$(restore_src "oa")
  for f in docker-compose.yml .env; do [[ -f "$src/$f" ]] && upload "$src/$f" "$REMOTE_BASE/oa-ai/"; done
  runcmd "cd $REMOTE_BASE/oa-ai && docker compose down 2>/dev/null || true && docker compose up -d"
  log "oa 恢复完成"
}

restore_monitor() {
  log "恢复 monitor..."; download_and_extract "monitor" "$TIMESTAMP"
  local src=$(restore_src "monitor")
  for f in docker-compose.yml .env; do [[ -f "$src/$f" ]] && upload "$src/$f" "$REMOTE_BASE/monitor-ai/"; done
  [[ -f "$src/data/monitor.db" ]] && runcmd "docker cp $src/data/monitor.db monitor-ai:/app/data/monitor.db"
  runcmd "cd $REMOTE_BASE/monitor-ai && docker compose down 2>/dev/null || true && docker compose up -d --build"
  log "monitor 恢复完成"
}

restore_generic() {
  local site="$1"; log "恢复 ${site}（通用模式）..."
  download_and_extract "${site}" "$TIMESTAMP"
  local src=$(restore_src "${site}")
  for f in docker-compose.yml .env Dockerfile; do [[ -f "$src/$f" ]] && upload "$src/$f" "$REMOTE_BASE/${site}-ai/"; done
  runcmd "cd $REMOTE_BASE/${site}-ai && docker compose down 2>/dev/null || true && docker compose up -d" 2>/dev/null || warn "${site} 容器重启失败"
}

# === 模式 2：数据库恢复 ===
restore_db() {
  local service="$1" file="$2"
  [[ -z "$service" || -z "$file" ]] && error "用法: bash restore.sh --db <assets|issue|ldap|gitea|monitor> <文件名>"

  local container="" dbpath="" project=""
  case "$service" in
    assets)  project="assets-ai"; container="assets-ai"; dbpath="/app/data/assets.db" ;;
    issue)   project="issue-ai"; container="issue-ai"; dbpath="/app/data/issue.db" ;;
    ldap)    project="ldap-ai"; container="lldap"; dbpath="/data/users.db" ;;
    gitea)   project="gitea-ai"; container="gitea-ai"; dbpath="/data/gitea.db" ;;
    monitor) project="monitor-ai"; container="monitor-ai"; dbpath="/app/data/monitor.db" ;;
    *) error "服务名必须是: assets, issue, ldap, gitea, monitor" ;;
  esac

  local src="$BACKUP_DIR/$file"; [[ ! -f "$src" ]] && src="$file"
  log "恢复 ${service} 数据库: ${src}..."
  upload "$src" "/tmp/db-restore.db"
  runcmd "docker cp /tmp/db-restore.db ${container}:${dbpath} && rm /tmp/db-restore.db"
  runcmd "cd $REMOTE_BASE/$project && docker compose restart" 2>/dev/null || true
  log "${service} 数据库恢复完成"
}

# === 模式 3：灾难恢复 ===
restore_full() {
  local file="$1"; [[ -z "$file" ]] && error "用法: bash restore.sh --full <backup.tar.gz>"
  local src="$BACKUP_DIR/$file"; [[ ! -f "$src" ]] && src="$file"
  [[ ! -f "$src" ]] && error "备份文件不存在: $file"

  log "完整恢复: $file (模式: ${MODE})"
  rm -rf "$RESTORE_DIR/full"; mkdir -p "$RESTORE_DIR/full"
  tar xzf "$src" -C "$RESTORE_DIR/full"
  local content="$RESTORE_DIR/full/full-backup"

  echo ""; echo "  1) assets    2) issue    3) ldap    4) gitea"
  echo "  5) monitor   6) nginx    7) 全部恢复"
  echo ""; printf "请选择 (1-7): "; read choice

  fdb() { [[ -f "$content/databases/$1.db" ]] && upload "$content/databases/$1.db" "/tmp/${1}.db" && runcmd "docker cp /tmp/${1}.db ${2}:/app/data/${1}.db; rm /tmp/${1}.db"; }
  local c="$content"
  case "$choice" in
    1) fdb assets assets-ai; runcmd "cd $REMOTE_BASE/assets-ai && docker compose down && docker compose up -d" ;;
    2) fdb issue issue-ai; runcmd "cd $REMOTE_BASE/issue-ai && docker compose down && docker compose up -d" ;;
    3) [[ -f "$c/databases/ldap.db" ]] && upload "$c/databases/ldap.db" "/tmp/ldap.db"; runcmd "docker cp /tmp/ldap.db lldap:/data/users.db; rm /tmp/ldap.db"; runcmd "cd $REMOTE_BASE/ldap-ai && docker compose down && docker compose up -d" ;;
    4) fdb gitea gitea-ai; runcmd "cd $REMOTE_BASE/gitea-ai && docker compose down && docker compose up -d" ;;
    5) fdb monitor monitor-ai; runcmd "cd $REMOTE_BASE/monitor-ai && docker compose down && docker compose up -d --build" ;;
    6) [[ -d "$c/nginx" ]] && { mkdir -p "$c/nginx"; upload "$c/nginx/"* "$REMOTE_BASE/nginx-proxy-ai/conf.d/"; }; runcmd "docker exec nginx-ai nginx -t && docker exec nginx-ai nginx -s reload" ;;
    7) for s in assets issue ldap gitea monitor; do fdb $s ${s}-ai 2>/dev/null || true; done
       runcmd "docker exec nginx-ai nginx -t && docker exec nginx-ai nginx -s reload"
       for s in assets-ai issue-ai ldap-ai gitea-ai monitor-ai; do runcmd "cd $REMOTE_BASE/$s && docker compose down && docker compose up -d" 2>/dev/null || true; done ;;
    *) error "无效选择" ;;
  esac
  log "灾难恢复完成"
}

# === 主流程 ===
case "$ACTION" in
  help)
    echo "用法: bash restore.sh [--local] [--db|--full] [参数...]"
    echo "  --local              本地模式（直接在服务器上运行）"
    echo "  无参数                交互式站点恢复"
    echo "  restore.sh TIMESTAMP  指定时间戳恢复"
    echo "  --db <site> <file>    数据库恢复"
    echo "  --full <file>         灾难恢复"
    ;;
  db)   restore_db "${ARGS[0]}" "${ARGS[1]}" ;;
  full) restore_full "${ARGS[0]}" ;;
  *)
    TIMESTAMP="${ARGS[0]:-}"
    if [[ -z "$TIMESTAMP" ]]; then
      echo ""; echo "可用的备份时间戳:"; echo ""
      timestamps=$(ls ${BACKUP_DIR}/assets-*.tar.gz 2>/dev/null | sed 's/.*assets-\(.*\)\.tar\.gz/\1/' | sort -r | head -10)
      [[ -z "$timestamps" ]] && error "未找到任何备份文件"
      declare -A ts_list; index=1
      while IFS= read -r ts; do echo "  $index) $ts"; ts_list[$index]="$ts"; ((index++)); done <<< "$timestamps"
      echo ""; printf "请选择 (1-$((index-1))): "; read choice
      TIMESTAMP="${ts_list[$choice]}"
      [[ -z "$TIMESTAMP" ]] && error "无效选择"
    fi

    echo ""; echo "可恢复的站点:"; echo ""
    declare -A site_list; index=1
    for site in $ALL_SITES; do
      bk=$(check_backup "$site" "$TIMESTAMP")
      [[ -n "$bk" ]] && echo "  $index) $site" && site_list[$index]="$site" || echo "  -) $site (无备份)"
      ((index++))
    done
    echo ""; echo "  a) 全部恢复"; printf "请选择 (1-$((index-1)) 或 a): "; read s_choice

    declare -a restore_list
    if [[ "$s_choice" == "a" ]]; then
      for idx in "${!site_list[@]}"; do restore_list+=("${site_list[$idx]}"); done
    else
      restore_list=("${site_list[$s_choice]}")
      [[ -z "${restore_list[0]}" ]] && error "无效选择或无备份"
    fi

    for site in "${restore_list[@]}"; do
      case "$site" in
        shared)  restore_shared ;;
        assets)  restore_assets ;;
        issue)   restore_issue ;;
        ldap)    restore_ldap ;;
        gitea)   restore_gitea ;;
        oa)      restore_oa ;;
        monitor) restore_monitor ;;
        *)       restore_generic "$site" ;;
      esac
    done

    runcmd "docker exec nginx-ai nginx -t && docker exec nginx-ai nginx -s reload" 2>/dev/null || true
    rm -rf "$RESTORE_DIR"
    echo ""; log "恢复完成！"
    runcmd "docker ps --format 'table {{.Names}}\t{{.Status}}' | grep -E '(assets|issue|ldap|gitea|oa|monitor|nginx|sso)'"
    ;;
esac