#!/bin/bash
# shellcheck disable=SC2090,SC2086,SC2089,SC1091

#===============================================================================
# mt_dao HPC Environment Installation Script
#===============================================================================
# Description: Sets up the development environment for the mt_dao project on HPC systems
# Author: mt_dao Team
# Version: 1.0
#
# This script is specifically designed for HPC environments and performs:
# 1. Module system configuration (purge and load required modules)
# 2. CUDA 12.1 and cuDNN setup
# 3. UV package manager installation and configuration
# 4. Python environment synchronization
# 5. Flash-attention installation with specific version
# 6. Environment validation and configuration
#
# Usage:
#   ./install_hpc_env.sh [OPTIONS]
#
# Options:
#   -p, --project_path PATH    Specify custom project path (default: $HOME/projects/mt_dao)
#   -h, --help                Show this help message
#
# Requirements:
#   - HPC system with module support
#   - CUDA 12.1 module available
#   - Python 3.10.13
#   - Singularity/Apptainer support
#===============================================================================
# CONSTANTS AND CONFIGURATION
#===============================================================================

SCRIPT_NAME=$(basename "$0")
readonly SCRIPT_NAME
readonly DEFAULT_PROJECT_PATH="$HOME/projects/mt_dao"
readonly REQUIRED_CUDA_VERSION="12.1"
readonly REQUIRED_PYTHON_VERSION="3.10.13"
readonly FLASH_ATTN_VERSION="2.3.2"

#===============================================================================
# UTILITY FUNCTIONS
#===============================================================================

# Print informational messages
log_info() {
	echo "$SCRIPT_NAME: $*"
}

# Print error messages to stderr
log_error() {
	echo "$SCRIPT_NAME: ERROR - $*" >&2
}

# Print warning messages
log_warn() {
	echo "$SCRIPT_NAME: WARNING - $*"
}

# Display usage information
show_usage() {
	cat <<EOF
Usage: $SCRIPT_NAME [OPTIONS]

Sets up the development environment for the mt_dao project on HPC systems.

Options:
    -p, --project_path PATH    Specify custom project path
                              (default: $DEFAULT_PROJECT_PATH)
    -h, --help                Show this help message

Requirements:
    - HPC system with module support
    - CUDA $REQUIRED_CUDA_VERSION module available
    - Python $REQUIRED_PYTHON_VERSION
    - Singularity/Apptainer support

Examples:
    $SCRIPT_NAME                              # Install with default settings
    $SCRIPT_NAME -p /custom/path             # Install with custom path
EOF
}

#===============================================================================
# ARGUMENT PARSING
#===============================================================================

# Initialize variables
PROJECT_PATH="$DEFAULT_PROJECT_PATH"

# Parse command-line options
if ! OPTIONS=$(getopt -o p:h --long project_path:,help -n "$SCRIPT_NAME" -- "$@"); then
	log_error "Failed to parse command-line options"
	show_usage
	exit 1
fi

eval set -- "$OPTIONS"

while true; do
	case "$1" in
	-p | --project_path)
		PROJECT_PATH="$2"
		shift 2
		;;
	-h | --help)
		show_usage
		exit 0
		;;
	--)
		shift
		break
		;;
	*)
		log_error "Unknown option: $1"
		show_usage
		exit 1
		;;
	esac
done

log_info "Installing HPC environment in: $PROJECT_PATH"

#===============================================================================
# MODULE SYSTEM CONFIGURATION
#===============================================================================

log_info "Configuring HPC module system..."

# Enable the module command
if [[ -f /etc/profile.d/modules.sh ]]; then
	# shellcheck disable=SC1091
	source /etc/profile.d/modules.sh
	log_info "Module system initialized ✓"
else
	log_error "Module system not found at /etc/profile.d/modules.sh"
	exit 1
fi

# Clean slate - remove all currently loaded modules
log_info "Purging existing modules..."
module purge

# Load required base modules
log_info "Loading base modules..."
module load singularity/current || log_warn "Could not load singularity/current"
module load rhel8/slurm || log_warn "Could not load rhel8/slurm"
module load dot || log_warn "Could not load dot"
module load rhel8/global || log_warn "Could not load rhel8/global"

# Load MPI (might not be available on all systems)
if module -s load openmpi/4.1.1/gcc-9.4.0-epagguv 2>/dev/null; then
	log_info "OpenMPI module loaded ✓"
else
	log_warn "OpenMPI module not available - this is optional"
fi

# Load CUDA and cuDNN
log_info "Loading CUDA and cuDNN modules..."
if ! module load cuda/12.1; then
	log_error "Failed to load CUDA 12.1 module"
	exit 1
fi

if ! module load cudnn/8.9_cuda-12.1; then
	log_error "Failed to load cuDNN 8.9 for CUDA 12.1 module"
	exit 1
fi

# Load additional modules
log_info "Loading additional modules..."
module load ceuadmin/gettext/0.20 || log_warn "Could not load gettext"
module load vgl/2.5.1/64 || log_warn "Could not load VGL"

log_info "Module configuration completed ✓"
#===============================================================================
# UV PACKAGE MANAGER INSTALLATION
#===============================================================================

# Function to check if UV is installed
check_uv_installation() {
	if command -v uv >/dev/null 2>&1; then
		local uv_version
		uv_version=$(uv --version 2>/dev/null || echo "unknown")
		log_info "UV is already installed: $uv_version ✓"
		return 0
	else
		return 1
	fi
}

# Function to install UV
install_uv() {
	log_info "Installing UV package manager..."

	if ! curl -LsSf https://astral.sh/uv/install.sh | sh; then
		log_error "Failed to install UV"
		exit 1
	fi

	# Reload shell environment to pick up UV
	# shellcheck disable=SC1090
	source "$HOME/.cargo/env" 2>/dev/null || true

	log_info "UV installation completed ✓"
}

