#!/usr/bin/env bash
# ACE-Step: настройка окружения и запуск веб-интерфейса на Linux / macOS.
# Windows-аналог — start.bat.
set -uo pipefail

cd "$(dirname "$0")"

# ============================ НАСТРОЙКИ ============================
VENV_DIR="${VENV_DIR:-venv}"          # каталог виртуального окружения
PORT="${PORT:-7865}"                  # порт веб-интерфейса
SERVER_NAME="${SERVER_NAME:-127.0.0.1}"  # 0.0.0.0 — доступ из локальной сети
DEVICE_ID="${DEVICE_ID:-0}"           # номер видеокарты
CHECKPOINT_PATH="${CHECKPOINT_PATH:-}"   # пусто — скачать в ~/.cache/ace-step
# Сборка PyTorch под CUDA (cu126 / cu124 / cu121). Только для Linux.
TORCH_INDEX="${TORCH_INDEX:-https://download.pytorch.org/whl/cu126}"
# ==================================================================

SCRIPT_NAME="$(basename "$0")"
OPT_LOWVRAM=0
OPT_CPU=0
OPT_SHARE=0
DO_SETUP_ONLY=0
DO_REINSTALL=0
DO_UPDATE=0

usage() {
    cat <<EOF

Использование: ./$SCRIPT_NAME [флаги]

  --lowvram      Режим экономии видеопамяти (~8 ГБ VRAM)
  --cpu          Запуск на процессоре, без CUDA (очень медленно)
  --share        Публичная ссылка Gradio
  --port <N>     Порт веб-интерфейса (по умолчанию $PORT)
  --device <N>   Номер видеокарты (по умолчанию $DEVICE_ID)
  --listen       Слушать 0.0.0.0 (доступ из локальной сети)
  --reinstall    Пересоздать виртуальное окружение с нуля
  --update       Обновить зависимости
  --setup        Только установка, без запуска
  --help         Эта справка

Настройки по умолчанию — в блоке НАСТРОЙКИ в начале файла,
их можно переопределить переменными окружения (PORT=7870 ./$SCRIPT_NAME).

EOF
}

die() {
    echo
    echo "[ОШИБКА] $*" >&2
    exit 1
}

while [ $# -gt 0 ]; do
    case "$1" in
        --help|-h)   usage; exit 0 ;;
        --lowvram)   OPT_LOWVRAM=1; shift ;;
        --cpu)       OPT_CPU=1; shift ;;
        --share)     OPT_SHARE=1; shift ;;
        --setup)     DO_SETUP_ONLY=1; shift ;;
        --reinstall) DO_REINSTALL=1; shift ;;
        --update)    DO_UPDATE=1; shift ;;
        --listen)    SERVER_NAME="0.0.0.0"; shift ;;
        --port)      PORT="${2:?--port требует значение}"; shift 2 ;;
        --device)    DEVICE_ID="${2:?--device требует значение}"; shift 2 ;;
        *)
            echo "[ОШИБКА] Неизвестный аргумент: $1" >&2
            echo "Запустите ./$SCRIPT_NAME --help для справки." >&2
            exit 1 ;;
    esac
done

echo
echo "=========================================="
echo "  ACE-Step - генерация музыки"
echo "=========================================="
echo

case "$(uname -s)" in
    Darwin) IS_MAC=1 ;;
    *)      IS_MAC=0 ;;
esac

VENV_PY="$PWD/$VENV_DIR/bin/python"

# Рабочим считаем окружение, в котором есть и python, и pip: каталог от
# прерванной установки (или venv без ensurepip) выглядит как готовый, но валится
# дальше с невнятной ошибкой.
venv_ok() {
    [ -x "$VENV_PY" ] && "$VENV_PY" -m pip --version >/dev/null 2>&1
}

venv_failed() {
    echo >&2
    echo "[ОШИБКА] Не удалось создать рабочее виртуальное окружение в \"$VENV_DIR\"." >&2
    if [ "$IS_MAC" = 0 ] && command -v apt >/dev/null 2>&1; then
        echo "На Debian/Ubuntu для этого нужен отдельный пакет:" >&2
        echo "    sudo apt install python3-venv" >&2
    fi
    exit 1
}

if [ "$DO_REINSTALL" = 1 ] && [ -d "$VENV_DIR" ]; then
    echo "[1/4] Удаляю старое окружение \"$VENV_DIR\"..."
    rm -rf "$VENV_DIR"
fi

# ------------------------- поиск Python -------------------------
if venv_ok; then
    echo "[1/4] Виртуальное окружение найдено: $VENV_DIR"
else
    if [ -d "$VENV_DIR" ]; then
        echo "[1/4] Окружение \"$VENV_DIR\" неработоспособно, пересоздаю..."
        rm -rf "$VENV_DIR"
    fi
    echo "[1/4] Ищу подходящий Python..."
    SYS_PY=""
    for candidate in python3.12 python3.11 python3.10 python3.13 python3; do
        if command -v "$candidate" >/dev/null 2>&1 &&
           "$candidate" -c 'import sys; sys.exit(0 if sys.version_info >= (3,10) else 1)' 2>/dev/null; then
            SYS_PY="$candidate"
            break
        fi
    done
    [ -n "$SYS_PY" ] || die "Не найден Python 3.10 или новее. Установите его и повторите."

    echo "      Найден Python $("$SYS_PY" -c 'import sys; print(sys.version.split()[0])') ($SYS_PY)"
    echo "[1/4] Создаю виртуальное окружение в \"$VENV_DIR\"..."
    "$SYS_PY" -m venv "$VENV_DIR" || venv_failed
    venv_ok || venv_failed
