#!/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                      # 忽略完成标记，强制重新同步
#
# 传输策略：双方均可用 rsync → 增量同步(断点续传、幂等)；否则退化为 scp -r 全量
# 拷贝(Windows 节点一般无 rsync，自动走 scp)。SSH 复用 deploy.sh 的 ControlMaster
# 连接复用，多次执行不重复输密码。
# ==============================================================================

set -eo 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
FORCE=false
INTERACTIVE=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)          FORCE=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
}

# 传输一个目录树（rsync 优先，退化 scp）
sync_dir() { # $1=user $2=ip $3=port $4=远端绝对目录 $5=本地目标
  local u="$1" ip="$2" port="$3" rdir="$4" dest="$5"
  local ssh_opts="-o ControlMaster=auto -o ControlPath=${SSH_CONTROL_DIR}/cm-${u}@${ip}:${port} -o ControlPersist=1800"

  mkdir -p "$(dirname "${dest}")"
  # 完成标记：已成功同步过且未 --force 时跳过
  if [ -f "${dest}/.fetch-complete" ] && [ "${FORCE}" != "true" ]; then
    echo -e "  ${CYAN}[i]${NC} 已有完成标记 ${dest}/.fetch-complete，跳过（--force 可重拉）"
    return 0
  fi

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

  # rsync 可用性（本机 + 远端）
  local rs_ok=false
  if command -v rsync >/dev/null 2>&1; then
    if ssh -p "${port}" ${ssh_opts} "${u}@${ip}" "command -v rsync" >/dev/null 2>&1; then
      rs_ok=true
    fi
  fi

  if [ "${rs_ok}" = "true" ]; then
    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
    echo -e "  ${YELLOW}[→]${NC} 远端无 rsync，退化 scp 全量拷贝: ${u}@${ip}:${rdir}/ → ${dest}/"
    scp -P "${port}" ${ssh_opts} -r "${u}@${ip}:${rdir}/" "${dest}/"
  fi
  touch "${dest}/.fetch-complete"
}

# 备份单个节点
backup_node() { # $1 = profile 文件
  local p="$1" node_id role env_mode u ip port dir
  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 "")

  # 跳过 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
      [ -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
    # 本机直接拷贝，不走网络
    echo -e "  ${BLUE}[→]${NC} 本机复制: ./data/result → ${dest_rc}"
    mkdir -p "${dest_rc}"
    cp -a ./data/result/. "${dest_rc}/"
    touch "${dest_rc}/.fetch-complete"
  else
    sync_dir "${u}" "${ip}" "${port}" "${dir}/data/result" "${dest_rc}" || return 1
  fi
  echo -e "  ${GREEN}[√]${NC} result 已同步，共 $(ls "${dest_rc}" | grep -cv '^\.fetch-complete$' || 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}"
      cp -a ./data/work/. "${dest_wk}/" 2>/dev/null || true
    else
      sync_dir "${u}" "${ip}" "${port}" "${dir}/data/work" "${dest_wk}" || return 1
    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

  # 用户显式勾选本机节点 → 自动纳入备份树（否则 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}"