#!/usr/bin/env bash
set -e

###############################################################################
# 0. Подготовка среды: обновление apt, установка unzip
###############################################################################
echo ">>> [Script 1] Step 0: Установка системных зависимостей ..."
apt-get update && apt-get install -y unzip


echo ">>> [Script 1] Step 1: Downloading MyRunpodLoraScripts.zip ..."
wget -O /workspace/MyRunpodLoraScripts.zip \
  "https://huggingface.co/Gerchegg/FeaturesFluxAI/resolve/main/MyRunpodLoraScripts.zip"

echo ">>> [Script 1] Step 2: Unzipping MyRunpodLoraScripts.zip ..."
unzip -o /workspace/MyRunpodLoraScripts.zip -d /workspace
###############################################################################
# 1. Копирование файлов из локальной папки в рабочую директорию
###############################################################################
echo ">>> [Script 1] Step 1: Копирование файлов из MyRunpodLoraScripts в /workspace ..."
mkdir -p /workspace/MyRunpodLoraScripts
cp -r /workspace/MyRunpodLoraScripts/* /workspace/

###############################################################################
# 2. Запускаем два "параллельных" потока (фона):
#   - Thread 1: Устанавливает пакеты + python Download_Train_Models.py
#   - Thread 2: Устанавливает нужные deps (kohya_ss, torch, xformers), запускает GUI
###############################################################################
echo ">>> [Script 1] Запуск параллельных потоков ..."

# ------------------ Поток 1 --------------------
(
  echo ">>> [Thread 1] Установка Python пакетов для huggingface ..."
  pip install huggingface_hub ipywidgets hf_transfer

  export HF_HUB_ENABLE_HF_TRANSFER=1

  echo ">>> [Thread 1] Запуск Download_Train_Models.py ..."
  dwn_path=$(find /workspace -name "Download_Train_Models.py" | head -n 1)
  if [ -z "$dwn_path" ]; then
    echo "[Thread 1] ОШИБКА: Download_Train_Models.py не найден!"
    exit 1
  fi

  python "$dwn_path" --dir /workspace

  echo ">>> [Thread 1] Готово."
) &

# ------------------ Поток 2 --------------------
(
  echo ">>> [Thread 2] Установка Python и OS зависимостей ..."
  apt update --yes
  yes | apt-get install python3.10-tk
  apt-get install psmisc --yes

  echo ">>> [Thread 2] Клонирование kohya_ss ..."
  cd /workspace
  if [ ! -d "/workspace/kohya_ss" ]; then
    git clone https://github.com/bmaltais/kohya_ss.git
  fi

  cd /workspace/kohya_ss

  echo ">>> [Thread 2] Создание Python venv ..."
  python3 -m venv venv
  source venv/bin/activate
  yes | apt-get install python3.10-tk

  echo ">>> [Thread 2] Запуск setup.sh ..."
  ./setup.sh -n -u

  apt update --yes
  yes | apt-get install python3.10-tk
  apt-get install psmisc --yes

  pip install hf_transfer
  export HF_HUB_ENABLE_HF_TRANSFER=0

  echo ">>> [Thread 2] Завершение процессов на порту 7860 ..."
  fuser -k 7860/tcp || true

  source venv/bin/activate

  echo ">>> [Thread 2] Обновление torch & xformers ..."
  pip uninstall -y xformers
  pip install torch==2.5.1+cu124 torchvision --index-url https://download.pytorch.org/whl/cu124
  pip install xformers==0.0.28.post3 --index-url https://download.pytorch.org/whl/cu124

  echo ">>> [Thread 2] Запуск kohya_ss GUI на 0.0.0.0:7860 ..."
  ./gui.sh --listen=0.0.0.0 --share --noverify
) &

###############################################################################
# 3. Создаём общую структуру папок (Models, Images, Logs, Configs)
###############################################################################
echo ">>> [Script 1] Создание структуры папок в /workspace/MyLoRADataset ..."
mkdir -p /workspace/MyLoRADataset/Models
mkdir -p /workspace/MyLoRADataset/Images
mkdir -p /workspace/MyLoRADataset/Logs
mkdir -p /workspace/MyLoRADataset/Configs

###############################################################################
# 4. Копирование конфигураций LoRA в папку Configs
###############################################################################
echo ">>> [Script 1] Копирование конфигураций LoRA в /workspace/MyLoRADataset/Configs ..."
cp -r /workspace/LoRA_Tab_LoRA_Training_Best_FLUX_Configs/* /workspace/MyLoRADataset/Configs/

###############################################################################
# 5. Ожидаем завершения потоков
###############################################################################
echo ">>> [Script 1] Ожидание завершения потоков ..."
wait
echo ">>> [Script 1] ГОТОВО! Теперь вы можете использовать Kohya SS GUI для обучения LoRA моделей."
echo ">>> [Script 1] Доступ к GUI: http://127.0.0.1:7860 или через внешнюю ссылку, предоставленную в выводе."
echo ">>> [Script 1] Конфигурации LoRA доступны в /workspace/MyLoRADataset/Configs"