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

# =============================================================================
# Flash Attention Installer Script
# =============================================================================
#
# This script automatically installs Flash Attention based on the detected
# hardware architecture and GPU type. It supports:
# - aarch64 architecture (ARM-based systems)
# - AMD GPUs with ROCm
# - NVIDIA GPUs with CUDA (standard x86_64)
#
# Usage:
#   ./flash_attn_installer.sh [options]
#
# Options:
#   -p, --project_path PATH    Specify the project path (default: $HOME/projects/mt_dao)
#
# Examples:
#   ./flash_attn_installer.sh
#   ./flash_attn_installer.sh -p /custom/path/to/mt_dao
#   ./flash_attn_installer.sh --project_path /custom/path/to/mt_dao
#
# =============================================================================

# Default configuration
readonly DEFAULT_PROJECT_PATH="$HOME/projects/mt_dao"
PROJECT_PATH="$DEFAULT_PROJECT_PATH"

# Script name for consistent logging
readonly SCRIPT_NAME="flash_attn_installer.sh"

# Flash Attention repository, branch, and PyTorch ROCm architecture
FA_REPO=${FA_REPO:-"https://github.com/Dao-AILab/flash-attention.git"}
AMD_FA_BRANCH=${AMD_FA_BRANCH:-"1a7f4dfa"}
FA_BRANCH=${FA_BRANCH:-"v2.8.2"}
PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH:-"gfx90a;gfx942;gfx1100;gfx1101;gfx1200;gfx1201"}
INSTALL_FAv3=$(uv run python -c "from mt_dao.utils import get_cuda_compute_capability; a = get_cuda_compute_capability(); print('true' if a and a[0] >= 9 else 'false')")
CUDA_HOME=/usr/local/cuda-12.8
export CUDA_HOME

# Function to log messages with script name prefix
log_info() {
	echo "$SCRIPT_NAME: $1"
}

log_error() {
	echo "$SCRIPT_NAME: ERROR: $1" >&2
}

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

Install Flash Attention based on detected hardware architecture.

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

EXAMPLES:
    $SCRIPT_NAME
    $SCRIPT_NAME -p /custom/path/to/mt_dao
    $SCRIPT_NAME --project_path /custom/path/to/mt_dao
EOF
}

# 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"

log_info "Parsing command-line options: $OPTIONS"

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

