#!/usr/bin/env bash
# ==============================================================================
# fetch_results.sh — 从全部计算节点同步 data/result 完整科学产物到本地备份
# ==============================================================================
# 背景：node 端 result 目录永久保留全部网格点的完整产物（.spec/.cont/.iden/
# 各阶段快照/conv.json/日志等，撤销了早期的 MAX_RESULT_MODELS=200 LRU 上限）。
# 为做异地备份、以及便于失败归因与收敛验证，把各节点 result 目录整树拉回本地，
# 按节点分目录存放。
#
# 用法:
#   ./scripts/fetch_results.sh                              # 无参数 → 交互式向导选择节点
#   ./scripts/fetch_results.sh -i                           # 显式进入交互式向导
#   ./scripts/fetch_results.sh -n node-fmq-dckj-02          # 非交互：仅同步指定节点(DCTS_NODE_ID)
#   ./scripts/fetch_results.sh -d /backup/dcts-salvage      # 覆盖备份根目录
#   ./scripts/fetch_results.sh --with-work                  # 连带拉取 data/work 沙盒残留
#   ./scripts/fetch_results.sh --include-local              # 把本机节点也复制进备份树
#   ./scripts/fetch_results.sh --force                      # 强制重拉：忽略本地已存在的同名目录，全部重新 scp 覆盖
#
# 传输策略（两类节点都走增量，重复运行只补差异、不重拷全量）：
#   - 双方均可用 rsync → rsync 增量（断点续传、幂等，自动补新增/变更文件）
#   - 否则（典型 Windows：无 rsync）→ scp 目录级增量：先列举远端顶层网格点子目录，
#     只 scp 本地缺失的目录，已存在的不重拷（避免 win-01 的 52GB 全量重拷）
# 远端 rsync/列目录探测均用多 shell 兼容命令（POSIX/command-v/where/Get-Command、
# ls/dir/Get-ChildItem），适配 Linux + Windows(cmd/PowerShell 默认 shell)。
# SSH 复用 deploy.sh 的 ControlMaster，多次执行不重复输密码。
# ==============================================================================

set -euo pipefail

GREEN='\033[0;32m'
BLUE='\033[0;34m'
RED='\033[0;31m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
NC='\033[0m'

# 定位至项目根目录（与 deploy.sh 一致）
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
WORK_DIR="$(cd "${SCRIPT_DIR}/.." && pwd)"
cd "${WORK_DIR}"

# =============================================================================
# 默认配置（可被命令行参数覆盖）
# =============================================================================
SSH_CONTROL_DIR="${SSH_CONTROL_DIR:-$HOME/.ssh/cm}"
BACKUP_ROOT="${BACKUP_ROOT:-${WORK_DIR}/data/salvage}"
ONLY_NODE=""
WITH_WORK=false
INCLUDE_LOCAL=false
INTERACTIVE=false
FORCE_REPULL=false
HAD_ARGS=false
SELECTED_PROFILES=()

# =============================================================================
# 解析命令行参数
# =============================================================================
while [[ $# -gt 0 ]]; do
  HAD_ARGS=true
  case "$1" in
    -n|--node)        ONLY_NODE="$2"; shift 2 ;;
    -d|--dest)        BACKUP_ROOT="$2"; shift 2 ;;
    -i|--interactive) INTERACTIVE=true; shift ;;
    --with-work)      WITH_WORK=true; shift ;;
    --include-local)  INCLUDE_LOCAL=true; shift ;;
    --force|-f)       FORCE_REPULL=true; shift ;;
    -h|--help)
      sed -n '2,32p' "$0"
      exit 0 ;;
    *) echo -e "${RED}未知参数: $1${NC}"; exit 1 ;;
  esac
done

mkdir -p "${BACKUP_ROOT}"

# =============================================================================
# 工具函数
# =============================================================================
# 从 profile 文件读取键值（避免 source 造成变量污染；写法对齐 deploy.sh）
read_profile_var() { # $1=profile $2=键 $3=默认值
  local p="$1" k="$2" d="${3:-}"
  local v
  v=$(grep -E "^${k}=" "$p" 2>/dev/null | head -n1 | cut -d'=' -f2- | tr -d '"' | tr -d "'" || true)
  [ -n "$v" ] && printf '%s' "$v" || printf '%s' "$d"
}

