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

release_root="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
python_bin="${PYTHON_BIN:-python3.12}"
venv_dir="${release_root}/.venv"
build_dir="${release_root}/.build"
scales_revision="a04c1b63b1a7a670840fb3e97a82c0dbe2a35ded"
scales_dir="${build_dir}/fish-scales-ops-${scales_revision}"
cuda_root="${CUDA_HOME:-/usr/local/cuda}"

for command_name in "${python_bin}" git; do
  if ! command -v "${command_name}" >/dev/null 2>&1; then
    echo "Missing required command: ${command_name}" >&2
    exit 1
  fi
done
if [[ ! -x "${cuda_root}/bin/nvcc" ]]; then
  echo "CUDA compiler not found at ${cuda_root}/bin/nvcc" >&2
  exit 1
fi

mkdir -p "${build_dir}" "${release_root}/runtime-data/torch-extensions"
if [[ ! -x "${venv_dir}/bin/python" ]]; then
  "${python_bin}" -m venv "${venv_dir}"
fi
"${venv_dir}/bin/python" -m pip install --upgrade "pip>=25" "uv==0.12.3"
uv_bin="${venv_dir}/bin/uv"

"${uv_bin}" pip install --python "${venv_dir}/bin/python" \
  "setuptools>=75" wheel ninja packaging \
  "torch==2.11.0" "torchaudio==2.11.0" \
  --torch-backend=cu130
"${uv_bin}" pip install --python "${venv_dir}/bin/python" \
  --overrides "${release_root}/runtime/torch-overrides.txt" \
  --torch-backend=cu130 \
  "${release_root}/vendor/fish-speech" \
  -r "${release_root}/requirements.txt"

if [[ ! -d "${scales_dir}/.git" ]]; then
  if [[ -e "${scales_dir}" ]]; then
    echo "${scales_dir} exists but is not the expected Git checkout; move it aside and retry." >&2
    exit 1
  fi
  git init "${scales_dir}"
  git -C "${scales_dir}" remote add origin https://github.com/fishaudio/fish-scales-ops.git
  git -C "${scales_dir}" fetch --depth 1 origin "${scales_revision}"
  git -C "${scales_dir}" checkout --detach FETCH_HEAD
  git -C "${scales_dir}" submodule update --init --recursive --depth 1
fi
if [[ "$(git -C "${scales_dir}" rev-parse HEAD)" != "${scales_revision}" ]]; then
  echo "Unexpected fish-scales-ops revision in ${scales_dir}" >&2
  exit 1
fi
git -C "${scales_dir}" submodule update --init --recursive --depth 1

CUTLASS_DIR="${scales_dir}/3rdparty/cutlass" \
CUDA_HOME="${cuda_root}" \
MAX_JOBS="${MAX_JOBS:-2}" \
TORCH_CUDA_ARCH_LIST=12.0a \
"${uv_bin}" pip install --python "${venv_dir}/bin/python" \
  --no-build-isolation "${scales_dir}/python"

"${venv_dir}/bin/python" -c \
  "import comfy_kitchen, fish_scales_ops, fish_speech, torch; print('installed', torch.__version__, torch.version.cuda)"
echo "Installation complete. Start the service with ./launch.sh"