#!/usr/bin/env bash
set -eu

TARGET_WEIGHTS_DIR="${TARGET_WEIGHTS_DIR:-/data/algorithmService/weights}"
FILES_VERSION="${FILES_VERSION:-1.0.0}"
FILES_REMOTE="${FILES_REMOTE:-screen@rsync.inner.xcastle.net::screen/apps/algorithm-python/${FILES_VERSION}/weights}"
RSYNC_PASSWORD="${RSYNC_PASSWORD:-txry2025}"
SOURCE_WEIGHTS_DIR="${SOURCE_WEIGHTS_DIR:-}"
TARGET_OWNER="${TARGET_OWNER:-txry:txry}"
DOWNLOAD_PARTS="${DOWNLOAD_PARTS:-8}"
DOWNLOAD_RETRIES="${DOWNLOAD_RETRIES:-3}"
DOWNLOAD_CONNECT_TIMEOUT="${DOWNLOAD_CONNECT_TIMEOUT:-10}"

readonly -a REQUIRED_ONNX_MODELS=(
  "2dfan4.onnx"
  "w600k_r50.onnx"
  "det_10g.onnx"
  "gender_age_res34.onnx"
  "inswapper_128.onnx"
  "gfpgan_1.4.onnx"
  "inswapper_emap.txt"
  "yoloface_8n.onnx"
)

readonly -a REQUIRED_EXTRA_FILES=(
  "inswapper_128.onnx__emap.npy"
)

declare -a TEMP_DOWNLOAD_PATHS=()

cleanupTempDownloads() {
  local path
  for path in "${TEMP_DOWNLOAD_PATHS[@]}"; do
    if [ -n "$path" ]; then
      rm -rf "$path"
    fi
  done
}

trap cleanupTempDownloads EXIT

trackTempPath() {
  TEMP_DOWNLOAD_PATHS+=("$1")
}

ensureCommand() {
  local cmd="$1"
  if ! command -v "$cmd" >/dev/null 2>&1; then
    echo "未找到命令: $cmd"
    exit 1
  fi
}

normalizeDownloadParts() {
  if [[ "$DOWNLOAD_PARTS" =~ ^[0-9]+$ ]] && [ "$DOWNLOAD_PARTS" -ge 1 ]; then
    echo "$DOWNLOAD_PARTS"
    return
  fi

  echo "DOWNLOAD_PARTS 无效，使用默认 8: $DOWNLOAD_PARTS" >&2
  echo 8
}

isHttpRemote() {
  case "$FILES_REMOTE" in
    http://*|https://*) return 0 ;;
    *) return 1 ;;
  esac
}

remoteUrlForFile() {
  local file="$1"
  echo "${FILES_REMOTE%/}/${file}"
}

getHttpContentLength() {
  local url="$1"

  curl -fsSLI --connect-timeout "$DOWNLOAD_CONNECT_TIMEOUT" "$url" \
    | tr -d '\r' \
    | awk 'tolower($1) == "content-length:" { size = $2 } END { if (size ~ /^[0-9]+$/) print size }'
}

supportsHttpRange() {
  local url="$1"
  local tmp_dir="$2"
  local probe_file="${tmp_dir}/range-probe"
  local http_code

  http_code="$(curl -fsSL \
    --connect-timeout "$DOWNLOAD_CONNECT_TIMEOUT" \
    --max-filesize 2 \
    --range 0-0 \
    --output "$probe_file" \
    --write-out "%{http_code}" \
    "$url" 2>/dev/null || true)"

  [ "$http_code" = "206" ] && [ "$(wc -c < "$probe_file" | tr -d ' ')" = "1" ]
}

downloadHttpSingle() {
  local url="$1"
  local target_file="$2"
  local tmp_file

  tmp_file="$(mktemp "${target_file}.download.XXXXXX")" || return 1
  trackTempPath "$tmp_file"

  curl -fL \
    --retry "$DOWNLOAD_RETRIES" \
    --connect-timeout "$DOWNLOAD_CONNECT_TIMEOUT" \
    --output "$tmp_file" \
    "$url" || return 1
  mv -f "$tmp_file" "$target_file" || return 1
}