# rsync 远端路径转义：含空格目录需单引号包住（现有路径均无空格，此处兜底）
rsync_remote_path() { # $1 = 远端绝对路径
  local p="$1"
  case "$p" in
    *\ *) printf "'%s'" "$p" ;;
    *)    printf '%s' "$p" ;;
  esac
}

# 远端目录列表（仅顶层条目名，用于 scp 目录级增量比对）。
# 按远端路径形态选择对应的列举命令：
#   - 盘符开头(如 E:/...) → Windows：powershell Get-ChildItem（兼容 cmd/PowerShell 默认 shell）
#   - 否则 → POSIX：ls -1
# 返回名字列表（每行一个），失败返回空。
remote_list_dir() { # $1=user $2=ip $3=port $4=ssh_opts $5=远端绝对目录
  local u="$1" ip="$2" port="$3" opts="$4" rdir="$5"
  # 统一 tr -d '\r'：Windows powershell/openssh 输出带 CRLF，会让 grep -Fxq 精确匹配
  # 永远失败（已同步目录每次重传）并产生带 CR 的畸形 scp 路径。
  if [[ "${rdir}" =~ ^[A-Za-z]:[\\/] ]]; then
    # Windows 路径：经 powershell 列举，-Name 直接返回每项一行
    ssh -p "${port}" ${opts} "${u}@${ip}" \
      "powershell -NoProfile -Command \"Get-ChildItem -Name -Path '${rdir}'\"" 2>/dev/null | tr -d '\r'
  else
    ssh -p "${port}" ${opts} "${u}@${ip}" "ls -1 '${rdir}'" 2>/dev/null | tr -d '\r'
  fi
}

# scp 单个远端目录到本地，采用"先落地 .partial 再原子 mv"模式：
#   1. 目标目录 ${dest}/${d} 已存在 → 视为已同步，跳过（保持原增量语义）；
#      但 force=true 时忽略该判断，强制重拉（用于同名目录内容已变更的场景）
#   2. 否则 scp 远端目录到 ${dest}/${d}.partial
#   3. scp 成功后 mv 为 ${dest}/${d}（同 dest 文件系统，原子语义）
#   4. scp 失败 → 清理残留 .partial，return 非零，由调用方决定是否继续
# 这样即便 scp 传输过程中断，也不会留下"伪完成"目录导致下次被跳过漏传。
# 强制模式下只在 scp 成功后才清空旧目录再 mv，避免传输失败时丢失旧备份。
scp_dir() { # $1=u $2=ip $3=port $4=ssh_opts $5=远端根目录 $6=本地目标 $7=目录名 $8=force
  local u="$1" ip="$2" port="$3" opts="$4" rdir="$5" dest="$6" d="$7" force="$8"
  local final="${dest}/${d}"
  if [ "${force}" != "true" ] && [ -e "${final}" ]; then
    return 0
  fi
  local tmp="${dest}/.${d}.partial"
  # 兜底清理可能的历史残留（上次中断留下的半成品）
  rm -rf "${tmp}"
  if scp -P "${port}" ${opts} -r "${u}@${ip}:${rdir}/${d}" "${tmp}"; then
    # 强制模式：先清旧目录再原子 mv（避免 mv 把新目录并入已存在的旧目录）
    [ "${force}" = "true" ] && rm -rf "${final}"
    mv "${tmp}" "${final}"
  else
    echo -e "  ${RED}[!]${NC} scp 失败: ${d}，清理残留 ${d}.partial"
    rm -rf "${tmp}"
    return 1
  fi
}

