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

# Auto-install helper for the project.
# Usage: scripts/auto_install.sh [--venv PATH] [--cpu|--gpu] [--local-tts] [--kokoro] [--no-editable]
# Examples:
#  ./scripts/auto_install.sh --cpu --local-tts
#  ./scripts/auto_install.sh --venv .venv311 --gpu --kokoro

VENV=".venv311"
MODE="cpu"
USE_LOCAL_TTS="false"
USE_KOKORO="false"
EDITABLE="true"
TORCH_VERSION=""

print_help(){
  cat <<'EOF'
Usage: auto_install.sh [options]

Options:
  --venv PATH         Path to virtualenv (default: .venv311)
  --cpu               Install CPU-only PyTorch (default)
  --gpu               Install PyTorch (GPU). You may need to pick the right CUDA wheel manually.
  --torch-version V   Optional torch version to install (e.g. 2.2.2)
  --local-tts         Install Coqui TTS and extras for local TTS
  --kokoro            Assume you'll use Kokoro endpoint (no local TTS extras)
  --no-editable       Do not install in editable mode; install packages normally
  -h, --help          Show this help

Examples:
  ./scripts/auto_install.sh --cpu --local-tts
  ./scripts/auto_install.sh --venv .venv311 --gpu --kokoro
EOF
}

while [[ ${#} -gt 0 ]]; do
  case "$1" in
    --venv) VENV="$2"; shift 2;;
    --cpu) MODE="cpu"; shift;;
    --gpu) MODE="gpu"; shift;;
    --torch-version) TORCH_VERSION="$2"; shift 2;;
    --local-tts) USE_LOCAL_TTS="false"; shift;;
    --kokoro) USE_KOKORO="true"; shift;;
    --no-editable) EDITABLE="false"; shift;;
    -h|--help) print_help; exit 0;;
    *) echo "Unknown arg: $1"; print_help; exit 2;;
  esac
done

echo "Installing into venv: ${VENV}"
if [[ ! -d "${VENV}" ]]; then
  echo "Creating virtualenv ${VENV} (Python 3.11 recommended)..."
  python3.11 -m venv "${VENV}" || python3 -m venv "${VENV}"
fi

echo "Activating virtualenv..."
source "${VENV}/bin/activate"

echo "Upgrading pip, setuptools and wheel..."
python -m pip install --upgrade pip setuptools wheel

install_torch_cpu(){
  if [[ -n "${TORCH_VERSION}" ]]; then
    echo "Installing CPU torch ${TORCH_VERSION} from PyTorch CPU index..."
    python -m pip install "torch==${TORCH_VERSION}" --index-url https://download.pytorch.org/whl/cpu
  else
    echo "Installing latest CPU torch from PyTorch CPU index..."
    python -m pip install torch --index-url https://download.pytorch.org/whl/cpu
  fi
}

install_torch_gpu(){
  if [[ -n "${TORCH_VERSION}" ]]; then
    echo "Installing torch ${TORCH_VERSION} (GPU) - pip will try to pick a matching wheel"
    python -m pip install "torch==${TORCH_VERSION}"
  else
    echo "Installing torch (GPU) - pip will try to pick a matching wheel. If you need a specific CUDA wheel, install it manually following https://pytorch.org/"
    python -m pip install torch
  fi
}

if [[ "${MODE}" == "cpu" ]]; then
  install_torch_cpu
else
  install_torch_gpu
fi

EXTRAS=()
if [[ "${USE_LOCAL_TTS}" == "true" ]]; then
  EXTRAS+=(tts)
fi
if [[ "${MODE}" == "cpu" ]]; then
  EXTRAS+=(cpu)
fi

if [[ "${EDITABLE}" == "true" ]]; then
  if [[ ${#EXTRAS[@]} -gt 0 ]]; then
    IFS=, extras_str="${EXTRAS[*]}"
    echo "Installing package editable with extras: ${extras_str// /,}"
    python -m pip install -e .[${extras_str// /,}]
  else
    echo "Installing package editable without extras"
    python -m pip install -e .
  fi
else
  echo "Installing dependencies from requirements.txt"
  python -m pip install -r requirements.txt
fi

echo "Optional: if you plan to use Kokoro remote endpoint, copy and edit config.toml.example -> config.toml and fill kokoro values."
if [[ -f config.toml.example && ! -f config.toml ]]; then
  echo "Copying config.toml.example -> config.toml (edit before running)"
  cp config.toml.example config.toml
fi

echo "Setup complete. Next steps:\n  source ${VENV}/bin/activate\n  ./${VENV}/bin/python -m whisper_project.main --help"

echo "Done."