#!/bin/bash
# 统一恢复脚本
# 用法：
#   bash restore.sh                           # 交互式站点恢复（默认）
#   bash restore.sh 20260708_120000           # 指定时间戳恢复
#   bash restore.sh --db issue issue.db       # 数据库恢复
#   bash restore.sh --full full-backup.tar.gz # 灾难恢复
set -e

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

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

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
  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
  ssh txjp "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 backup_info=$(check_backup "$site" "$ts")
  [[ -z "$backup_info" ]] && error "${site} 备份文件不存在: ${ts}"

  local src="${backup_info#*:}"
  rm -rf "$RESTORE_DIR/backup-${site}"; mkdir -p "$RESTORE_DIR/backup-${site}"

  if [[ "$backup_info" == local:* ]]; then
    log "从本地解压: ${BACKUP_DIR}/${site}-${ts}.tar.gz"
    tar xzf "${BACKUP_DIR}/${site}-${ts}.tar.gz" -C "$RESTORE_DIR/backup-${site}" 2>/dev/null || tar xzf "${BACKUP_DIR}/${site}-${ts}.tar.gz" -C "$RESTORE_DIR/backup-${site}" --strip-components=1 2>/dev/null || tar tzf "${BACKUP_DIR}/${site}-${ts}.tar.gz" --show-transformed-names 2>/dev/null | head -1
    tar xzf "${BACKUP_DIR}/${site}-${ts}.tar.gz" -C "$RESTORE_DIR/backup-${site}"
  else
    log "从云端下载: ${ts}"
    scp "txjp:${src}" "$RESTORE_DIR/"
    local f="$RESTORE_DIR/${site}-${ts}.tar.gz"
    mkdir -p "$RESTORE_DIR/backup-${site}"
    tar xzf "$f" -C "$RESTORE_DIR/backup-${site}"
  fi
}