# 传输一个目录树（rsync 优先，退化 scp 目录级增量）
# $6=可选跳板 ([user@]host[:port]，来自 profile 的 REMOTE_PROXY_JUMP）：
#   目标节点不与本机直连、需经中转节点（如 dckj-01 → dckj-02/03）时启用 ProxyJump。
sync_dir() { # $1=user $2=ip $3=port $4=远端绝对目录 $5=本地目标 $6=jump
  local u="$1" ip="$2" port="$3" rdir="$4" dest="$5" jump="${6:-}"
  local ssh_opts="-o ControlMaster=auto -o ControlPath=${SSH_CONTROL_DIR}/cm-${u}@${ip}:${port} -o ControlPersist=1800"
  [ -n "${jump}" ] && ssh_opts="${ssh_opts} -o ProxyJump=${jump}"
  local jump_note=""
  [ -n "${jump}" ] && jump_note="（经跳板 ${jump}）"

  mkdir -p "${dest}"

  # 连通性与主连接建立（BatchMode=no 允许交互输密码）。
  # 先显式建立 ControlMaster 主连接（-fN 后台保活）：跳板机/目标机的密码集中在此
  # 交互输入并各只输一次；成功后后续的探测/rsync/scp 全部复用 socket，不再要密码。
  # 若主连接已存在（如刚同步过其他目录），-O check 直接复用。
  if ! ssh -p "${port}" ${ssh_opts} -O check "${u}@${ip}" >/dev/null 2>&1; then
    echo -e "  ${BLUE}[→]${NC} 建立 SSH 主连接${jump_note}（需分别输入跳板机/目标机密码）..."
    if ! ssh -p "${port}" ${ssh_opts} -fN "${u}@${ip}"; then
      echo -e "  ${RED}[!]${NC} SSH 主连接建立失败: ${u}@${ip}${jump_note}（密码错误或链路不通），跳过本节点"
      return 1
    fi
  fi

  # 探测命令用 `echo ok` 而非 `true`：Windows OpenSSH 默认 shell 是 cmd/PowerShell，
  # 没有 Unix 的 true，会导致"密码正确但仍判连接失败"。
  if ! ssh -p "${port}" ${ssh_opts} "${u}@${ip}" "echo ok" >/dev/null 2>&1; then
    echo -e "  ${RED}[!] 无法连接 ${u}@${ip}，跳过本节点${NC}"
    return 1
  fi

  # rsync 可用性（本机 + 远端）。远端探测按平台分支：
  #   Windows 路径 → powershell Get-Command；POSIX 路径 → command -v
  local rs_ok=false
  if command -v rsync >/dev/null 2>&1; then
    local probe
    if [[ "${rdir}" =~ ^[A-Za-z]:[\\/] ]]; then
      probe="powershell -NoProfile -Command \"Get-Command rsync -ErrorAction SilentlyContinue | Out-Null\""
    else
      probe="command -v rsync"
    fi
    if ssh -p "${port}" ${ssh_opts} "${u}@${ip}" "${probe}" >/dev/null 2>&1; then
      rs_ok=true
    fi
  fi

  if [ "${rs_ok}" = "true" ]; then
    # rsync 天然增量幂等：每次都跑，自动补齐新增/变更的文件。
    echo -e "  ${BLUE}[→]${NC} rsync 增量同步: ${u}@${ip}:${rdir}/ → ${dest}/"
    rsync -a --partial --info=progress2 \
      -e "ssh -p ${port} ${ssh_opts}" \
      "${u}@${ip}:$(rsync_remote_path "${rdir}")/" "${dest}/"
  else
    # 无 rsync（典型 Windows）：scp 目录级增量——只拉本地缺失的网格点子目录，
    # 不重复拷贝已存在的目录（避免 win-01 的 52GB 全量重拷）。
    local remote_dirs existing missing
    remote_dirs=$(remote_list_dir "${u}" "${ip}" "${port}" "${ssh_opts}" "${rdir}")
    if [ -z "${remote_dirs}" ]; then
      echo -e "  ${YELLOW}[!]${NC} 远端目录 ${rdir} 为空或无法列举，跳过"
      return 1
    fi
    existing=$(ls -1 "${dest}" 2>/dev/null)
    missing=""
    local d
    # 远端目录名白名单校验：scp 路径会把 ${d} 直接拼进 scp 的 remote 路径，
    # 这里只允许字母/数字/._-，禁止 / 与 . .. —— 防御性过滤畸形或注入目录名。
    while IFS= read -r d; do
      [ -n "$d" ] || continue
      case "$d" in
        "."|"..")
          echo -e "  ${YELLOW}[!]${NC} 跳过可疑远端目录名: ${d}"
          continue ;;
      esac
      if [[ "$d" == */* ]] || ! [[ "$d" =~ ^[A-Za-z0-9._-]+$ ]]; then
        echo -e "  ${YELLOW}[!]${NC} 跳过不符合白名单的远端目录名: ${d}"
        continue
      fi
      if [ "${FORCE_REPULL}" = "true" ] || ! grep -Fxq -- "$d" <<<"${existing}"; then
        missing="${missing}${d}"$'\n'
      fi
    done <<<"${remote_dirs}"

    local total new
    total=$(printf '%s\n' "${remote_dirs}" | grep -c .)
    new=$(printf '%s' "${missing}" | grep -c .)
    if [ "${FORCE_REPULL}" = "true" ]; then
      echo -e "  ${YELLOW}[→]${NC} scp 强制重拉: ${u}@${ip}:${rdir}/ 共 ${total} 个，全部重拉 ${new} 个 → ${dest}/"
    else
      echo -e "  ${YELLOW}[→]${NC} scp 目录级增量: ${u}@${ip}:${rdir}/ 共 ${total} 个，缺失 ${new} 个 → ${dest}/"
    fi
    if [ "${new}" -eq 0 ]; then
      if [ "${FORCE_REPULL}" = "true" ]; then
        echo -e "  ${GREEN}[√]${NC} 强制重拉完成，无远端目录"
      else
        echo -e "  ${GREEN}[√]${NC} 本地已是最新，无新增目录"
      fi
      return 0
    fi
    # scp 原子写入：先落地到 .partial 临时目录，成功后同文件系统原子 mv 为最终目录名。
    # 开拷前先清理可能的历史 .partial 残留（上次中断留下的半成品），scp_dir 自身失败
    # 时也会清理各自的 .partial；这样避免下次因"目录已存在"误判为已同步而漏传。
    # （原先用 trap RETURN 兜底，但 RETURN 陷阱会被外层 backup_node 继承，函数返回时
    #   dest 已出作用域，触发 set -u 的"未绑定变量"报错，故改为前置清理。）
    rm -rf "${dest}"/.*.partial 2>/dev/null || true
    while IFS= read -r d; do
      [ -n "$d" ] || continue
      if [ "${FORCE_REPULL}" = "true" ]; then
        echo -e "  ${YELLOW}[→]${NC} scp 强制重拉: ${d}"
      else
        echo -e "  ${YELLOW}[→]${NC} scp 新增: ${d}"
      fi
      scp_dir "${u}" "${ip}" "${port}" "${ssh_opts}" "${rdir}" "${dest}" "$d" "${FORCE_REPULL}"
    done <<<"${missing}"
  fi
}

# 备份单个节点
backup_node() { # $1 = profile 文件
  local p="$1" node_id role env_mode u ip port dir jump
  node_id=$(read_profile_var "$p" DCTS_NODE_ID)
  role=$(read_profile_var "$p" DEPLOY_ROLE node)
  env_mode=$(read_profile_var "$p" DEPLOY_ENV "")
  u=$(read_profile_var "$p" REMOTE_USER "fmq")
  ip=$(read_profile_var "$p" REMOTE_IP "")
  port=$(read_profile_var "$p" REMOTE_PORT 22)
  dir=$(read_profile_var "$p" REMOTE_DIR "")
  jump=$(read_profile_var "$p" REMOTE_PROXY_JUMP "")

  # 跳过 server 角色（server 的 result 不是计算产物，且 seeds 另行备份）
  if [ "${role}" = "server" ]; then
    echo -e "${YELLOW}== ${p##*/}: 角色 server，跳过（seed 备份不属本脚本职责）${NC}"
    return 0
  fi
  # 本地节点：result 就在本机 ./data/result，默认不重复备份
  if [ "${env_mode}" = "local" ]; then
    if [ "${INCLUDE_LOCAL}" = "true" ]; then
      # 本机节点 ID 优先取根目录 .env（与节点注册、看板展示的 DCTS_NODE_ID 一致）
      [ -n "${node_id}" ] || node_id=$(grep -E "^DCTS_NODE_ID=" .env 2>/dev/null | head -n1 | cut -d'=' -f2- | tr -d '"' | tr -d "'")
      [ -n "${node_id}" ] || node_id="node-local"
    else
      echo -e "${YELLOW}== ${p##*/}: 本地节点(${node_id:-localhost})，result 已在 ./data/result，跳过（--include-local 可纳入备份树）${NC}"
      return 0
    fi
  fi
  # 按节点过滤
  if [ -n "${ONLY_NODE}" ] && [ "${node_id}" != "${ONLY_NODE}" ]; then
    return 0
  fi
  if [ -z "${node_id}" ]; then
    node_id="${u}@${ip}"
  fi

  local loc=""
  if [ "${env_mode}" = "local" ]; then
    loc="（本机）"
  else
    loc=" ${u}@${ip}:${dir}"
  fi
  echo -e "\n${GREEN}== ${p##*/} → 节点 ${node_id}${NC} ${loc}"

  local dest_rc="${BACKUP_ROOT}/${node_id}/result"
  if [ "${env_mode}" = "local" ]; then
    # 本机节点：优先 rsync 增量（避免每次 cp -a 全量重拷），无 rsync 才退化 cp -a。
    echo -e "  ${BLUE}[→]${NC} 本机复制: ./data/result → ${dest_rc}"
    mkdir -p "${dest_rc}"
    if command -v rsync >/dev/null 2>&1; then
      rsync -a --info=progress2 ./data/result/ "${dest_rc}/"
    else
      cp -a ./data/result/. "${dest_rc}/"
    fi
  else
    sync_dir "${u}" "${ip}" "${port}" "${dir}/data/result" "${dest_rc}" "${jump}" || return 1
  fi
  echo -e "  ${GREEN}[√]${NC} result 已同步，共 $(ls -1 "${dest_rc}" 2>/dev/null | grep -c . || true) 个网格点子目录"

  # 可选：一并拉取 data/work 沙盒残留（未清理的计算现场，含完整过程文件）
  if [ "${WITH_WORK}" = "true" ]; then
    local dest_wk="${BACKUP_ROOT}/${node_id}/work"
    if [ "${env_mode}" = "local" ]; then
      echo -e "  ${BLUE}[→]${NC} 本机复制: ./data/work → ${dest_wk}"
      mkdir -p "${dest_wk}"
      if command -v rsync >/dev/null 2>&1; then
        rsync -a ./data/work/ "${dest_wk}/" 2>/dev/null || true
      else
        cp -a ./data/work/. "${dest_wk}/" 2>/dev/null || true
      fi
    else
      sync_dir "${u}" "${ip}" "${port}" "${dir}/data/work" "${dest_wk}" "${jump}" || true
    fi
  fi
}

