#!/bin/bash
# Flux Identity LoRA Training Environment - Complete Setup Script
# Target: RunPod with NVIDIA A100 80GB

set -e

echo "=============================================="
echo "Flux Identity LoRA Training Environment Setup"
echo "=============================================="

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

log_info() { echo -e "${GREEN}[INFO]${NC} $1"; }
log_warn() { echo -e "${YELLOW}[WARN]${NC} $1"; }
log_error() { echo -e "${RED}[ERROR]${NC} $1"; }

# ============================================
# Phase 1: System Inspection
# ============================================
echo ""
log_info "Phase 1: System Inspection"
echo "--------------------------------------------"

echo "GPU Information:"
nvidia-smi --query-gpu=name,memory.total,memory.free,driver_version --format=csv || log_warn "nvidia-smi not available"

echo ""
echo "CUDA Version:"
nvcc --version 2>/dev/null || log_warn "nvcc not available"

echo ""
echo "CPU Information:"
cat /proc/cpuinfo | grep "model name" | head -1 || echo "CPU info not available"
echo "CPU Cores: $(nproc)"

echo ""
echo "Memory:"
free -h

echo ""
echo "Disk Space:"
df -h /workspace 2>/dev/null || df -h .

# ============================================
# Phase 2: Project Structure
# ============================================
echo ""
log_info "Phase 2: Creating Project Structure"
echo "--------------------------------------------"

PROJECT_DIR="/workspace/flux-project"
mkdir -p ${PROJECT_DIR}/{models/{flux,lora,vae,clip,controlnet,ip-adapter},datasets/identity/{images,processed},output,scripts,logs,configs}

log_info "Project structure created at ${PROJECT_DIR}"

# ============================================
# Phase 3: System Dependencies
# ============================================
echo ""
log_info "Phase 3: Installing System Dependencies"
echo "--------------------------------------------"

apt-get update && apt-get install -y \
    git git-lfs wget curl \
    libgl1-mesa-glx libglib2.0-0 \
    python3.10 python3.10-venv python3-pip \
    aria2 ffmpeg

git lfs install

# ============================================
# Phase 4: Python Virtual Environment
# ============================================
echo ""
log_info "Phase 4: Setting Up Python Virtual Environment"
echo "--------------------------------------------"

VENV_PATH="/workspace/venv"
python3.10 -m venv ${VENV_PATH}
source ${VENV_PATH}/bin/activate

pip install --upgrade pip wheel setuptools

# ============================================
# Phase 5: Core ML Packages
# ============================================
echo ""
log_info "Phase 5: Installing Core ML Packages"
echo "--------------------------------------------"

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
pip install xformers

# ============================================
# Phase 6: Training Dependencies
# ============================================
echo ""
log_info "Phase 6: Installing Training Dependencies"
echo "--------------------------------------------"

pip install accelerate transformers diffusers safetensors
pip install bitsandbytes scipy tensorboard
pip install lion-pytorch prodigyopt dadaptation
pip install opencv-python pillow numpy
pip install toml pyyaml tqdm rich
pip install imagehash
pip install huggingface_hub

# ============================================
# Phase 7: CUDA Verification
# ============================================
echo ""
log_info "Phase 7: Verifying CUDA Installation"
echo "--------------------------------------------"

python -c "
import torch
print(f'PyTorch version: {torch.__version__}')
print(f'CUDA available: {torch.cuda.is_available()}')
if torch.cuda.is_available():
    print(f'CUDA version: {torch.version.cuda}')
    print(f'GPU: {torch.cuda.get_device_name(0)}')
    print(f'VRAM: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB')
else:
    print('WARNING: CUDA not available!')
"

# ============================================
# Phase 8: ComfyUI Installation
# ============================================
echo ""
log_info "Phase 8: Installing ComfyUI"
echo "--------------------------------------------"

cd /workspace
if [ ! -d "ComfyUI" ]; then
    git clone https://github.com/comfyanonymous/ComfyUI.git
fi
cd ComfyUI
pip install -r requirements.txt

# Install ComfyUI Manager
cd custom_nodes
if [ ! -d "ComfyUI-Manager" ]; then
    git clone https://github.com/ltdrdata/ComfyUI-Manager.git
fi

# Create model symlinks
log_info "Creating model symlinks..."
ln -sf ${PROJECT_DIR}/models/flux /workspace/ComfyUI/models/unet
ln -sf ${PROJECT_DIR}/models/lora /workspace/ComfyUI/models/loras
ln -sf ${PROJECT_DIR}/models/vae /workspace/ComfyUI/models/vae
ln -sf ${PROJECT_DIR}/models/clip /workspace/ComfyUI/models/clip
ln -sf ${PROJECT_DIR}/models/controlnet /workspace/ComfyUI/models/controlnet

log_info "ComfyUI installed successfully"

# ============================================
# Phase 9: kohya_ss/sd-scripts Installation
# ============================================
echo ""
log_info "Phase 9: Installing kohya_ss/sd-scripts"
echo "--------------------------------------------"

cd /workspace
if [ ! -d "sd-scripts" ]; then
    git clone https://github.com/kohya-ss/sd-scripts.git
fi
cd sd-scripts
pip install -r requirements.txt

# Verify installation
python -c "from library import flux_train_utils; print('sd-scripts Flux support verified')" || log_warn "sd-scripts Flux verification failed"

log_info "sd-scripts installed successfully"

# ============================================
# Phase 10: Copy Scripts to Project
# ============================================
echo ""
log_info "Phase 10: Setting Up Project Scripts"
echo "--------------------------------------------"

# Make scripts executable
chmod +x ${PROJECT_DIR}/scripts/*.sh 2>/dev/null || true
chmod +x ${PROJECT_DIR}/scripts/*.py 2>/dev/null || true

# ============================================
# Final Verification
# ============================================
echo ""
log_info "Running Final Verification"
echo "--------------------------------------------"

echo "1. CUDA Check:"
python -c "import torch; print(f'   CUDA: {torch.cuda.is_available()}')"

echo "2. sd-scripts Check:"
python -c "from library import flux_train_utils; print('   sd-scripts: OK')" 2>/dev/null || echo "   sd-scripts: FAILED"

echo "3. Directory Structure:"
ls -la ${PROJECT_DIR}/

echo ""
echo "=============================================="
log_info "Setup Complete!"
echo "=============================================="
echo ""
echo "Next Steps:"
echo "1. Download FLUX.1 Dev model:"
echo "   huggingface-cli login"
echo "   huggingface-cli download black-forest-labs/FLUX.1-dev --local-dir ${PROJECT_DIR}/models/flux"
echo ""
echo "2. Download text encoders:"
echo "   huggingface-cli download openai/clip-vit-large-patch14 --local-dir ${PROJECT_DIR}/models/clip/clip-l"
echo "   huggingface-cli download google/t5-v1_1-xxl --local-dir ${PROJECT_DIR}/models/clip/t5-xxl"
echo ""
echo "3. Place your training images in:"
echo "   ${PROJECT_DIR}/datasets/identity/images/"
echo ""
echo "4. Start ComfyUI:"
echo "   ${PROJECT_DIR}/scripts/start_comfyui.sh"
echo ""
echo "5. Start Training:"
echo "   ${PROJECT_DIR}/scripts/start_training.sh"
echo ""