#!/usr/bin/env bash
set -euo pipefail

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
MODEL_DIR="${XGB_MODEL_DIR:-$ROOT_DIR/ai-engine/models/xgboost}"

mkdir -p "$MODEL_DIR"

download_model() {
  local file_name="$1"
  local url="${2:-}"
  local expected_sha="${3:-}"

  if [[ -z "$url" ]]; then
    echo "⚠️  Skip ${file_name}: URL not provided"
    return 0
  fi

  local target_path="${MODEL_DIR}/${file_name}"
  local tmp_path="${target_path}.tmp"

  echo "⬇️  Downloading ${file_name}..."
  curl -fL --retry 3 --retry-delay 2 "$url" -o "$tmp_path"

  if [[ -n "$expected_sha" ]]; then
    local actual_sha
    actual_sha="$(sha256sum "$tmp_path" | awk '{print $1}')"
    if [[ "$actual_sha" != "$expected_sha" ]]; then
      echo "❌ SHA256 mismatch for ${file_name}"
      echo "   expected: ${expected_sha}"
      echo "   actual  : ${actual_sha}"
      rm -f "$tmp_path"
      exit 1
    fi
  fi

  mv "$tmp_path" "$target_path"
  echo "✅ Ready: ${file_name}"
}

download_model "xgb_ht_ft.pkl" "${MODEL_XGB_HT_FT_URL:-}" "${MODEL_XGB_HT_FT_SHA256:-}"
download_model "xgb_ms.pkl" "${MODEL_XGB_MS_URL:-}" "${MODEL_XGB_MS_SHA256:-}"
download_model "xgb_ou25.pkl" "${MODEL_XGB_OU25_URL:-}" "${MODEL_XGB_OU25_SHA256:-}"
download_model "xgb_btts.pkl" "${MODEL_XGB_BTTS_URL:-}" "${MODEL_XGB_BTTS_SHA256:-}"
download_model "xgb_ou15.pkl" "${MODEL_XGB_OU15_URL:-}" "${MODEL_XGB_OU15_SHA256:-}"
download_model "xgb_ou35.pkl" "${MODEL_XGB_OU35_URL:-}" "${MODEL_XGB_OU35_SHA256:-}"

echo "📦 XGBoost model bootstrap completed."