# Validate remaining arguments
if [[ $# -gt 0 ]]; then
	log_info "Additional arguments provided: $*"
fi

log_info "Installing Flash Attention in PROJECT_PATH=$PROJECT_PATH"

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

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

# =============================================================================
# Hardware Detection and Installation Logic
# =============================================================================

# Function to detect system architecture
detect_architecture() {
	uname -m
}

# Function to detect AMD GPU presence
detect_amd_gpu() {
	if command -v rocm-smi >/dev/null 2>&1; then
		local amd_gpu_count
		amd_gpu_count=$(rocm-smi --showproductname 2>/dev/null | grep -c -i 'amd' || echo "0")
		echo "$amd_gpu_count"
	else
		echo "0"
	fi
}

# Function to build Flash Attention for AMD GPUs from source
build_flash_attention_amd() {

	log_info "Building Flash Attention for AMD/ROCm from source"
	log_info "Repository: $FA_REPO"
	log_info "Branch: $FA_BRANCH"
	log_info "ROCm architectures: $PYTORCH_ROCM_ARCH"

	# Clean up any existing flash-attention directory
	rm -rf flash-attention

	log_info "Cloning Flash Attention repository"
	if ! git clone "$FA_REPO" flash-attention; then
		log_error "Failed to clone Flash Attention repository"
		return 1
	fi

	log_info "Switching to branch: $FA_BRANCH"
	if ! (cd flash-attention && git checkout "$FA_BRANCH"); then
		log_error "Failed to switch to branch: $FA_BRANCH"
		rm -rf flash-attention
		return 1
	fi

	log_info "Updating submodules"
	if ! (cd flash-attention && git submodule update --init); then
		log_error "Failed to update submodules"
		rm -rf flash-attention
		return 1
	fi

	log_info "Building Flash Attention wheel with $max_jobs parallel jobs"

	# Set GPU architectures for AMD ROCm (filter out unsupported architectures)
	local gpu_archs
	gpu_archs="${PYTORCH_ROCM_ARCH//;gfx1[0-9][0-9][0-9]/}"
	log_info "Using GPU architectures: $gpu_archs"

	# Build Flash Attention wheel using original commands
	cd flash-attention || {
		log_error "Failed to change to flash-attention directory"
		rm -rf flash-attention
		return 1
	}

	# Create dist directory
	mkdir -p dist

	if ! eval "GPU_ARCHS='$gpu_archs' uv run python3 setup.py bdist_wheel --dist-dir=dist"; then
		log_error "Failed to build Flash Attention wheel"
		cd .. && rm -rf flash-attention
		return 1
	fi

	# Install the built wheel
	log_info "Installing Flash Attention from built wheel"
	if ! uv run pip install dist/*.whl; then
		log_error "Failed to install Flash Attention wheel"
		cd .. && rm -rf flash-attention
		return 1
	fi

	log_info "Cleaning up build directory"
	cd .. && rm -rf flash-attention

	log_info "Flash Attention for AMD/ROCm built and installed successfully"
	return 0
}

# Function to install Flash Attention v3 from source
install_flash_attention_v3_from_source() {
	local branch="$1"
	local max_jobs="$2"
	local env_vars="$3"

	log_info "Cloning Flash Attention repository (branch: $branch)"
	rm -rf flash-attention
	# if ! git clone --recursive $FA_REPO; then
	if ! git clone $FA_REPO; then
		log_error "Failed to clone Flash Attention repository"
		return 1
	fi

	log_info "Switching to branch: $branch"
	if ! (cd flash-attention && git switch -C "$branch"); then
		log_error "Failed to switch to branch: $branch"
		rm -rf flash-attention
		return 1
	fi

	log_info "Updating submodules"
	if ! (cd flash-attention && git submodule update --init); then
		log_error "Failed to update submodules"
		rm -rf flash-attention
		return 1
	fi

	log_info "Installing Flash Attention v3 from source with $max_jobs parallel jobs"
	if [[ $INSTALL_FAv3 == true ]]; then
		log_info "Installing FlashAttention v3"
		if ! (cd flash-attention/hopper && uv run python setup.py install); then
			log_error "Failed to install FlashAttention v3"
			return 1
		fi
	fi

	log_info "Cleaning up cloned repository"
	rm -rf flash-attention
	return 0
}

# Function to install Flash Attention from source
install_flash_attention_from_source() {
	local branch="$1"
	local max_jobs="$2"
	local env_vars="$3"

	log_info "Cloning Flash Attention repository (branch: $branch)"
	rm -rf flash-attention
	# if ! git clone --recursive $FA_REPO; then
	if ! git clone $FA_REPO; then
		log_error "Failed to clone Flash Attention repository"
		return 1
	fi

	log_info "Switching to branch: $branch"
	if ! (cd flash-attention && git switch -C "$branch"); then
		log_error "Failed to switch to branch: $branch"
		rm -rf flash-attention
		return 1
	fi

	log_info "Updating submodules"
	if ! (cd flash-attention && git submodule update --init); then
		log_error "Failed to update submodules"
		rm -rf flash-attention
		return 1
	fi

	log_info "Installing Flash Attention from source with $max_jobs parallel jobs"
	if [[ -n $env_vars ]]; then
		log_info "Using environment variables: $env_vars"
		if ! (cd flash-attention && $env_vars MAX_JOBS=$max_jobs uv run python setup.py install); then
			log_error "Failed to install Flash Attention from source"
			rm -rf flash-attention
			return 1
		fi
	else
		if ! (cd flash-attention && MAX_JOBS=$max_jobs uv run python setup.py install); then
			log_error "Failed to install Flash Attention from source"
			rm -rf flash-attention
			return 1
		fi
	fi

	log_info "Cleaning up cloned repository"
	rm -rf flash-attention
	return 0
}

# Function to install Flash Attention via pre-built wheel
install_flash_attention_wheel() {
	local wheel_url="$1"

	log_info "Installing Flash Attention from pre-built wheel"
	if uv pip install -q "$wheel_url"; then
		return 0
	else
		log_error "Failed to install Flash Attention from wheel"
		return 1
	fi
}

# Function to verify Flash Attention installation
verify_installation() {
	log_info "Verifying Flash Attention installation"
	if uv pip list | grep -q "flash"; then
		log_info "Flash Attention installation verified successfully"
		return 0
	else
		log_error "Flash Attention installation verification failed"
		return 1
	fi
}

# Function to reinstall pre-commit hooks
reinstall_precommit_hooks() {
	log_info "Reinstalling pre-commit hooks after Flash Attention setup"
	if ! uv run pre-commit install; then
		log_error "Failed to reinstall pre-commit hooks"
		return 1
	fi
	return 0
}

# =============================================================================
# Main Installation Logic
# =============================================================================

log_info "Starting Flash Attention installation process"

# Detect system architecture
ARCHITECTURE=$(detect_architecture)
log_info "Detected architecture: $ARCHITECTURE"

# Install Flash Attention based on architecture and GPU type
if [[ $ARCHITECTURE == 'aarch64' ]]; then
	# ARM64 Architecture (e.g., Apple Silicon, ARM servers)
	log_info "Installing Flash Attention for aarch64 architecture"
	log_info "This process may take a while as it requires compilation from source..."

	if install_flash_attention_from_source $FA_BRANCH "10" "FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE"; then
		log_info "Flash Attention successfully installed for aarch64"
	else
		log_error "Flash Attention installation failed for aarch64"
		exit 1
	fi

else
	# x86_64 Architecture - Check for GPU type
	AMD_GPU_COUNT=$(detect_amd_gpu)

	# Use the minimum between (nproc - 1) and 16
	MAX_PARALLEL_JOBS=$(($(nproc) - 1))
	if ((MAX_PARALLEL_JOBS > 16)); then
		MAX_PARALLEL_JOBS=16
	fi
	log_info "Using $MAX_PARALLEL_JOBS parallel compilation jobs"

	if [[ $AMD_GPU_COUNT -gt 0 ]]; then
		# AMD GPU Environment
		log_info "Detected $AMD_GPU_COUNT AMD GPU(s) - installing Flash Attention for AMD/ROCm"
		log_info "This process may take a while as it requires compilation from source..."

		# Use the dedicated AMD build function
		if build_flash_attention_amd; then
			log_info "Flash Attention successfully installed for AMD/ROCm"
		else
			log_error "Flash Attention installation failed for AMD/ROCm"
			exit 1
		fi

	else
		# NVIDIA GPU Environment (standard x86_64)
		log_info "Installing Flash Attention for standard x86_64 environment with NVIDIA GPUs"

		# Pre-built wheel URL for CUDA 12.x + PyTorch 2.7
		WHEEL_URL="https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.2/flash_attn-2.8.2+cu12torch2.7cxx11abiTRUE-cp312-cp312-linux_x86_64.whl"

		if install_flash_attention_wheel "$WHEEL_URL"; then
			log_info "Flash Attention successfully installed from pre-built wheel"
		else
			log_info "Pre-built wheel installation failed, falling back to source compilation"
			log_info "This process may take a while as it requires compilation from source..."

			MAX_PARALLEL_JOBS=$(($(nproc) - 1))
			log_info "Using $MAX_PARALLEL_JOBS parallel compilation jobs"

			if install_flash_attention_from_source $FA_BRANCH "$MAX_PARALLEL_JOBS" ""; then
				log_info "Flash Attention successfully installed from source as fallback"
			else
				log_error "Both wheel and source installation methods failed"
				exit 1
			fi
		fi

		if install_flash_attention_v3_from_source $FA_BRANCH "$MAX_PARALLEL_JOBS" ""; then
			log_info "Flash Attention v3 successfully installed from source as fallback"
		else
			log_error "Flash Attention v3 source installation method failed"
			exit 1
		fi
	fi
fi

# Verify installation success
if verify_installation; then
	log_info "Flash Attention installation completed successfully"

	# Reinstall pre-commit hooks
	if reinstall_precommit_hooks; then
		log_info "Pre-commit hooks reinstalled successfully"
	else
		log_error "Failed to reinstall pre-commit hooks (non-critical)"
	fi

	log_info "Flash Attention installer finished successfully"
else
	log_error "Flash Attention installation failed final verification"
	exit 1
fi