fi

# ------------------------- зависимости -------------------------
echo "[2/4] Проверяю зависимости..."

NEED_TORCH=0
"$VENV_PY" -c "import importlib.util,sys; sys.exit(0 if importlib.util.find_spec('torch') else 1)" 2>/dev/null || NEED_TORCH=1

# Проверяем именно установку пакета: каталог acestep лежит рядом со скриптом,
# поэтому find_spec('acestep') сработал бы и без установленных зависимостей.
NEED_ACESTEP=0
"$VENV_PY" -c "import importlib.metadata as md, importlib.util as u, sys; md.version('ace_step'); sys.exit(0 if u.find_spec('click') and u.find_spec('gradio') else 1)" 2>/dev/null || NEED_ACESTEP=1

if [ "$DO_UPDATE" = 1 ]; then
    NEED_TORCH=1
    NEED_ACESTEP=1
fi

if [ "$NEED_TORCH" = 1 ] || [ "$NEED_ACESTEP" = 1 ]; then
    echo "      Обновляю pip..."
    "$VENV_PY" -m pip install --upgrade pip setuptools wheel --quiet ||
        die "Не удалось обновить pip. Проверьте подключение к интернету."
fi

if [ "$NEED_TORCH" = 1 ]; then
    if [ "$IS_MAC" = 1 ]; then
        # На macOS колёса с PyPI уже собраны с поддержкой MPS.
        echo "      Устанавливаю PyTorch (macOS/MPS). Это займёт несколько минут..."
        "$VENV_PY" -m pip install --upgrade torch torchvision torchaudio
    elif [ "$OPT_CPU" = 1 ]; then
        echo "      Устанавливаю PyTorch для CPU. Это займёт несколько минут..."
        "$VENV_PY" -m pip install --upgrade torch torchvision torchaudio \
            --index-url https://download.pytorch.org/whl/cpu
    else
        echo "      Устанавливаю PyTorch с поддержкой CUDA. Это займёт несколько минут..."
        "$VENV_PY" -m pip install --upgrade torch torchvision torchaudio --index-url "$TORCH_INDEX"
    fi || die "Не удалось установить PyTorch. Если у вас другая версия CUDA, поменяйте TORCH_INDEX."
fi

if [ "$NEED_ACESTEP" = 1 ]; then
    echo "      Устанавливаю ACE-Step и зависимости..."
    "$VENV_PY" -m pip install -e . || die "Не удалось установить ACE-Step."
fi

if [ "$OPT_LOWVRAM" = 1 ] && [ "$IS_MAC" = 0 ]; then
    if ! "$VENV_PY" -c "import importlib.util,sys; sys.exit(0 if importlib.util.find_spec('triton') else 1)" 2>/dev/null; then
        echo "      Режим экономии VRAM: устанавливаю triton..."
        "$VENV_PY" -m pip install triton --quiet ||
            echo "[ВНИМАНИЕ] triton не установился, --torch_compile может не заработать."
    fi
fi

# ------------------------- проверка устройства -------------------------
echo "[3/4] Проверяю устройство..."
BF16=true
if [ "$OPT_CPU" = 1 ]; then
    echo "      Принудительный режим CPU. Генерация будет очень медленной."
    DEVICE_ID=-1
    BF16=false
elif [ "$IS_MAC" = 1 ]; then
    # На MPS bfloat16 приводит к ошибкам, см. README.
    echo "      macOS: используется MPS, bf16 отключён."
    BF16=false
else
    GPUINFO="$("$VENV_PY" -c "import torch; print((torch.cuda.get_device_name(0) + ' - ' + str(round(torch.cuda.get_device_properties(0).total_memory/1073741824, 1)) + ' GB VRAM') if torch.cuda.is_available() else 'NO_CUDA')" 2>/dev/null || echo NO_CUDA)"
    if [ "$GPUINFO" = "NO_CUDA" ]; then
        echo
        echo "[ВНИМАНИЕ] CUDA недоступна - модель будет работать на процессоре (очень медленно)."
        echo "Если у вас есть видеокарта NVIDIA, обновите драйвер и переустановите PyTorch:"
        echo "    ./$SCRIPT_NAME --update"
        echo
    else
        echo "      $GPUINFO"
    fi
fi

if [ "$DO_SETUP_ONLY" = 1 ]; then
    echo
    echo "Установка завершена. Для запуска выполните ./$SCRIPT_NAME"
    exit 0
fi

# ------------------------- запуск -------------------------
ARGS=(--port "$PORT" --server_name "$SERVER_NAME" "--device_id=$DEVICE_ID" --bf16 "$BF16")
[ -n "$CHECKPOINT_PATH" ] && ARGS+=(--checkpoint_path "$CHECKPOINT_PATH")
[ "$OPT_SHARE" = 1 ] && ARGS+=(--share true)
if [ "$OPT_LOWVRAM" = 1 ]; then
    ARGS+=(--cpu_offload true --overlapped_decode true)
    [ "$IS_MAC" = 0 ] && ARGS+=(--torch_compile true)
fi

echo "[4/4] Запускаю веб-интерфейс..."
echo
[ "$OPT_LOWVRAM" = 1 ] && echo "      Режим: экономия видеопамяти"
echo "      Идёт загрузка моделей, это занимает время."
echo "      При первом запуске дополнительно скачиваются веса (~8 ГБ)."
echo
echo "      Интерфейс будет доступен на http://$SERVER_NAME:$PORT"
echo "      когда ниже появится строка \"Running on local URL\"."
echo "      Остановить: Ctrl+C в этом окне."
echo

exec "$VENV_PY" -m acestep.gui "${ARGS[@]}"