#!/usr/bin/env bash
# Reinstall the AMD ROCm runtime stack (torch + CTranslate2 ROCm wheels) into the repo .venv.
# Never creates or re-syncs the venv; run `uv sync` first, then this script.

set -euo pipefail

# Resolve the repo root from the script location
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(dirname "$SCRIPT_DIR")"

CT2_VERSION="4.8.4"
ROCM_INDEX="https://download.pytorch.org/whl/rocm7.2"

help() {
    cat << 'EOF'
ROCm runtime installer (torch + CTranslate2 ROCm wheels into .venv)

Usage: scripts/install-rocm.sh [--ct2-version X.Y.Z] [--rocm-index <torch index url>]

Options:
  --ct2-version X.Y.Z       CTranslate2 release version (default: 4.8.4)
  --rocm-index <url>        Torch ROCm index URL (default: https://download.pytorch.org/whl/rocm7.2)
  -h, --help                Show this message

Example:
  scripts/install-rocm.sh
  scripts/install-rocm.sh --ct2-version 4.8.4 --rocm-index https://download.pytorch.org/whl/rocm7.2

Note: any later `uv sync` reverts these wheels, so re-run this script afterwards.
EOF
}

main() {
    while [[ $# -gt 0 ]]; do
        case "${1:-}" in
            --ct2-version)
                CT2_VERSION="${2:-}"
                if [[ -z "$CT2_VERSION" ]]; then
                    echo "Missing value for --ct2-version" >&2
                    return 1
                fi
                shift 2
                ;;
            --rocm-index)
                ROCM_INDEX="${2:-}"
                if [[ -z "$ROCM_INDEX" ]]; then
                    echo "Missing value for --rocm-index" >&2
                    return 1
                fi
                shift 2
                ;;
            -h|--help|help)
                help
                return 0
                ;;
            *)
                echo "Unknown argument: ${1:-}" >&2
                echo "Use 'scripts/install-rocm.sh --help' for usage" >&2
                return 1
                ;;
        esac
    done

    local venv_python="$REPO_ROOT/.venv/bin/python"
    if [[ ! -x "$venv_python" ]]; then
        echo "Error: $venv_python not found." >&2
        echo "Run 'uv sync' first ('uv sync --extra speaker --extra triage' if you use the optional lanes), then re-run this script." >&2
        return 1
    fi

    export TMPDIR="${TMPDIR:-$HOME/.tmp}"
    mkdir -p "$TMPDIR"

    local venv_pip="$REPO_ROOT/.venv/bin/pip"
    "$venv_pip" install --force-reinstall torch torchvision torchaudio \
        --index-url "$ROCM_INDEX" \
        --extra-index-url "https://pypi.org/simple"

    local workdir
    workdir="$(mktemp -d "$TMPDIR/ct2-rocm-XXXXXX")"
    trap 'rm -rf "${workdir:-}"' EXIT
    mkdir -p "$workdir"

    local bundle_url="https://github.com/OpenNMT/CTranslate2/releases/download/v${CT2_VERSION}/rocm-python-wheels-Linux.zip"
    local bundle_zip="$workdir/rocm-python-wheels-Linux.zip"
    curl -fsSL -o "$bundle_zip" "$bundle_url"
    unzip -q -o "$bundle_zip" -d "$workdir"

    local py_tag
    py_tag="$("$venv_python" -c 'import sys; print(f"cp{sys.version_info.major}{sys.version_info.minor}")')"

    # The release archive layout is not stable across versions: v4.8.0 puts the
    # wheels at the top level, v4.8.4 nests them under temp-linux/. Search deeply.
    local wheel
    wheel="$(find "$workdir" -type f -name "ctranslate2-*-${py_tag}-*.whl" -print -quit)"
    if [[ -z "$wheel" ]]; then
        echo "Error: no CTranslate2 wheel for $py_tag found in $bundle_url" >&2
        return 1
    fi
    "$venv_pip" install --force-reinstall "$wheel"

    rm -rf "$workdir"
    trap - EXIT

    "$venv_python" -c 'import ctranslate2, torch; print(torch.__version__); print(torch.version.hip); print(ctranslate2.get_cuda_device_count()); import sys; sys.exit(1 if torch.version.hip is None or ctranslate2.get_cuda_device_count() == 0 else 0)'

    cat << 'EOF'
ROCm stack installed. Next step: run `uv run stt-daemon --check` to confirm
(the transcriber/CLI now report the resolved device).
Note: any later `uv sync` will revert these wheels, so re-run this script afterwards.
EOF
}

main "$@"