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

#===============================================================================
# mt_dao Environment Installation Script
#===============================================================================
# Description: Sets up the development environment for the mt_dao project
# Author: mt_dao Team
# Version: 1.0
#
# This script performs the following operations:
# 1. Parses command-line arguments
# 2. Synchronizes the Python environment using UV
# 3. Installs pre-commit hooks
# 4. Detects and configures CUDA (if available)
# 5. Installs flash-attention (GPU mode only)
# 6. Sets up environment variables for Triton and Ray
#
# Usage:
#   ./install_env.sh [OPTIONS] [no_cuda]
#
# Options:
#   -p, --project_path PATH    Specify custom project path (default: $HOME/projects/mt_dao)
#
# Arguments:
#   no_cuda                    Skip CUDA-related installations and configurations
#
# Examples:
#   ./install_env.sh                             # Install with default settings
#   ./install_env.sh -p /custom/path             # Install with custom path
#   ./install_env.sh no_cuda                     # Install without CUDA support
#===============================================================================
# CONSTANTS AND CONFIGURATION
#===============================================================================

# Default configuration
SCRIPT_NAME=$(basename "$0")
readonly SCRIPT_NAME
readonly DEFAULT_PROJECT_PATH="$HOME/projects/mt_dao"
readonly MIN_CUDA_MAJOR=12
readonly MIN_CUDA_MINOR=4

# CUDA installation paths (ordered by preference - latest versions first)
readonly CUDA_PATHS=(
	"/usr/local/cuda-12.9"
	"/usr/local/cuda-12.8"
	"/usr/local/cuda-12.7"
	"/usr/local/cuda-12.6"
	"/usr/local/cuda-12.5"
	"/usr/local/cuda-12.4"
	"/usr/local/cuda"
	"/opt/cuda"
)

#===============================================================================
# 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] [no_cuda]

Sets up the development environment for the mt_dao project.

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

Arguments:
    no_cuda                   Skip CUDA-related installations and configurations

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

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

# Initialize variables
PROJECT_PATH="$DEFAULT_PROJECT_PATH"
SKIP_CUDA=false

# 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 "PROJECT_PATH=$PROJECT_PATH"
export VIRTUAL_ENV="$PROJECT_PATH/.venv"
log_info "VIRTUAL_ENV=$VIRTUAL_ENV"

# 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
}
AMD_GPU_COUNT=$(detect_amd_gpu)
# Check for no_cuda argument or AMD GPU presence
if [[ ${1:-} == "no_cuda" ]] || [[ $AMD_GPU_COUNT -gt 0 ]]; then
	SKIP_CUDA=true
	log_info "CUDA support will be skipped"
fi

log_info "Installing environment in: $PROJECT_PATH"
log_info "CUDA support: $([ "$SKIP_CUDA" = true ] && echo "disabled" || echo "enabled")"

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

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

#===============================================================================
# PYTHON ENVIRONMENT SYNCHRONIZATION
#===============================================================================

log_info "Synchronizing Python environment with UV..."

# Sync environment based on flash-attention installation status
# Note: GIT_LFS_SKIP_SMUDGE prevents downloading large files during sync
sync_environment() {
	local flash_installed=false

	if uv pip list | grep -q "flash-attn\|flash_attn"; then
		flash_installed=true
	fi

	if [[ $flash_installed == true ]]; then
		log_info "Flash-attention detected, syncing without reinstalling..."
		if uv pip list | grep -q "flash-attn"; then
			GIT_LFS_SKIP_SMUDGE=true uv sync --inexact --no-install-package flash-attn
		else
			GIT_LFS_SKIP_SMUDGE=true uv sync --inexact --no-install-package flash_attn
		fi
	else
		log_info "No flash-attention detected, performing full sync..."
		GIT_LFS_SKIP_SMUDGE=true uv sync
	fi
}

sync_environment

#===============================================================================
# PRE-COMMIT HOOKS INSTALLATION
#===============================================================================

log_info "Installing pre-commit hooks..."
if ! uv run pre-commit install; then
	log_error "Failed to install pre-commit hooks"
	exit 1
fi

#===============================================================================
# CUDA DETECTION AND CONFIGURATION
#===============================================================================