downloadHttpSegmented() {
  local url="$1"
  local target_file="$2"
  local part_count="$3"
  local remote_size tmp_dir tmp_file

  remote_size="$(getHttpContentLength "$url" || true)"
  if [ -z "$remote_size" ] || [ "$remote_size" -le 0 ]; then
    echo "无法获取文件大小，改用普通下载: $url"
    downloadHttpSingle "$url" "$target_file" || return 1
    return
  fi

  if [ "$part_count" -le 1 ]; then
    downloadHttpSingle "$url" "$target_file" || return 1
    return
  fi

  if [ "$remote_size" -lt "$part_count" ]; then
    part_count="$remote_size"
  fi

  tmp_dir="$(mktemp -d "${TMPDIR:-/tmp}/onnx-download.XXXXXX")" || return 1
  tmp_file="$(mktemp "${target_file}.download.XXXXXX")" || return 1
  trackTempPath "$tmp_dir"
  trackTempPath "$tmp_file"

  if ! supportsHttpRange "$url" "$tmp_dir"; then
    echo "远端不支持 Range 分片，改用普通下载: $url"
    rm -rf "$tmp_dir"
    downloadHttpSingle "$url" "$target_file" || return 1
    return
  fi

  echo "分片下载: $url -> $target_file，分片数: $part_count，大小: $remote_size bytes"

  local i start end part_file
  local -a pids=()
  for ((i = 0; i < part_count; i++)); do
    start=$((i * remote_size / part_count))
    end=$(((i + 1) * remote_size / part_count - 1))
    if [ "$i" -eq "$((part_count - 1))" ]; then
      end=$((remote_size - 1))
    fi

    part_file="${tmp_dir}/part-${i}"
    curl -fsSL \
      --retry "$DOWNLOAD_RETRIES" \
      --connect-timeout "$DOWNLOAD_CONNECT_TIMEOUT" \
      --range "${start}-${end}" \
      --output "$part_file" \
      "$url" &
    pids+=("$!")
  done

  local failed=0
  local pid
  for pid in "${pids[@]}"; do
    if ! wait "$pid"; then
      failed=1
    fi
  done

  if [ "$failed" -ne 0 ]; then
    echo "分片下载失败: $url" >&2
    return 1
  fi

  : > "$tmp_file" || return 1
  for ((i = 0; i < part_count; i++)); do
    cat "${tmp_dir}/part-${i}" >> "$tmp_file" || return 1
  done

  local downloaded_size
  downloaded_size="$(wc -c < "$tmp_file" | tr -d ' ')"
  if [ "$downloaded_size" != "$remote_size" ]; then
    echo "分片合并后大小不一致: $downloaded_size != $remote_size，$url" >&2
    return 1
  fi

  mv -f "$tmp_file" "$target_file" || return 1
  rm -rf "$tmp_dir"
}

fixOwner() {
  if [ -z "$TARGET_OWNER" ]; then
    return
  fi

  local owner_user="${TARGET_OWNER%%:*}"
  if [ "$(id -u)" -ne 0 ]; then
    echo "当前不是 root 用户，跳过 chown: $TARGET_OWNER"
    return
  fi

  if ! id "$owner_user" >/dev/null 2>&1; then
    echo "用户不存在，跳过 chown: $owner_user"
    return
  fi

  if ! chown -R "$TARGET_OWNER" "$TARGET_WEIGHTS_DIR"; then
    echo "chown 失败，跳过权限修正: $TARGET_OWNER"
  fi
}

copyFromLocal() {
  local file="$1"
  local target_file="$2"
  local source_file="${SOURCE_WEIGHTS_DIR%/}/${file}"

  if [ ! -f "$source_file" ]; then
    echo "本地源文件不存在: $source_file"
    exit 1
  fi

  cp -p "$source_file" "$target_file" || return 1
}

copyFromRemote() {
  local file="$1"
  local target_file="$2"

  if isHttpRemote; then
    ensureCommand curl
    downloadHttpSegmented "$(remoteUrlForFile "$file")" "$target_file" "$(normalizeDownloadParts)" || return 1
    return
  fi

  # rsync does not expose HTTP byte ranges, so keep the original rsync download path.
  ensureCommand rsync
  export RSYNC_PASSWORD
  rsync -v -P --timeout=60 -a --chown=txry:txry "${FILES_REMOTE%/}/${file}" "$target_file" || return 1
}

copyMissingFile() {
  local file="$1"
  local target_file="${TARGET_WEIGHTS_DIR%/}/${file}"

  if [ -e "$target_file" ] || [ -L "$target_file" ]; then
    echo "已存在，跳过: $target_file"
    return 1
  fi

  mkdir -p "$(dirname "$target_file")" || return 2
  if [ -n "$SOURCE_WEIGHTS_DIR" ]; then
    if ! copyFromLocal "$file" "$target_file"; then
      echo "本地拷贝失败: $file" >&2
      return 2
    fi
  else
    if ! copyFromRemote "$file" "$target_file"; then
      echo "远端下载失败: $file" >&2
      return 2
    fi
  fi
  echo "已拷贝: $target_file"
  return 0
}

copyMissingModels() {
  local total copied skipped file status
  total=0
  copied=0
  skipped=0

  mkdir -p "$TARGET_WEIGHTS_DIR"

  for file in "${REQUIRED_ONNX_MODELS[@]}" "${REQUIRED_EXTRA_FILES[@]}"; do
    total=$((total + 1))
    set +e
    copyMissingFile "$file"
    status=$?
    set -e

    case "$status" in
      0) copied=$((copied + 1)) ;;
      1) skipped=$((skipped + 1)) ;;
      *) exit "$status" ;;
    esac
  done

  if [ "$copied" -gt 0 ]; then
    fixOwner
  fi

  echo "模型检查完成，总数: $total，已拷贝: $copied，已存在: $skipped"
}

copyMissingModels