# Install UV if not present
if ! check_uv_installation; then
	install_uv
fi

#===============================================================================
# PYTHON ENVIRONMENT SETUP
#===============================================================================

log_info "Setting up Python environment..."

# Upgrade pip first
log_info "Upgrading pip..."
if ! uv pip install --upgrade pip; then
	log_error "Failed to upgrade pip"
	exit 1
fi

# Install monitoring utilities
log_info "Installing monitoring utilities..."
if command -v snap >/dev/null 2>&1; then
	sudo snap install bpytop 2>/dev/null || log_warn "Failed to install bpytop via snap"
else
	log_warn "Snap not available, skipping bpytop installation"
fi

if ! uv pip install nvitop; then
	log_warn "Failed to install nvitop"
fi

# Install cmake
log_info "Installing cmake..."
if ! uv pip install cmake; then
	log_error "Failed to install cmake"
	exit 1
fi

#===============================================================================
# PROJECT ENVIRONMENT SYNCHRONIZATION
#===============================================================================

log_info "Synchronizing project environment..."

# Change to project directory
if [[ ! -d $PROJECT_PATH ]]; then
	log_error "Project directory does not exist: $PROJECT_PATH"
	exit 1
fi

cd "$PROJECT_PATH" || {
	log_error "Failed to change to project directory: $PROJECT_PATH"
	exit 1
}

# Synchronize UV environment
# Note: GIT_LFS_SKIP_SMUDGE prevents downloading large files during sync
log_info "Running UV sync (this may take a few minutes)..."
if ! GIT_LFS_SKIP_SMUDGE=true uv sync -q; then
	log_error "Failed to synchronize UV environment"
	exit 1
fi

# Activate virtual environment if available
if [[ -n ${VIRTUAL_ENV:-} && -f "$VIRTUAL_ENV/bin/activate" ]]; then
	# shellcheck disable=SC1091
	source "$VIRTUAL_ENV/bin/activate"
	log_info "Virtual environment activated ✓"
fi
#===============================================================================
# SYSTEM VALIDATION
#===============================================================================

# Function to validate CUDA installation
validate_cuda() {
	log_info "Validating CUDA installation..."

	local nvcc_output
	if ! nvcc_output=$(nvcc -V 2>/dev/null); then
		log_error "NVCC not found or not working"
		return 1
	fi

	if [[ $nvcc_output == *"release $REQUIRED_CUDA_VERSION"* ]]; then
		log_info "CUDA $REQUIRED_CUDA_VERSION detected ✓"
		return 0
	else
		log_error "CUDA $REQUIRED_CUDA_VERSION not detected"
		log_error "Current NVCC output: $nvcc_output"
		return 1
	fi
}

# Function to validate Python version
validate_python() {
	log_info "Validating Python installation..."

	local python_output
	if ! python_output=$(python --version 2>/dev/null); then
		log_error "Python not found or not working"
		return 1
	fi

	if [[ $python_output == *"$REQUIRED_PYTHON_VERSION"* ]]; then
		log_info "Python $REQUIRED_PYTHON_VERSION detected ✓"
		return 0
	else
		log_error "Python $REQUIRED_PYTHON_VERSION not detected"
		log_error "Current Python version: $python_output"
		return 1
	fi
}

# Run validations
if ! validate_cuda; then
	log_error "CUDA validation failed"
	exit 1
fi

if ! validate_python; then
	log_error "Python validation failed"
	exit 1
fi

#===============================================================================
# FLASH-ATTENTION INSTALLATION
#===============================================================================

log_info "Setting up Flash-Attention..."

# Function to install flash-attention
install_flash_attention() {
	if uv run pip list | grep -q flash-attn; then
		log_info "Flash-attention is already installed ✓"
		return 0
	fi

	log_info "Installing flash-attention $FLASH_ATTN_VERSION..."
	log_info "Note: This compilation may take 10-15 minutes..."

	if ! uv run uv pip install -q "flash-attn==$FLASH_ATTN_VERSION" --no-build-isolation; then
		log_error "Failed to install flash-attention"
		return 1
	fi

	log_info "Flash-attention installation completed ✓"
	return 0
}

if ! install_flash_attention; then
	log_error "Flash-attention setup failed"
	exit 1
fi

#===============================================================================
# ENVIRONMENT CONFIGURATION
#===============================================================================

log_info "Configuring environment variables..."

# Configure Python warnings for HPC environment
# Reduce noise from deprecation and resource warnings common in HPC setups
export PYTHONWARNINGS="ignore::DeprecationWarning,ignore::ResourceWarning"
log_info "Python warnings configured for HPC environment"

# Set TRITON_CACHE_DIR to be user-specific
TRITON_CACHE_DIR="/home/$(whoami)/.triton_cache"
export TRITON_CACHE_DIR
log_info "Set TRITON_CACHE_DIR to: $TRITON_CACHE_DIR"

#===============================================================================
# INSTALLATION COMPLETION
#===============================================================================

log_info "HPC environment installation completed successfully! ✅"
log_info ""
log_info "Summary:"
log_info "  Project path: $PROJECT_PATH"
log_info "  CUDA version: $REQUIRED_CUDA_VERSION"
log_info "  Python version: $REQUIRED_PYTHON_VERSION"
log_info "  Flash-attention: $FLASH_ATTN_VERSION"
log_info "  Python environment: synchronized"
log_info "  HPC modules: loaded"
log_info "  Environment variables: configured"
log_info ""
log_info "You can now start using the mt_dao development environment on HPC!"
log_info "For more information, see the project documentation."
