#!/bin/bash
#
# mt_dao Base Launcher Script
# ===========================
#
# This script orchestrates the launch of a mt_dao federated learning environment.
# It handles environment setup, GPU allocation, and component orchestration including:
# - Superlink (coordination service)
# - ServerApp (federated learning server)
# - ClientApp(s) (federated learning clients)
# - Flower run execution
#
# Usage:
#   ./mt_dao_base.sh [-p|--project_path PATH]
#
# Environment Variables:
#   RUN_UUID                           - Unique identifier for this run
#   SAVE_PATH                         - Base save path for outputs
#   mt_dao_SAVE_PATH                  - mt_dao-specific save path
#   CLIENTAPP_LAUNCHER_INDEPENDENT    - Launch one ClientApp per GPU
#   CLIENTAPP_LAUNCHER_BI_INDEPENDENT - Launch one ClientApp per 2 GPUs
#   CLIENTAPP_LAUNCHER_QUADRI_INDEPENDENT - Launch one ClientApp per 4 GPUs
#   mt_dao_LAUNCHER_NO_SUPERLINK      - Skip Superlink launch
#   mt_dao_LAUNCHER_NO_SERVERAPP      - Skip ServerApp launch
#   SUPERLINK_ADDRESS                 - Superlink service address
#   EXEC_API_ADDRESS                  - Execution API address
#   SERVERAPP_API_ADDRESS             - ServerApp API address
#   EXTERNAL_CONFIGS                  - Additional configuration parameters
#
# shellcheck disable=SC2090,SC2086,SC2089,SC1091

set -eo pipefail # Enable strict error handling

# =============================================================================
# CONFIGURATION AND ARGUMENT PARSING
# =============================================================================

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

# Parse command-line options
if ! OPTIONS=$(getopt -o p: --long project_path: -n 'parse-options' -- "$@"); then
	echo "mt_dao_base.sh: Error parsing options" >&2
	exit 1
fi

eval set -- "$OPTIONS"

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

echo "mt_dao_base.sh: PROJECT_PATH=$PROJECT_PATH"

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

#! Moving to the project folder
cd "$PROJECT_PATH" || {
	echo "mt_dao_base.sh: ERROR - Cannot access project directory: $PROJECT_PATH" >&2
	exit 1
}

#! Preparing environment
echo "mt_dao_base.sh: Setting up environment..."
if [[ $(hostname) == *'gpu-q'* ]]; then
	echo "mt_dao_base.sh: Detected CSD3 environment"
	#! Executing the environment preparation script
	#! NOTE: Must use "." to execute, "sh" doesn't work
	bash "$PROJECT_PATH"/scripts/install_hpc_env.sh -p "$PROJECT_PATH"
else
	echo "mt_dao_base.sh: Detected local/non-CSD3 environment"
	#! Executing the environment preparation script
	#! NOTE: Must use "." to execute, "sh" doesn't work
	bash "$PROJECT_PATH"/scripts/install_env.sh -p "$PROJECT_PATH"
fi

#! Export the endpoint of the S3 object store
export S3_ENDPOINT_URL='http://anon.anon:9000'

# =============================================================================
# PATH AND DIRECTORY SETUP
# =============================================================================

#! Set the date and time if not already set
if [ -z "${DATETIME:-}" ]; then
	DATETIME=$(date '+%Y%m%d_%H%M%S')
	export DATETIME
	echo "mt_dao_base.sh: Generated DATETIME: $DATETIME"
else
	echo "mt_dao_base.sh: Using existing DATETIME: $DATETIME"
fi

#! If RUN_UUID hasn't been set, set it to the default value
if [ -z "${RUN_UUID:-}" ]; then
	export RUN_UUID="mt_dao-$DATETIME"
	echo "mt_dao_base.sh: Generated RUN_UUID: $RUN_UUID"
else
	echo "mt_dao_base.sh: Using existing RUN_UUID: $RUN_UUID"
fi

#! If SAVE_PATH hasn't been set, set it to the default value
if [ -z "${SAVE_PATH:-}" ]; then
	export SAVE_PATH="$PROJECT_PATH/$RUN_UUID"
	mkdir -p "$SAVE_PATH"
	echo "mt_dao_base.sh: Created SAVE_PATH: $SAVE_PATH"
else
	echo "mt_dao_base.sh: Using existing SAVE_PATH: $SAVE_PATH"
fi

#! If mt_dao_SAVE_PATH hasn't been set, set it to the default value
if [ -z "${mt_dao_SAVE_PATH:-}" ]; then
	export mt_dao_SAVE_PATH="$PROJECT_PATH/runs/$RUN_UUID"
	echo "mt_dao_base.sh: Set mt_dao_SAVE_PATH: $mt_dao_SAVE_PATH"
