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

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

###############################################################################
# 1. Скачиваем zip-файл с Hugging Face и распаковываем
###############################################################################
echo ">>> [Script 1] Step 1: Downloading MyRunpodFinetuneScripts.zip ..."
wget -O /workspace/MyRunpodFinetuneScripts.zip \
  "https://huggingface.co/Gerchegg/FeaturesFluxAI/resolve/main/MyRunpodFinetuneScriptsv17.zip"

echo ">>> [Script 1] Step 2: Unzipping MyRunpodFinetuneScripts.zip ..."
unzip -o /workspace/MyRunpodFinetuneScripts.zip -d /workspace
# Предполагается, что внутри появится папка Kohya_Flux*** с нужными файлами.

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

# ------------------ Поток 1 --------------------
(
  echo ">>> [Thread 1] Installing Python packages for huggingface ..."
  pip install huggingface_hub ipywidgets hf_transfer

  export HF_HUB_ENABLE_HF_TRANSFER=1

  echo ">>> [Thread 1] Running Download_Train_Models.py ..."
  dwn_path=$(find /workspace -name "Download_Train_Models.py" | head -n 1)
  if [ -z "$dwn_path" ]; then
    echo "[Thread 1] ERROR: Download_Train_Models.py not found!"
    exit 1
  fi

  python "$dwn_path" --dir /workspace

  echo ">>> [Thread 1] Done."
) &

# ------------------ Поток 2 --------------------
(
  echo ">>> [Thread 2] Installing Python and OS dependencies ..."
  apt update --yes
  yes | apt-get install python3.10-tk
  apt-get install psmisc --yes

  echo ">>> [Thread 2] Cloning 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] Creating Python venv ..."
  python3 -m venv venv
  source venv/bin/activate
  yes | apt-get install python3.10-tk

  echo ">>> [Thread 2] Running 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] Killing processes on port 7860 ..."
  fuser -k 7860/tcp || true

  source venv/bin/activate

  echo ">>> [Thread 2] Updating 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
  pip install pydantic==2.10.6

  echo ">>> [Thread 2] Launching kohya_ss GUI on 0.0.0.0:7860 ..."
  ./gui.sh --listen=0.0.0.0 --share --noverify
) &

###############################################################################
# 3. Создаём общую структуру папок (Models, Images, Logs)
#    по заданию, делаем это в первой части
###############################################################################
echo ">>> [Script 1] Creating folder structure in /workspace/MyLearningDataset ..."
mkdir -p /workspace/MyLearningDataset/Models
mkdir -p /workspace/MyLearningDataset/Images
mkdir -p /workspace/MyLearningDataset/Logs

###############################################################################
# 4. Ожидаем завершения потоков
###############################################################################
echo ">>> [Script 1] Waiting for threads to finish ..."
wait
echo ">>> [Script 1] ALL DONE! Now run 'finetune_analysis.sh' to proceed."