#!/usr/bin/env bash
#
# Incremental model downloader.
# Downloads missing models into an existing models directory without deleting
# existing files and without requiring uv.
#

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"
if [[ -f "${PROJECT_ROOT}/app/utils/download_models.py" ]]; then
  DEFAULT_MODEL_DIR="${PROJECT_ROOT}/models"
else
  DEFAULT_MODEL_DIR="/opt/dep/asr/models"
fi
MODEL_DIR="${MODEL_STORAGE_DIR:-${MODELS_DIR:-$DEFAULT_MODEL_DIR}}"
PYTHON_BIN="${PYTHON_BIN:-${PYTHON:-python3}}"
AUTO_MODE="false"
MODE="auto"
DOCKER_IMAGE="${ASR_IMAGE:-}"
ARG_COUNT="$#"
RUNTIME_MODELS_DIR=""
INTERACTIVE_DOCKER_TARGET="false"

info() { echo "[INFO] $1"; }
die() { echo "[ERROR] $1" >&2; exit 1; }

usage() {
  cat <<EOF
Usage:
  ./scripts/download-models.sh [options]

Options:
  --models-dir DIR     Model directory to fill; default: MODEL_STORAGE_DIR, MODELS_DIR, or project ./models
  --python-bin BIN     Python executable; default: PYTHON_BIN, PYTHON, or python3
  --mode MODE          auto, local, or docker; default: interactive local when no args
  --docker-image IMAGE Image used by docker mode; default: ASR_IMAGE from env or .env
  --runtime-models-dir DIR
                       Runtime path written into CAM++ config; default: models-dir
  --auto-mode          Reduce Python downloader output
  -h, --help           Show help

Examples:
  ./scripts/download-models.sh
  ./scripts/download-models.sh --models-dir /opt/dep/asr/models
  MODEL_STORAGE_DIR=/opt/dep/asr/models ./scripts/download-models.sh
  ./download-models.sh --mode docker --models-dir /opt/dep/asr/models
EOF
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --models-dir) MODEL_DIR="$2"; shift 2 ;;
    --python-bin) PYTHON_BIN="$2"; shift 2 ;;
    --mode) MODE="$2"; shift 2 ;;
    --docker-image) DOCKER_IMAGE="$2"; shift 2 ;;
    --runtime-models-dir) RUNTIME_MODELS_DIR="$2"; shift 2 ;;
    --auto-mode) AUTO_MODE="true"; shift ;;
    -h|--help) usage; exit 0 ;;
    *) die "Unknown argument: $1" ;;
  esac
done

case "$MODE" in
  auto|local|docker) ;;
  *) die "Invalid --mode: ${MODE}; expected auto, local, or docker" ;;
esac

if [[ -z "$DOCKER_IMAGE" && -f "${SCRIPT_DIR}/.env" ]]; then
  DOCKER_IMAGE="$(
    awk -F= '
      /^[[:space:]]*ASR_IMAGE[[:space:]]*=/ {
        value=$2
        sub(/^[[:space:]]*/, "", value)
        sub(/[[:space:]]*$/, "", value)
        gsub(/^"|"$/, "", value)
        gsub(/^'\''|'\''$/, "", value)
        print value
      }
    ' "${SCRIPT_DIR}/.env" | tail -n 1
  )"
fi

prompt_for_target_if_no_args() {
  [[ "$ARG_COUNT" -eq 0 ]] || return 0

  if [[ ! -t 0 || ! -t 1 ]]; then
    MODE="local"
    info "No arguments in non-interactive shell; defaulting to local mode"
    return 0
  fi

  local docker_model_dir="${MODEL_STORAGE_DIR:-/opt/dep/asr/models}"
  local choice=""
  local input=""

  cat <<EOF
Select model download target:
  1) Local run (default)
     download dir: ${MODEL_DIR}
     runtime dir:  ${MODEL_DIR}

  2) Docker/offline run
     host dir:     ${docker_model_dir}
     runtime dir:  /app/models (fixed default)
     note: CAM++ nested model paths will be written for /app/models.
EOF

  read -r -p "Select [1/2] (default 1): " choice
  case "${choice:-1}" in
    1)
      MODE="local"
      ;;
    2)
      MODE="local"
      INTERACTIVE_DOCKER_TARGET="true"
      read -r -p "Host models dir [${docker_model_dir}]: " input
      MODEL_DIR="${input:-$docker_model_dir}"
      RUNTIME_MODELS_DIR="/app/models"
      ;;
    *)
      die "Invalid selection: ${choice}"
      ;;
  esac
}

prompt_for_target_if_no_args

mkdir -p "$MODEL_DIR"
MODEL_DIR="$(cd "$MODEL_DIR" && pwd)"
MODEL_PARENT_DIR="$(dirname "$MODEL_DIR")"

info "Download settings:"
info "  Models: ${MODEL_DIR}"
if [[ -n "$RUNTIME_MODELS_DIR" ]]; then
  info "  Runtime models: ${RUNTIME_MODELS_DIR}"
fi
info "  Mode: incremental, existing files are kept"

run_local_download() {
  local downloader="${SCRIPT_DIR}/download_models_standalone.py"
  if [[ ! -f "$downloader" ]]; then
    return 2
  fi
  command -v "$PYTHON_BIN" >/dev/null 2>&1 || die "Python not found: ${PYTHON_BIN}"
  info "  Runner: local python (${PYTHON_BIN})"

  local args=("$downloader" --models-dir "$MODEL_DIR" --cache-dir "${MODELSCOPE_CACHE:-$MODEL_PARENT_DIR}")
  if [[ -n "$RUNTIME_MODELS_DIR" ]]; then
    args+=(--runtime-models-dir "$RUNTIME_MODELS_DIR")
  fi
  if [[ "$AUTO_MODE" == "true" ]]; then
    args+=(--auto-mode)
  fi
  "$PYTHON_BIN" "${args[@]}"
}

run_docker_download() {
  command -v docker >/dev/null 2>&1 || die "Docker not found; cannot use docker mode"
  [[ -n "$DOCKER_IMAGE" ]] || die "Docker image not set; pass --docker-image IMAGE or set ASR_IMAGE in .env"
  [[ -f "${SCRIPT_DIR}/download_models_standalone.py" ]] || die "Cannot find download_models_standalone.py"
  info "  Runner: docker image (${DOCKER_IMAGE})"

  local args=(/tmp/download_models_standalone.py --models-dir /app/models --cache-dir /app)
  if [[ -n "$RUNTIME_MODELS_DIR" ]]; then
    args+=(--runtime-models-dir "$RUNTIME_MODELS_DIR")
  fi
  if [[ "$AUTO_MODE" == "true" ]]; then
    args+=(--auto-mode)
  fi

  docker run --rm \
    -e MODELS_DIR=/app/models \
    -e MODELSCOPE_PATH=/app/models \
    -e MODELSCOPE_CACHE=/app \
    -e DATA_DIR=/app/data \
    -e TEMP_DIR=/app/data/temp \
    -e LOG_FILE=/app/data/logs/qwen3-asr.log \
    -v "${MODEL_DIR}:/app/models" \
    -v "${SCRIPT_DIR}/download_models_standalone.py:/tmp/download_models_standalone.py:ro" \
    "${DOCKER_IMAGE}" \
    python3 "${args[@]}"
}

case "$MODE" in
  local)
    run_local_download
    ;;
  docker)
    run_docker_download
    ;;
  auto)
    if ! run_local_download; then
      run_docker_download
    fi
    ;;
esac

info "Model download complete"
info "Models directory: ${MODEL_DIR}"