#!/bin/bash
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

# FastGen Offline Environment Setup Script
#
# This script sets up the FastGen environment on an offline machine using
# pre-downloaded packages from prepare_offline_package.py
#
# Usage:
#     cd /path/to/offline_package
#     bash setup_offline_env.sh
#
# Options:
#     --env-name NAME     Conda environment name (default: sf)
#     --python-version    Python version (default: 3.12)
#     --skip-install      Skip pip package installation
#     --skip-symlinks     Skip HuggingFace cache symlink setup
#     --verify-only       Only run verification, don't install

set -e

# Default values
ENV_NAME="sf"
PYTHON_VERSION="3.12"
SKIP_INSTALL=false
SKIP_SYMLINKS=false
VERIFY_ONLY=false

# Parse arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --env-name)
            ENV_NAME="$2"
            shift 2
            ;;
        --python-version)
            PYTHON_VERSION="$2"
            shift 2
            ;;
        --skip-install)
            SKIP_INSTALL=true
            shift
            ;;
        --skip-symlinks)
            SKIP_SYMLINKS=true
            shift
            ;;
        --verify-only)
            VERIFY_ONLY=true
            shift
            ;;
        -h|--help)
            echo "FastGen Offline Environment Setup"
            echo ""
            echo "Usage: bash setup_offline_env.sh [OPTIONS]"
            echo ""
            echo "Options:"
            echo "  --env-name NAME      Conda environment name (default: sf)"
            echo "  --python-version     Python version (default: 3.12)"
            echo "  --skip-install       Skip pip package installation"
            echo "  --skip-symlinks      Skip HuggingFace cache symlink setup"
            echo "  --verify-only        Only run verification, don't install"
            echo "  -h, --help           Show this help message"
            exit 0
            ;;
        *)
            echo "Unknown option: $1"
            exit 1
            ;;
    esac
done

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