# =============================================================================
# 交互式选择（无参数或 -i 时进入，风格对齐 deploy.sh 的 Profile 向导）
# =============================================================================
interactive_select() {
  # 构建可同步列表：非 *.example 且 DEPLOY_ROLE=node 的 profile（server 不在菜单）
  local menu=()
  for p in deploy.env.d/*.env; do
    [ -f "$p" ] || continue
    [[ "$p" != *.example ]] || continue
    if [ "$(read_profile_var "$p" DEPLOY_ROLE node)" = "node" ]; then
      menu+=("$p")
    fi
  done
  if [ ${#menu[@]} -eq 0 ]; then
    echo -e "${RED}错误: deploy.env.d/ 下没有可用的 node profile${NC}"
    exit 1
  fi

  echo -e "${BLUE}==============================================================${NC}"
  echo -e "${CYAN}             🚀 DCTS 计算结果产物下载控制台                       ${NC}"
  echo -e "${BLUE}==============================================================${NC}"
  echo -e "${YELLOW}【节点选择】检测到 ${#menu[@]} 个可用的计算节点 Profile 文件:${NC}"
  local i p node_id env_mode ip user
  i=1
  for p in "${menu[@]}"; do
    node_id=$(read_profile_var "$p" DCTS_NODE_ID)
    env_mode=$(read_profile_var "$p" DEPLOY_ENV "")
    ip=$(read_profile_var "$p" REMOTE_IP "")
    user=$(read_profile_var "$p" REMOTE_USER "fmq")
    [ -n "$node_id" ] || node_id="${user}@${ip:-localhost}"
    local tag="远程" ipshow="$ip"
    if [ "${env_mode}" = "local" ]; then tag="本机"; ipshow="(本机)"; fi
    printf "  %d) %s (%s)  %s@%s  [%s]\n" "$i" "$node_id" "${p##*/}" "$user" "$ipshow" "$tag"
    ((i++))
  done
  echo -e "  a) [全部节点] 依次同步上述全部 ${#menu[@]} 个节点"
  read -r -p "请选择节点 [序号如 1,2 或 1-2; 全部如 a; 默认 0=退出]: " CHOICE
  CHOICE=${CHOICE:-0}

  SELECTED_PROFILES=()
  if [[ "${CHOICE,,}" == "a" ]] || [[ "${CHOICE,,}" == "all" ]]; then
    SELECTED_PROFILES=("${menu[@]}")
  elif [[ "${CHOICE}" =~ ^[0-9] ]]; then
    local part from to n
    IFS=',' read -r -a ADDR <<< "${CHOICE}"
    for part in "${ADDR[@]}"; do
      if [[ "$part" =~ ^([0-9]+)-([0-9]+)$ ]]; then
        from="${BASH_REMATCH[1]}"; to="${BASH_REMATCH[2]}"
        for ((n=from; n<=to; n++)); do
          if [ "$n" -ge 1 ] && [ "$n" -le "${#menu[@]}" ]; then
            SELECTED_PROFILES+=("${menu[$((n-1))]}")
          fi
        done
      elif [[ "$part" =~ ^[0-9]+$ ]]; then
        if [ "$part" -ge 1 ] && [ "$part" -le "${#menu[@]}" ]; then
          SELECTED_PROFILES+=("${menu[$((part-1))]}")
        fi
      fi
    done
  fi

  if [ ${#SELECTED_PROFILES[@]} -eq 0 ]; then
    echo -e "${YELLOW}未选择任何节点，退出。${NC}"
    exit 0
  fi

  # 是否连带同步 data/work 沙盒残留
  if [ "${WITH_WORK}" != "true" ]; then
    read -r -p "是否连带同步 data/work 沙盒残留（未清理的计算现场，含完整过程文件）? [y/N]: " WK_CHOICE
    [[ "${WK_CHOICE,,}" == "y" || "${WK_CHOICE,,}" == "yes" ]] && WITH_WORK=true
  fi

  # Windows(scp)节点是否强制重拉：本地即使已有同名目录也全部重新下载覆盖
  if [ "${FORCE_REPULL}" != "true" ]; then
    read -r -p "是否强制重拉（忽略本地已有同名目录，全部重新 scp 覆盖，仅对 Windows/scp 节点生效）? [y/N]: " FORCE_CHOICE
    [[ "${FORCE_CHOICE,,}" == "y" || "${FORCE_CHOICE,,}" == "yes" ]] && FORCE_REPULL=true
  fi

  # 用户显式勾选本机节点 → 自动纳入备份树（否则 backup_node 会跳过本机）
  local p2 env2
  for p2 in "${SELECTED_PROFILES[@]}"; do
    env2=$(read_profile_var "$p2" DEPLOY_ENV "")
    [ "${env2}" = "local" ] && INCLUDE_LOCAL=true
  done

  echo -e "${CYAN}将同步 ${#SELECTED_PROFILES[@]} 个节点 → ${BACKUP_ROOT}${NC}"
}

if [ "${HAD_ARGS}" = "false" ] || [ "${INTERACTIVE}" = "true" ]; then
  interactive_select
fi

# =============================================================================
# 主流程：遍历 deploy.env.d/*.env（排除 *.example）
# =============================================================================
echo -e "${CYAN}备份根目录: ${BACKUP_ROOT}${NC}"
echo -e "${CYAN}SSH 复用目录: ${SSH_CONTROL_DIR}${NC}"
mkdir -p "${SSH_CONTROL_DIR}"

profiles=()
for p in deploy.env.d/*.env; do
  [ -f "$p" ] || continue
  [[ "$p" != *.example ]] || continue
  profiles+=("$p")
done

if [ ${#profiles[@]} -eq 0 ]; then
  echo -e "${RED}错误: deploy.env.d/ 下没有可用的 node profile${NC}"
  exit 1
fi

# 交互式模式下仅同步用户勾选的节点
if [ ${#SELECTED_PROFILES[@]} -gt 0 ]; then
  profiles=("${SELECTED_PROFILES[@]}")
fi

for p in "${profiles[@]}"; do
  backup_node "$p" || true
done

echo -e "\n${GREEN}===== 全部节点处理完成，备份位于 ${BACKUP_ROOT} =====${NC}"