#!/usr/bin/env bash
# Download OCR/object models used by the service.

set -euo pipefail

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
MODELS_DIR="${MODELS_DIR:-$ROOT_DIR/models}"
PADDLE_DIR="${PADDLE_DIR:-$MODELS_DIR/paddle}"
TESSDATA_LOCAL_DIR="${TESSDATA_LOCAL_DIR:-$MODELS_DIR/tessdata_best}"

DOWNLOAD_CORAL=1
DOWNLOAD_EAST=1
DOWNLOAD_PADDLE=1
DOWNLOAD_TESSDATA=1

usage() {
  cat <<'EOF'
Usage:
  ./scripts/download_models.sh [options]

Options:
  --only-coral       Download only Coral model + labels
  --only-east        Download only EAST text detector model
  --only-paddle      Download only PaddleOCR models
  --only-tessdata    Download only local tessdata_best (eng+rus)
  --no-coral         Skip Coral model download
  --no-east          Skip EAST model download
  --no-paddle        Skip PaddleOCR model download
  --no-tessdata      Skip local tessdata_best download
  -h, --help         Show this help

Environment overrides:
  MODELS_DIR=/abs/path/to/models
  PADDLE_DIR=/abs/path/to/models/paddle
  TESSDATA_LOCAL_DIR=/abs/path/to/tessdata_best
EOF
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --only-coral)
      DOWNLOAD_CORAL=1; DOWNLOAD_EAST=0; DOWNLOAD_PADDLE=0; DOWNLOAD_TESSDATA=0
      shift
      ;;
    --only-east)
      DOWNLOAD_CORAL=0; DOWNLOAD_EAST=1; DOWNLOAD_PADDLE=0; DOWNLOAD_TESSDATA=0
      shift
      ;;
    --only-paddle)
      DOWNLOAD_CORAL=0; DOWNLOAD_EAST=0; DOWNLOAD_PADDLE=1; DOWNLOAD_TESSDATA=0
      shift
      ;;
    --only-tessdata)
      DOWNLOAD_CORAL=0; DOWNLOAD_EAST=0; DOWNLOAD_PADDLE=0; DOWNLOAD_TESSDATA=1
      shift
      ;;
    --no-coral)
      DOWNLOAD_CORAL=0
      shift
      ;;
    --no-east)
      DOWNLOAD_EAST=0
      shift
      ;;
    --no-paddle)
      DOWNLOAD_PADDLE=0
      shift
      ;;
    --no-tessdata)
      DOWNLOAD_TESSDATA=0
      shift
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      echo "[ERROR] Unknown option: $1" >&2
      usage
      exit 1
      ;;
  esac
done

mkdir -p "$MODELS_DIR" "$PADDLE_DIR" "$TESSDATA_LOCAL_DIR"

if command -v curl >/dev/null 2>&1; then
  DOWNLOADER="curl"
elif command -v wget >/dev/null 2>&1; then
  DOWNLOADER="wget"
else
  echo "[ERROR] curl or wget is required" >&2
  exit 1
fi

download_file() {
  local url="$1"
  local dst="$2"
  mkdir -p "$(dirname "$dst")"
  if [[ -s "$dst" ]]; then
    echo "[SKIP] $dst already exists"
    return 0
  fi
  local tmp="${dst}.part"
  echo "[GET] $url"
  if [[ "$DOWNLOADER" == "curl" ]]; then
    curl -L --fail --retry 3 --connect-timeout 20 --max-time 600 "$url" -o "$tmp"
  else
    wget -O "$tmp" "$url"
  fi
  mv "$tmp" "$dst"
}

download_and_extract_tar() {
  local url="$1"
  local tag="$2"
  local out_dir="$3"
  local cache_dir="$MODELS_DIR/.downloads"
  local archive="$cache_dir/${tag}.tar"
  mkdir -p "$cache_dir" "$out_dir"
  download_file "$url" "$archive"
  echo "[EXTRACT] $archive -> $out_dir"
  tar -xf "$archive" -C "$out_dir"
}

# ---- URLs ----
CORAL_MODEL_URL="https://github.com/google-coral/test_data/raw/master/ssd_mobilenet_v2_coco_quant_postprocess_edgetpu.tflite"
CORAL_LABELS_URL="https://github.com/google-coral/test_data/raw/master/coco_labels.txt"
EAST_URL="https://github.com/argman/EAST/releases/download/model/frozen_east_text_detection.pb"
TESS_ENG_URL="https://github.com/tesseract-ocr/tessdata_best/raw/main/eng.traineddata"
TESS_RUS_URL="https://github.com/tesseract-ocr/tessdata_best/raw/main/rus.traineddata"

# PaddleOCR v5 mobile models (detection + RU/EN recognition)
PADDLE_DET_URL="https://paddle-model-ecology.bj.bcebos.com/paddlex/official_inference_model/paddle3.0.0/PP-OCRv5_mobile_det_infer.tar"
PADDLE_REC_URL="https://paddle-model-ecology.bj.bcebos.com/paddlex/official_inference_model/paddle3.0.0/eslav_PP-OCRv5_mobile_rec_infer.tar"

if [[ "$DOWNLOAD_CORAL" == "1" ]]; then
  download_file "$CORAL_MODEL_URL" "$MODELS_DIR/ssd_mobilenet_v2_coco_quant_postprocess_edgetpu.tflite"
  download_file "$CORAL_LABELS_URL" "$MODELS_DIR/coco_labels.txt"
fi

if [[ "$DOWNLOAD_EAST" == "1" ]]; then
  download_file "$EAST_URL" "$MODELS_DIR/frozen_east_text_detection.pb"
fi

if [[ "$DOWNLOAD_PADDLE" == "1" ]]; then
  download_and_extract_tar "$PADDLE_DET_URL" "PP-OCRv5_mobile_det_infer" "$PADDLE_DIR"
  download_and_extract_tar "$PADDLE_REC_URL" "eslav_PP-OCRv5_mobile_rec_infer" "$PADDLE_DIR"
fi

if [[ "$DOWNLOAD_TESSDATA" == "1" ]]; then
  download_file "$TESS_ENG_URL" "$TESSDATA_LOCAL_DIR/eng.traineddata"
  download_file "$TESS_RUS_URL" "$TESSDATA_LOCAL_DIR/rus.traineddata"
fi

cat <<EOF

[DONE] Model download complete.

Recommended env:
  export MODEL_EDGETPU_PATH="$MODELS_DIR/ssd_mobilenet_v2_coco_quant_postprocess_edgetpu.tflite"
  export MODEL_LABELS_PATH="$MODELS_DIR/coco_labels.txt"
  export MODEL_EAST_PATH="$MODELS_DIR/frozen_east_text_detection.pb"
  export TESSDATA_PREFIX="$TESSDATA_LOCAL_DIR"
  export OCR_BACKEND="tesseract"

For PaddleOCR backend:
  export OCR_BACKEND="paddle"
  export PADDLE_LANG="ru"
  export PADDLE_DET_MODEL_DIR="$PADDLE_DIR/PP-OCRv5_mobile_det_infer"
  export PADDLE_REC_MODEL_DIR="$PADDLE_DIR/eslav_PP-OCRv5_mobile_rec_infer"
EOF