# === 站点恢复函数 ===
restore_shared() {
  local ts="$TIMESTAMP"
  local src="${RESTORE_DIR}/backup-shared"
  [[ -f "$src/nginx"/*.conf ]] && { mkdir -p "$REMOTE_BASE/nginx-proxy-ai/conf.d"; scp "$src/nginx/"* txjp:"$REMOTE_BASE/nginx-proxy-ai/conf.d/"; }
  ssh txjp "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_DIR}/backup-assets/backup-assets"
  [[ ! -d "$src" ]] && src="${RESTORE_DIR}/backup-assets"
  for f in docker-compose.yml .env Dockerfile; do
    [[ -f "$src/$f" ]] && scp "$src/$f" txjp:$REMOTE_BASE/assets-ai/
  done
  [[ -d "$src/data/uploads" ]] && ssh txjp "rm -rf /root/docker/assets-ai/data/uploads && cp -a $src/data/uploads /root/docker/assets-ai/data/"
  [[ -f "$src/data/assets.db" ]] && ssh txjp "docker cp $src/data/assets.db assets-ai:/app/data/assets.db"
  ssh txjp "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_DIR}/backup-issue/backup-issue"
  [[ ! -d "$src" ]] && src="${RESTORE_DIR}/backup-issue"
  for f in docker-compose.yml .env Dockerfile; do
    [[ -f "$src/$f" ]] && scp "$src/$f" txjp:$REMOTE_BASE/issue-ai/
  done
  [[ -d "$src/data/uploads" ]] && ssh txjp "rm -rf /root/docker/issue-ai/data/uploads && cp -a $src/data/uploads /root/docker/issue-ai/data/"
  [[ -d "$src/data/reports" ]] && ssh txjp "rm -rf /root/docker/issue-ai/data/reports && cp -a $src/data/reports /root/docker/issue-ai/data/"
  [[ -f "$src/data/issue.db" ]] && ssh txjp "docker cp $src/data/issue.db issue-ai:/app/data/issue.db"
  ssh txjp "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_DIR}/backup-ldap/backup-ldap"
  [[ ! -d "$src" ]] && src="${RESTORE_DIR}/backup-ldap"
  for f in docker-compose.yml .env Dockerfile; do
    [[ -f "$src/$f" ]] && scp "$src/$f" txjp:$REMOTE_BASE/ldap-ai/
  done
  [[ -f "$src/data/users.db" ]] && ssh txjp "docker cp $src/data/users.db lldap:/data/users.db"
  ssh txjp "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_DIR}/backup-gitea/backup-gitea"
  [[ ! -d "$src" ]] && src="${RESTORE_DIR}/backup-gitea"
  [[ -f "$src/docker-compose.yml" ]] && scp "$src/docker-compose.yml" txjp:$REMOTE_BASE/gitea-ai/
  [[ -f "$src/data/gitea.db" ]] && ssh txjp "docker cp $src/data/gitea.db gitea-ai:/data/gitea.db"
  ssh txjp "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_DIR}/backup-oa/backup-oa"
  [[ ! -d "$src" ]] && src="${RESTORE_DIR}/backup-oa"
  for f in docker-compose.yml .env; do
    [[ -f "$src/$f" ]] && scp "$src/$f" txjp:$REMOTE_BASE/oa-ai/
  done
  ssh txjp "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_DIR}/backup-monitor/backup-monitor"
  [[ ! -d "$src" ]] && src="${RESTORE_DIR}/backup-monitor"
  for f in docker-compose.yml .env; do
    [[ -f "$src/$f" ]] && scp "$src/$f" txjp:$REMOTE_BASE/monitor-ai/
  done
  [[ -f "$src/data/monitor.db" ]] && ssh txjp "docker cp $src/data/monitor.db monitor-ai:/app/data/monitor.db"
  ssh txjp "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_DIR}/backup-${site}/backup-${site}"
  [[ ! -d "$src" ]] && src="${RESTORE_DIR}/backup-${site}"
  for f in docker-compose.yml .env Dockerfile; do
    [[ -f "$src/$f" ]] && scp "$src/$f" txjp:$REMOTE_BASE/${site}/
  done
  ssh txjp "cd $REMOTE_BASE/${site} && 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}..."
  scp "$src" "txjp:/tmp/db-restore.db"
  ssh txjp "docker cp /tmp/db-restore.db ${container}:${dbpath} && rm /tmp/db-restore.db"
  ssh txjp "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"

  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

  case "$choice" in
    1) restore_assets_from_full "$content" ;;
    2) restore_issue_from_full "$content" ;;
    3) restore_ldap_from_full "$content" ;;
    4) restore_gitea_from_full "$content" ;;
    5) restore_monitor_from_full "$content" ;;
    6) restore_nginx_from_full "$content" ;;
    7) restore_all_from_full "$content" ;;
    *) error "无效选择" ;;
  esac
}

restore_assets_from_full() {
  log "恢复 assets..."
  [[ -f "$1/databases/assets.db" ]] && ssh txjp "docker cp $1/databases/assets.db assets-ai:/app/data/assets.db"
  ssh txjp "cd $REMOTE_BASE/assets-ai && docker compose down && docker compose up -d"; log "完成"
}
restore_issue_from_full() {
  log "恢复 issue..."
  [[ -f "$1/databases/issue.db" ]] && ssh txjp "docker cp $1/databases/issue.db issue-ai:/app/data/issue.db"
  ssh txjp "cd $REMOTE_BASE/issue-ai && docker compose down && docker compose up -d"; log "完成"
}
restore_ldap_from_full() {
  log "恢复 ldap..."
  [[ -f "$1/databases/ldap.db" ]] && ssh txjp "docker cp $1/databases/ldap.db lldap:/data/users.db"
  ssh txjp "cd $REMOTE_BASE/ldap-ai && docker compose down && docker compose up -d"; log "完成"
}
restore_gitea_from_full() {
  log "恢复 gitea..."
  [[ -f "$1/databases/gitea.db" ]] && ssh txjp "docker cp $1/databases/gitea.db gitea-ai:/data/gitea.db"
  ssh txjp "cd $REMOTE_BASE/gitea-ai && docker compose down && docker compose up -d"; log "完成"
}
restore_monitor_from_full() {
  log "恢复 monitor..."
  [[ -f "$1/databases/monitor.db" ]] && ssh txjp "docker cp $1/databases/monitor.db monitor-ai:/app/data/monitor.db"
  ssh txjp "cd $REMOTE_BASE/monitor-ai && docker compose down && docker compose up -d --build"; log "完成"
}
restore_nginx_from_full() {
  log "恢复 nginx..."
  [[ -d "$1/nginx" ]] && ssh txjp "mkdir -p $REMOTE_BASE/nginx-proxy-ai/conf.d" && scp "$1/nginx/"* txjp:$REMOTE_BASE/nginx-proxy-ai/conf.d/
  ssh txjp "docker exec nginx-ai nginx -t && docker exec nginx-ai nginx -s reload"; log "完成"
}
restore_all_from_full() {
  local c="$1"
  restore_assets_from_full "$c"
  restore_issue_from_full "$c"
  restore_ldap_from_full "$c"
  restore_gitea_from_full "$c"
  restore_monitor_from_full "$c"
  restore_nginx_from_full "$c"
}

# === 主流程 ===
case "${1:-}" in
  --db)
    restore_db "$2" "$3" ;;
  --full)
    restore_full "$2" ;;
  --help|-h)
    echo "用法: bash restore.sh [--db|--full] [参数...]"
    echo "  无参数              交互式站点恢复"
    echo "  restore.sh TIMESTAMP 指定时间戳恢复"
    echo "  --db <site> <file>   数据库恢复"
    echo "  --full <file>        灾难恢复"
    ;;
  *)
    # 交互式站点恢复
    TIMESTAMP="${1:-}"
    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_status; index=1
    for site in $ALL_SITES; do
      local 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 ;;
        *)       warn "${site}: 无专用恢复逻辑"; restore_generic "$site" ;;
      esac
    done

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