print_step() {
    echo -e "\n${GREEN}[STEP]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

print_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

# Determine script directory (where offline_package contents are)
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# If we're inside offline_package/FastGen/scripts, go up to offline_package
if [[ -d "$SCRIPT_DIR/../pip_wheels" ]]; then
    PACKAGE_DIR="$SCRIPT_DIR/.."
elif [[ -d "$SCRIPT_DIR/pip_wheels" ]]; then
    PACKAGE_DIR="$SCRIPT_DIR"
elif [[ -d "$SCRIPT_DIR/FastGen/scripts" ]]; then
    # Script is in offline_package root
    PACKAGE_DIR="$SCRIPT_DIR"
else
    # Assume we're in the offline_package directory
    PACKAGE_DIR="$(pwd)"
fi

PACKAGE_DIR="$(cd "$PACKAGE_DIR" && pwd)"

echo "============================================================"
echo "FastGen Offline Environment Setup"
echo "============================================================"
echo "Package directory: $PACKAGE_DIR"
echo "Environment name: $ENV_NAME"
echo "Python version: $PYTHON_VERSION"
echo ""

# Verify package structure
print_step "Verifying package structure..."

REQUIRED_DIRS=(
    "pip_wheels"
    "hf_models/Wan2.1-T2V-1.3B-Diffusers"
    "checkpoints/Self-Forcing"
    "FastGen"
)

MISSING_DIRS=()
for dir in "${REQUIRED_DIRS[@]}"; do
    if [[ ! -d "$PACKAGE_DIR/$dir" ]]; then
        MISSING_DIRS+=("$dir")
    fi
done

if [[ ${#MISSING_DIRS[@]} -gt 0 ]]; then
    print_warning "Missing directories:"
    for dir in "${MISSING_DIRS[@]}"; do
        echo "  - $dir"
    done
    echo ""
    echo "The package may be incomplete. Some components may need to be downloaded separately."
fi

if [[ "$VERIFY_ONLY" == true ]]; then
    print_step "Running verification only..."
    # Jump to verification section
    if command -v python &> /dev/null; then
        echo "Testing Python imports..."
        python -c "import torch; print(f'PyTorch: {torch.__version__}')" 2>/dev/null || echo "PyTorch not available"
        python -c "from diffusers import WanPipeline; print('WanPipeline: OK')" 2>/dev/null || echo "WanPipeline not available"
        python -c "from fastgen.networks.Wan.network import Wan; print('FastGen Wan: OK')" 2>/dev/null || echo "FastGen Wan not available"
    fi
    exit 0
fi

# Check for conda
print_step "Checking for conda..."
if command -v conda &> /dev/null; then
    print_success "Conda found: $(conda --version)"
else
    print_error "Conda not found. Please install Miniconda or Anaconda first."
    echo "Download from: https://docs.conda.io/en/latest/miniconda.html"
    exit 1
fi

# Create conda environment
print_step "Creating conda environment '$ENV_NAME' with Python $PYTHON_VERSION..."

if conda env list | grep -q "^$ENV_NAME "; then
    print_warning "Environment '$ENV_NAME' already exists."
    read -p "Do you want to remove and recreate it? (y/N) " -n 1 -r
    echo
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        conda env remove -n "$ENV_NAME" -y
        conda create -n "$ENV_NAME" python="$PYTHON_VERSION" -y
    fi
else
    conda create -n "$ENV_NAME" python="$PYTHON_VERSION" -y
fi

# Activate environment
print_step "Activating conda environment..."

# Source conda.sh to enable conda activate in script
CONDA_BASE=$(conda info --base)
source "$CONDA_BASE/etc/profile.d/conda.sh"
conda activate "$ENV_NAME"

echo "Python location: $(which python)"
echo "Python version: $(python --version)"

# Install pip packages
if [[ "$SKIP_INSTALL" == false ]] && [[ -d "$PACKAGE_DIR/pip_wheels" ]]; then
    print_step "Installing pip packages from local wheels..."

    # Install torch first (important for dependency resolution)
    TORCH_WHEEL=$(find "$PACKAGE_DIR/pip_wheels" -name "torch-*.whl" | head -1)
    if [[ -n "$TORCH_WHEEL" ]]; then
        echo "Installing PyTorch first..."
        pip install --no-index --find-links="$PACKAGE_DIR/pip_wheels" torch
    fi

    # Install all packages
    pip install --no-index --find-links="$PACKAGE_DIR/pip_wheels" \
        diffusers transformers accelerate safetensors \
        scipy einops hydra-core imageio tqdm webdataset av \
        loguru boto3 timm ftfy opencv-python-headless \
        sentencepiece nvidia-ml-py 'numpy<2.0.0' wandb || {
            print_warning "Some packages may have failed to install from wheels."
            echo "You may need to install them manually or download source distributions."
        }

    print_success "Pip packages installed"
else
    print_warning "Skipping pip package installation"
fi

# Install FastGen
if [[ -d "$PACKAGE_DIR/FastGen" ]]; then
    print_step "Installing FastGen in editable mode..."
    cd "$PACKAGE_DIR/FastGen"
    pip install -e . --no-deps
    cd "$PACKAGE_DIR"
    print_success "FastGen installed"
else
    print_warning "FastGen source not found, skipping installation"
fi

# Setup HuggingFace cache symlinks
if [[ "$SKIP_SYMLINKS" == false ]]; then
    print_step "Setting up HuggingFace cache structure..."

    HF_HOME="$PACKAGE_DIR/hf_cache"
    mkdir -p "$HF_HOME/hub"

    # Create symlinks with HuggingFace naming convention
    if [[ -d "$PACKAGE_DIR/hf_models/Wan2.1-T2V-1.3B-Diffusers" ]]; then
        LINK_PATH="$HF_HOME/hub/models--Wan-AI--Wan2.1-T2V-1.3B-Diffusers"
        if [[ ! -L "$LINK_PATH" ]]; then
            ln -sf "$PACKAGE_DIR/hf_models/Wan2.1-T2V-1.3B-Diffusers" "$LINK_PATH"
            echo "Created symlink: $LINK_PATH"
        fi
    fi

    print_success "HuggingFace cache structure created"
fi

# Create environment setup script
print_step "Creating environment activation script..."

ENV_SCRIPT="$PACKAGE_DIR/activate_fastgen.sh"
cat > "$ENV_SCRIPT" << EOF
#!/bin/bash
# FastGen Environment Activation Script
# Source this file to set up the environment:
#   source activate_fastgen.sh

# Activate conda environment
CONDA_BASE=\$(conda info --base)
source "\$CONDA_BASE/etc/profile.d/conda.sh"
conda activate $ENV_NAME

# Set environment variables for offline mode
export HF_HOME="$PACKAGE_DIR/hf_cache"
export HF_HUB_OFFLINE=1
export TRANSFORMERS_OFFLINE=1
export CKPT_ROOT_DIR="$PACKAGE_DIR/checkpoints"
export FASTGEN_OUTPUT_ROOT="$PACKAGE_DIR/outputs"

# Create outputs directory if it doesn't exist
mkdir -p "\$FASTGEN_OUTPUT_ROOT"

echo "FastGen environment activated"
echo "  HF_HOME: \$HF_HOME"
echo "  CKPT_ROOT_DIR: \$CKPT_ROOT_DIR"
echo "  FASTGEN_OUTPUT_ROOT: \$FASTGEN_OUTPUT_ROOT"
EOF

chmod +x "$ENV_SCRIPT"
print_success "Created: $ENV_SCRIPT"

# Add to .bashrc suggestion
echo ""
echo "To activate automatically, add to your ~/.bashrc:"
echo "  source $ENV_SCRIPT"

# Run verification
print_step "Running verification..."

# Set environment variables for verification
export HF_HOME="$PACKAGE_DIR/hf_cache"
export HF_HUB_OFFLINE=1
export TRANSFORMERS_OFFLINE=1
export CKPT_ROOT_DIR="$PACKAGE_DIR/checkpoints"
export FASTGEN_OUTPUT_ROOT="$PACKAGE_DIR/outputs"
mkdir -p "$FASTGEN_OUTPUT_ROOT"

echo ""
echo "Testing Python imports..."

# Test PyTorch
python -c "import torch; print(f'  PyTorch: {torch.__version__}')" 2>/dev/null && \
    print_success "PyTorch OK" || print_warning "PyTorch import failed"

# Test CUDA
python -c "import torch; print(f'  CUDA available: {torch.cuda.is_available()}')" 2>/dev/null || true

# Test diffusers
python -c "from diffusers import WanPipeline; print('  WanPipeline: importable')" 2>/dev/null && \
    print_success "Diffusers OK" || print_warning "Diffusers import failed"

# Test FastGen
python -c "from fastgen.networks.Wan.network import Wan; print('  FastGen Wan: importable')" 2>/dev/null && \
    print_success "FastGen OK" || print_warning "FastGen import failed"

# Check checkpoint file
if [[ -f "$CKPT_ROOT_DIR/Self-Forcing/checkpoints/ode_init.pt" ]]; then
    print_success "Self-Forcing checkpoint found"
else
    print_warning "Self-Forcing checkpoint not found at expected path"
    echo "  Expected: $CKPT_ROOT_DIR/Self-Forcing/checkpoints/ode_init.pt"
fi

echo ""
echo "============================================================"
print_success "Setup complete!"
echo "============================================================"
echo ""
echo "To use FastGen:"
echo "  1. Activate the environment:"
echo "     source $ENV_SCRIPT"
echo ""
echo "  2. Run data-free training:"
echo "     cd $PACKAGE_DIR/FastGen"
echo "     python train.py --config=fastgen/configs/experiments/WanT2V/config_sf_datafree.py"
echo ""
echo "============================================================"