# Function to detect CUDA installation
detect_cuda() {
	local cuda_found=false
	local cuda_path=""

	log_info "Searching for CUDA installation..."

	for path in "${CUDA_PATHS[@]}"; do
		if [[ -d "$path/bin" && -f "$path/bin/nvcc" ]]; then
			log_info "Found CUDA installation at: $path"
			cuda_path="$path"
			cuda_found=true
			break
		fi
	done

	if [[ $cuda_found == false ]]; then
		log_error "No CUDA installation found in standard paths"
		log_error "Searched paths: ${CUDA_PATHS[*]}"
		log_error "Please ensure CUDA is properly installed and accessible"
		return 1
	fi

	# Export CUDA paths
	export PATH="$cuda_path/bin${PATH:+:${PATH}}"
	export LD_LIBRARY_PATH="$cuda_path/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}"

	return 0
}

# Function to validate CUDA version
validate_cuda_version() {
	local nvcc_output
	local cuda_version
	local cuda_major
	local cuda_minor

	log_info "Validating CUDA version..."

	if ! nvcc_output=$(nvcc -V 2>/dev/null); then
		log_error "Failed to execute 'nvcc -V'"
		return 1
	fi

	# Extract CUDA version from nvcc output
	if ! cuda_version=$(echo "$nvcc_output" | grep -oP 'release \K[0-9]+\.[0-9]+' | head -1); then
		log_error "Could not determine CUDA version from nvcc output"
		return 1
	fi

	# Parse version components
	cuda_major=$(echo "$cuda_version" | cut -d. -f1)
	cuda_minor=$(echo "$cuda_version" | cut -d. -f2)

	# Validate version requirements
	if [[ $cuda_major -gt $MIN_CUDA_MAJOR ]] ||
		[[ $cuda_major -eq $MIN_CUDA_MAJOR && $cuda_minor -ge $MIN_CUDA_MINOR ]]; then
		log_info "CUDA $cuda_version detected: >= $MIN_CUDA_MAJOR.$MIN_CUDA_MINOR required ✓"
		return 0
	else
		log_error "CUDA $cuda_version detected, but CUDA >= $MIN_CUDA_MAJOR.$MIN_CUDA_MINOR is required"
		return 1
	fi
}

# Function to install flash-attention
install_flash_attention() {
	log_info "Checking flash-attention installation..."

	if uv pip list | grep -q "flash"; then
		log_info "Flash-attention is already installed ✓"
		return 0
	fi

	log_info "Installing flash-attention..."
	log_info "Note: This may take several minutes to compile..."

	# Reference links for known issues:
	# - Poetry bug: https://github.com/python-poetry/poetry/issues/8427
	# - UV variant: https://github.com/astral-sh/uv/issues/6437

	local flash_installer="$PROJECT_PATH/scripts/system_setup/flash_attn_installer.sh"

	if [[ ! -f $flash_installer ]]; then
		log_error "Flash-attention installer not found: $flash_installer"
		return 1
	fi

	if ! uv run bash "$flash_installer" -p "$PROJECT_PATH"; then
		log_error "Failed to install flash-attention"
		return 1
	fi

	log_info "Flash-attention installed successfully ✓"
}

# Install Flash Attention
install_flash_attention

# Main CUDA setup logic
if [[ $SKIP_CUDA == true ]]; then
	log_info "Skipping CUDA-related configurations as requested"
else
	if detect_cuda && validate_cuda_version; then
		log_info "CUDA environment setup completed successfully ✓"
	else
		log_error "CUDA setup failed"
		exit 1
	fi
fi

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

log_info "Configuring environment variables..."

# Set TRITON_CACHE_DIR to be user-specific for proper isolation
# This prevents permission issues in multi-user environments
TRITON_CACHE_DIR="/home/$(whoami)/.triton_cache"
export TRITON_CACHE_DIR
log_info "Set TRITON_CACHE_DIR to: $TRITON_CACHE_DIR"

# Configure Ray temporary directory
# Using /tmp ensures cleanup on system restart
RAY_TEMP_DIR="/tmp/ray"
export RAY_TEMP_DIR
log_info "Set RAY_TEMP_DIR to: $RAY_TEMP_DIR"

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

log_info "Environment installation completed successfully! ✅"
log_info ""
log_info "Summary:"
log_info "  Project path: $PROJECT_PATH"
log_info "  CUDA support: $([ "$SKIP_CUDA" = true ] && echo "disabled" || echo "enabled")"
log_info "  Python environment: synchronized"
log_info "  Pre-commit hooks: installed"
log_info "  Environment variables: configured"

if [[ $SKIP_CUDA == false ]]; then
	log_info "  Flash-attention: installed"
fi

log_info ""
log_info "You can now start using the mt_dao development environment!"
log_info "For more information, see the project documentation."