else
	echo "mt_dao_base.sh: Using existing mt_dao_SAVE_PATH: $mt_dao_SAVE_PATH"
fi

# Create necessary directories
mkdir -p "$mt_dao_SAVE_PATH"
mkdir -p "$mt_dao_SAVE_PATH/$DATETIME"

# =============================================================================
# GPU DETECTION AND CONFIGURATION
# =============================================================================

#! Getting visible GPUs
echo "mt_dao_base.sh: Detecting available GPUs..."
N_GPUS=$(uv run python -c 'import torch; print(torch.cuda.device_count())')

if [ "$N_GPUS" -eq 0 ]; then
	echo "mt_dao_base.sh: WARNING - No GPUs found. Running in CPU mode."
	CUDA_VISIBLE_DEVICES=""
else
	CUDA_VISIBLE_DEVICES=$(seq -s, 0 $((N_GPUS - 1)))
	echo "mt_dao_base.sh: Detected $N_GPUS GPU(s)"
fi
echo "mt_dao_base.sh: CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES"

# =============================================================================
# CONFIGURATION RESOLUTION
# =============================================================================

#! Run Hydra resolver
echo "mt_dao_base.sh: Running Hydra configuration resolver..."
echo "mt_dao_base.sh: EXTERNAL_CONFIGS=$EXTERNAL_CONFIGS"
HYDRA_FULL_ERROR=1 EXTERNAL_CONFIGS=$EXTERNAL_CONFIGS . "$PROJECT_PATH"/scripts/base_launchers/hydra_resolver.sh

# =============================================================================
# FLOWER SERVICE CONFIGURATION
# =============================================================================

#! Set the Flower driver and fleet API addresses if they haven't been set
export SUPERLINK_ADDRESS=${SUPERLINK_ADDRESS:-"[::]:9092"}
export EXEC_API_ADDRESS=${EXEC_API_ADDRESS:-"[::]:9093"}
export SERVERAPP_API_ADDRESS=${SERVERAPP_API_ADDRESS:-"[::]:9091"}

echo "mt_dao_base.sh: Flower service addresses configured:"
echo "  SUPERLINK_ADDRESS=$SUPERLINK_ADDRESS"
echo "  EXEC_API_ADDRESS=$EXEC_API_ADDRESS"
echo "  SERVERAPP_API_ADDRESS=$SERVERAPP_API_ADDRESS"

# =============================================================================
# SERVICE ORCHESTRATION
# =============================================================================

#! Start a Superlink
echo "mt_dao_base.sh: Launching Superlink service..."
if [[ -n ${mt_dao_LAUNCHER_NO_SUPERLINK:-} ]]; then
	echo "mt_dao_base.sh: mt_dao_LAUNCHER_NO_SUPERLINK is set, so skipping Superlink launch."
else
	. "$PROJECT_PATH"/scripts/base_launchers/superlink.sh -p "$PROJECT_PATH"
	echo "mt_dao_base.sh: Superlink launched successfully"
fi

#! Launch ServerApp
echo "mt_dao_base.sh: Launching ServerApp..."
if [[ -n ${mt_dao_LAUNCHER_NO_SERVERAPP:-} ]]; then
	echo "mt_dao_base.sh: mt_dao_LAUNCHER_NO_SERVERAPP is set, so skipping ServerApp launch."
else
	. "$PROJECT_PATH"/scripts/base_launchers/serverapp.sh -p "$PROJECT_PATH"
	echo "mt_dao_base.sh: ServerApp launched successfully"
fi

# =============================================================================
# CLIENTAPP LAUNCH ORCHESTRATION
# =============================================================================

# Function to launch a single ClientApp instance
launch_clientapp() {
	local cuda_devices="$1"
	local description="$2"

	echo "mt_dao_base.sh: Launching ClientApp with $description (CUDA_VISIBLE_DEVICES=$cuda_devices)"
	CUDA_VISIBLE_DEVICES="$cuda_devices" . "$PROJECT_PATH"/scripts/base_launchers/clientapp.sh -p "$PROJECT_PATH"
	echo "mt_dao_base.sh: ClientApp launched successfully with CUDA_VISIBLE_DEVICES=$cuda_devices"
}

#! Launch ClientApp(s) based on the configured launch mode
echo "mt_dao_base.sh: Configuring ClientApp launch strategy..."

if [[ -n ${CLIENTAPP_LAUNCHER_INDEPENDENT:-} ]]; then
	# Independent mode: 1 ClientApp per GPU
	echo "mt_dao_base.sh: CLIENTAPP_LAUNCHER_INDEPENDENT is set - launching 1 ClientApp per GPU"

	#! Make the CUDA_VISIBLE_DEVICES into a list of int
	LIST_CUDA_VISIBLE_DEVICES=$(echo "$CUDA_VISIBLE_DEVICES" | tr "," "\n")

	#! Loop over the CUDA_VISIBLE_DEVICES
	for CUDA_VISIBLE_DEVICE in $LIST_CUDA_VISIBLE_DEVICES; do
		#! Launch one instance of ClientApp for each CUDA device
		launch_clientapp "$CUDA_VISIBLE_DEVICE" "single GPU $CUDA_VISIBLE_DEVICE"
	done

elif [[ -n ${CLIENTAPP_LAUNCHER_BI_INDEPENDENT:-} ]]; then
	# Bi-independent mode: 1 ClientApp per 2 GPUs
	echo "mt_dao_base.sh: CLIENTAPP_LAUNCHER_BI_INDEPENDENT is set - launching 1 ClientApp per 2 GPUs"

	#! Get the number of CUDA devices
	IFS=',' read -ra CUDA_DEVICES <<<"$CUDA_VISIBLE_DEVICES"
	NUM_DEVICES=${#CUDA_DEVICES[@]}
	echo "mt_dao_base.sh: Total CUDA devices: $NUM_DEVICES"

	#! Launch one instance of ClientApp for each pair of CUDA devices
	for ((i = 0; i < NUM_DEVICES; i += 2)); do
		if [ $((i + 1)) -lt $NUM_DEVICES ]; then
			# Launch with device pair
			DEVICE_PAIR="${CUDA_DEVICES[$i]},${CUDA_DEVICES[i + 1]}"
			launch_clientapp "$DEVICE_PAIR" "GPU pair ${CUDA_DEVICES[$i]},${CUDA_DEVICES[i + 1]}"
		else
			# Launch with single remaining device
			launch_clientapp "${CUDA_DEVICES[$i]}" "remaining single GPU ${CUDA_DEVICES[$i]}"
		fi
	done

elif [[ -n ${CLIENTAPP_LAUNCHER_QUADRI_INDEPENDENT:-} ]]; then
	# Quadri-independent mode: 1 ClientApp per 4 GPUs
	echo "mt_dao_base.sh: CLIENTAPP_LAUNCHER_QUADRI_INDEPENDENT is set - launching 1 ClientApp per 4 GPUs"

	#! Get the number of CUDA devices
	IFS=',' read -ra CUDA_DEVICES <<<"$CUDA_VISIBLE_DEVICES"
	NUM_DEVICES=${#CUDA_DEVICES[@]}
	echo "mt_dao_base.sh: Total CUDA devices: $NUM_DEVICES"

	#! Launch one instance of ClientApp for each quad of CUDA devices
	for ((i = 0; i < NUM_DEVICES; i += 4)); do
		if [ $((i + 3)) -lt $NUM_DEVICES ]; then
			# Launch with device quad
			DEVICE_QUAD="${CUDA_DEVICES[$i]},${CUDA_DEVICES[i + 1]},${CUDA_DEVICES[i + 2]},${CUDA_DEVICES[i + 3]}"
			launch_clientapp "$DEVICE_QUAD" "GPU quad ${CUDA_DEVICES[$i]},${CUDA_DEVICES[i + 1]},${CUDA_DEVICES[i + 2]},${CUDA_DEVICES[i + 3]}"
		else
			# Launch with remaining single device
			launch_clientapp "${CUDA_DEVICES[$i]}" "remaining single GPU ${CUDA_DEVICES[$i]}"
		fi
	done

else
	# Standard mode: 1 ClientApp with all available GPUs
	echo "mt_dao_base.sh: Using standard launcher - launching single ClientApp with all GPUs"
	launch_clientapp "$CUDA_VISIBLE_DEVICES" "all available GPUs"
fi

# =============================================================================
# FLOWER EXECUTION
# =============================================================================

#! Launch Flower run
echo "mt_dao_base.sh: Starting Flower federated learning execution..."
echo "mt_dao_base.sh: Logs will be saved to: $mt_dao_SAVE_PATH/$DATETIME/flwr_run.log"

FLWR_HOME="$PROJECT_PATH" uv run flwr run . local-deployment \
	--stream \
	--federation-config "address='${EXEC_API_ADDRESS}'" 2>&1 | tee "$mt_dao_SAVE_PATH/$DATETIME/flwr_run.log"

echo "mt_dao_base.sh: Flower execution completed"
