dyinit

#!/usr/bin/env bash

# dyinit: Initialize the Daylily environment
# This script must be sourced, not executed directly.

# Ensure the script is sourced
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    echo "Error: This script must be sourced, not executed directly."
    echo "Usage: source $0 [--project <project_name>] [--skip-project-check]"
    return 3
fi


# Set default contact email if not already set
export DAY_CONTACT_EMAIL="${DAY_CONTACT_EMAIL:-john@daylilyinformatics.com}"

# Source AWS ParallelCluster configuration to get the region
if [[ -f "/etc/parallelcluster/cfnconfig" ]]; then
    source "/etc/parallelcluster/cfnconfig"
    region="$cfn_region"
else
    echo "Error: AWS ParallelCluster config not found at /etc/parallelcluster/cfnconfig"
fi

# Check if region is set
if [[ -z "$region" ]]; then
    echo "Error: cfn_region is not set or is empty."

fi

# Function to display usage
usage() {
    echo "This script initializes the Day CLI (commands aliased to dy-* )."
    echo "Region is autodetected as $region"
    echo "Usage: source $0 [--project <project_name>] [--skip-project-check]"
    echo "If --project is not specified, the default budget for the region is used: daylily-omics-analysis-${region}."
    echo "Valid projects/budgets and authorized users:"
    if [[ -f "/fsx/data/budget_tags/pcluster-project-budget-tags.tsv" ]]; then
        cat "/fsx/data/budget_tags/pcluster-project-budget-tags.tsv"
    else
        echo "Budget tags file not found."
    fi
    echo ""
    echo "Valid regions: Use a valid AWS region (e.g., us-east-1, us-west-2)."
    return 3
}

# Email validation regex reused from cluster creation script
EMAIL_REGEX='^[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}$'

# Function to check if a command exists
command_exists() {
    command -v "$1" &> /dev/null
}

is_valid_email() {
    local email="$1"
    [[ -n "$email" ]] || return 1
    [[ $email =~ $EMAIL_REGEX ]]
}

prompt_for_valid_email() {
    local __resultvar="$1"
    local __prompt="$2"
    local __input=""

    while true; do
        read -r -p "$__prompt" __input
        if [[ -z "$__input" ]]; then
            echo "Email cannot be empty. Please try again."
            continue
        fi

        if is_valid_email "$__input"; then
            printf -v "$__resultvar" '%s' "$__input"
            break
        else
            echo "Invalid email format. Please use string@string.string format."
        fi
    done
}

prompt_with_default() {
    local __resultvar="$1"
    local __prompt="$2"
    local __default="$3"
    local __input=""

    if [[ -n "$__default" ]]; then
        read -r -p "$__prompt [${__default}]: " __input
        __input="${__input:-$__default}"
    else
        read -r -p "$__prompt" __input
    fi

    printf -v "$__resultvar" '%s' "$__input"
}

# Ensure AWS CLI is installed
if ! command_exists aws; then
    echo "Error: AWS CLI is not installed. Please install it first."
    return 3
fi

# Parse input arguments
SKIP_PROJECT_CHECK=false
DEACTIVATE=false
while [[ "$#" -gt 0 ]]; do
    case "$1" in
        --project)
            PROJECT="$2"
            shift 2
            ;;
        --skip-project-check)
            SKIP_PROJECT_CHECK=true
            shift
            ;;
        --deactivate)
            DEACTIVATE=true
            shift
            ;;          
        -h|--help|help)
            usage
            return 0
            ;;
        *)
            echo "Unknown parameter passed: $1"
            usage
            return 3
            ;;
    esac
done


# Check if DEACTIVATE flag is set
if [[ "$DEACTIVATE" == true ]]; then
    echo "Deactivating existing environments..."
    source bin/day_deactivate &> /dev/null || {
        echo "Error during environment deactivation."
        return 4
    }
    echo "Environment deactivated successfully."
fi


# Set default project name if not provided
if [[ -z "$PROJECT" ]]; then
    if [[ -f "/opt/parallelcluster/shared/cluster-config.yaml" ]]; then
        PROJECT=$(awk '/Key: aws-parallelcluster-project/ {getline; print $2}' /opt/parallelcluster/shared/cluster-config.yaml)
        echo "Notice: --project not set. Using default project name: $PROJECT"
    else
        echo "Error: Cluster config file not found at /opt/parallelcluster/shared/cluster-config.yaml"

    fi
fi

# Display the parsed values
echo "Project: $PROJECT"
echo "Skip Project Check: $SKIP_PROJECT_CHECK"

# Export environment variables
export DAY_PROJECT="$PROJECT"
export DAY_AWS_REGION="$region"


# Get AWS Account ID
export AWS_ACCOUNT_ID=$(aws sts get-caller-identity --query "Account" --output text)
if [[ -z "$AWS_ACCOUNT_ID" ]]; then
    echo "Error: Unable to retrieve AWS Account ID."
    echo " is AWS_PROFILE=<yorprofile> set? or ... have you run 'aws configure --profile $AWS_PROFILE' to set up your AWS credentials and/or are they valid?"
    return 3
fi

# Skip project check if the flag is provided
if [[ "$SKIP_PROJECT_CHECK" == false ]]; then
    # Path to the budget file
    budget_file="/fsx/data/budget_tags/pcluster-project-budget-tags.tsv"

    # Ensure the budget file exists
    if [[ ! -f "$budget_file" ]]; then
        echo "Error: Budget file '$budget_file' not found."
        return 3
    fi

    # Validate the project for the current user
    USER_PROJECTS=$(awk -F'\t' -v user="$USER" '$2 ~ user {print $1}' "$budget_file" | tr '\n' ',' | sed 's/,$//')
    if [[ ! ",$USER_PROJECTS," =~ ",$PROJECT," ]]; then
        echo "Error: Project '$PROJECT' is not valid for user '$USER'."
        echo "Valid projects for '$USER': $USER_PROJECTS"
        PROJECT="daylily-global"
        echo ""
        echo "DEFAULTING TO PROJECT=daylily-global"
        echo ""
    fi
else
    echo "Skipping project validation as --skip-project-check was passed."
fi

derive_region_az_and_cluster_name() {
    local project_name="$1"
    local __region_az_var="$2"
    local __cluster_name_var="$3"

    local region_az_guess=""
    local cluster_name_guess=""

    if [[ "$project_name" =~ ^da-([^-]+-[^-]+-[^-]+)-(.*)$ ]]; then
        region_az_guess="${BASH_REMATCH[1]}"
        cluster_name_guess="${BASH_REMATCH[2]}"
    fi

    if [[ -n "$__region_az_var" ]]; then
        printf -v "$__region_az_var" '%s' "$region_az_guess"
    fi
    if [[ -n "$__cluster_name_var" ]]; then
        printf -v "$__cluster_name_var" '%s' "$cluster_name_guess"
    fi
}

derive_bucket_url() {
    local __resultvar="$1"
    local cluster_config="/opt/parallelcluster/shared/cluster-config.yaml"
    local bucket_url_guess=""

    if [[ -f "$cluster_config" ]]; then
        bucket_url_guess=$(awk '/ImportPath:/ {print $2; exit}' "$cluster_config")
        if [[ -n "$bucket_url_guess" ]]; then
            bucket_url_guess="${bucket_url_guess%/data/}"
            bucket_url_guess="${bucket_url_guess%/}"
        fi
    fi

    if [[ -z "$bucket_url_guess" ]]; then
        read -r -p "Enter S3 bucket URL (e.g., s3://example-bucket): " bucket_url_guess
    fi

    printf -v "$__resultvar" '%s' "$bucket_url_guess"
}

ensure_global_budget() {
    local region="$1"
    local region_az="$2"
    local cluster_name="$3"
    local bucket_url="$4"
    local budget_email="$5"
    local budgets_json="$6"

    local global_budget_name="daylily-global"
    local global_budget_exists
    if [[ -z "$budgets_json" ]]; then
        budgets_json='{"Budgets":[]}'
    fi
    global_budget_exists=$(echo "$budgets_json" | jq -r ".Budgets[] | select(.BudgetName==\"$global_budget_name\") | .BudgetName")

    if [[ -n "$global_budget_exists" && "$global_budget_exists" != "None" ]]; then
        echo "✅ Default budget '$global_budget_name' exists. Proceeding"
        return 0
    fi

    echo "Budget '$global_budget_name' not found."

    local gallowed_users=""
    prompt_with_default gallowed_users "Enter csv of allowed user names for global budget (ubuntu is allowed for all budgets)" ""
    gallowed_users="${gallowed_users//[[:space:]]/}"

    local default_global_budget_amount="200"
    local gamount=""
    prompt_with_default gamount "Enter global budget amount" "$default_global_budget_amount"

    echo "Creating the budget..."
    echo "bin/create_budget.sh -p $global_budget_name -r $region -t 25,50,75,99 -e $budget_email -a $gamount -c $cluster_name -z $region_az -b $bucket_url -u $gallowed_users"
    if ! bin/create_budget.sh -p "$global_budget_name" -r "$region" -t 25,50,75,99 -e "$budget_email" -a "$gamount" -c "$cluster_name" -z "$region_az" -b "$bucket_url" -u "$gallowed_users"; then
        echo "❌ Error: Failed to create the global budget."
        return 1
    fi

    echo "✅  Budget '$global_budget_name' created successfully."
}

create_project_budget() {
    local project="$1"
    local region="$2"
    local region_az="$3"
    local cluster_name="$4"
    local bucket_url="$5"
    local budget_email="$6"

    local allowed_users=""
    prompt_with_default allowed_users "Enter csv string of allowed user names for cluster budget (ubuntu - default)" ""
    allowed_users="${allowed_users//[[:space:]]/}"

    local default_budget_amount="200"
    local amount=""
    prompt_with_default amount "Enter cluster budget amount" "$default_budget_amount"

    echo "Creating the budget..."
    echo "bin/create_budget.sh -p $project -r $region -t 75 -e $budget_email -a $amount -c $cluster_name -z $region_az -b $bucket_url -u $allowed_users"
    if ! bin/create_budget.sh -p "$project" -r "$region" -t 75 -e "$budget_email" -a "$amount" -c "$cluster_name" -z "$region_az" -b "$bucket_url" -u "$allowed_users"; then
        echo "❌ Error: Failed to create the budget."
        return 1
    fi

    echo "✅ Budget '$project' created successfully."
}

# Query AWS budgets
BUDGETS=$(aws budgets describe-budgets --account-id "$AWS_ACCOUNT_ID" --region "$region" 2>/dev/null)

if [[ -z "$BUDGETS" && "$SKIP_PROJECT_CHECK" != "true" ]]; then
    echo "Error: Unable to retrieve any budgets from AWS. Please check your AWS permissions or configuration."

fi

# Check if the specified project budget exists
MATCHING_BUDGET=$(echo "$BUDGETS" | jq -r ".Budgets[] | select(.BudgetName==\"$PROJECT\")")

if [[ -z "$MATCHING_BUDGET" && "$SKIP_PROJECT_CHECK" != "true" ]]; then
    echo "No matching AWS budget found for project '$PROJECT' in region '$region'."
    echo "Available AWS Budgets (specify with the --project flag):"
    echo "$BUDGETS" | jq -r '.Budgets[].BudgetName'

    read -p "Would you like to create a new budget? (y/n): " RESPONSE
    if [[ "$RESPONSE" =~ ^[Yy]$ ]]; then
        echo "Starting budget creation flow..."

        local_region_az=""
        local_cluster_name=""
        derive_region_az_and_cluster_name "$PROJECT" local_region_az local_cluster_name

        if [[ -z "$local_region_az" ]]; then
            prompt_with_default local_region_az "Enter the region+AZ for the budget (e.g., us-west-2d)" ""
        fi

        if [[ -z "$local_cluster_name" ]]; then
            prompt_with_default local_cluster_name "Enter the cluster name for budget tagging" ""
        fi

        if [[ -z "$local_region_az" || -z "$local_cluster_name" ]]; then
            echo "❌ Error: Unable to determine region AZ or cluster name for budget creation."
            return 3
        fi

        bucket_url=""
        derive_bucket_url bucket_url
        if [[ -z "$bucket_url" ]]; then
            echo "❌ Error: Unable to determine bucket URL for budget creation."
            return 3
        fi

        budget_email="${DAY_CONTACT_EMAIL:-}"
        if [[ -z "$budget_email" ]]; then
            prompt_for_valid_email budget_email "Enter an email for budget alerts: "
        elif ! is_valid_email "$budget_email"; then
            echo "Contact email '$budget_email' is not a valid email format."
            prompt_for_valid_email budget_email "Enter an email for budget alerts: "
        else
            echo "Using budget email: $budget_email"
        fi

        if ! ensure_global_budget "$region" "$local_region_az" "$local_cluster_name" "$bucket_url" "$budget_email" "$BUDGETS"; then
            return 3
        fi

        if ! create_project_budget "$PROJECT" "$region" "$local_region_az" "$local_cluster_name" "$bucket_url" "$budget_email"; then
            return 3
        fi

        BUDGETS=$(aws budgets describe-budgets --account-id "$AWS_ACCOUNT_ID" --region "$region" 2>/dev/null)
        MATCHING_BUDGET=$(echo "$BUDGETS" | jq -r ".Budgets[] | select(.BudgetName==\"$PROJECT\")")
        if [[ -z "$MATCHING_BUDGET" ]]; then
            echo "❌ Error: Budget creation reported success but the budget is still not visible via AWS CLI."
            return 3
        fi
    else
        echo "Exiting without creating a budget."
        return 3
    fi
fi

# Extract budget details
export TOTAL_BUDGET=$(echo "$MATCHING_BUDGET" | jq -r ".BudgetLimit.Amount")
export USED_BUDGET=$(echo "$MATCHING_BUDGET" | jq -r ".CalculatedSpend.ActualSpend.Amount")

if [[ (-z "$TOTAL_BUDGET" || -z "$USED_BUDGET") && "$SKIP_PROJECT_CHECK" != "true" ]]; then
    echo "Error: Unable to calculate budget details for project '$PROJECT'."
    return 3
fi

# Calculate usage percentage
export PERCENT_USED=$(awk "BEGIN {print ($USED_BUDGET / $TOTAL_BUDGET) * 100}")

# Display budget information
echo ""
echo "________________________________________________________"
echo "AWS Budget for project '$PROJECT' in region '$region':"
echo "  Total: $TOTAL_BUDGET USD"
echo "  Used: $USED_BUDGET USD"
echo "  Percent Used: $PERCENT_USED%"
echo "________________________________________________________"

# Check if the budget is exhausted
if (( $(echo "$PERCENT_USED >= 100" | bc -l) )); then
    echo "Warning: Budget for project '$PROJECT' is exhausted!"
fi

# Set environment variables
export APPTAINER_HOME="/fsx/resources/environments/containers/$USER/$(hostname)/"
export DAY_BIOME="AWSPC"
export DAY_ROOT="$PWD"

# Increase file descriptor limit
ulimit -n 16384

# Set Sentieon variables

# Backup original PATH and PS1 if not already set
export ORIG_PATH="${ORIG_PATH:-$PATH}"
export ORIG_PS1="${ORIG_PS1:-$PS1}"

# Source color themes
if [[ -f "config/cli/colr_themes.sh" ]]; then
    source "config/cli/colr_themes.sh"
else
    echo "Warning: Color themes file not found at config/cli/colr_themes.sh"
fi

# Warn if shell is not bash
if [[ "$SHELL" != *"bash"* ]]; then
    echo "Warning: This script has only been tested with bash."
    sleep 10
fi

# Update PATH
export PATH="$PATH:$PWD/bin"

# Define aliases
alias dy-h="echo hello"
alias dy-b="bin/init_dayec"
alias day-build-env="bin/init_dayec"

# Display help if requested
if [[ "$1" =~ ^(-h|--help|help)$ ]]; then
    if [[ -f "docs/markdown/cli_help_brief.md" ]]; then
        hlp=$(cat docs/markdown/cli_help_brief.md)
        if command_exists colr; then
            colr """$hlp""" "floralwhite" "midnightblue" "b"
        else
            echo "$hlp"
        fi
    else
        echo "Help file not found at docs/markdown/cli_help_brief.md"
    fi
    return 0
fi

# Define 'sq' alias for squeue with specific output format
sq_cmd="squeue -o '%i  %P  %C  %t  %N  %c  %T  %m  %M  %D  %j'"
sqq() {
    if [ $# -eq 0 ]; then
        eval "$sq_cmd"
    else
        eval "$sq_cmd -j \"$(echo "$@" | tr ' ' ',')\""
    fi
}
alias sq=sqq

echo "squeue helper alias 'sq' activated"


reference_bucket=$(awk '/^[[:space:]]*Script: s3:/{match($0,/s3:\/\/([^/]+)/,a); print a[1]; exit}' /opt/parallelcluster/shared/cluster-config.yaml)


# Activate or create the Daylily CLI conda environment
if conda env list | grep -q "^DAY-EC "; then
    echo "Conda environment 'DAY-EC' already exists!"
    ENV_ALREADY_EXISTS=true
else
    echo "Creating 'DAY-EC' environment."
    cd ~/projects/daylily-ephemeral-cluster || {
        echo "Error: Unable to change directory to ~/projects/daylily-ephemeral-cluster"
        return 3
    }
    bin/init_dayec
    if [ $? -ne 0 ]; then
        echo "Error: Failed to create 'DAY-EC' environment."
        return 
    fi
fi
conda activate DAY-EC || {
    echo "Error: Failed to activate 'DAY-EC' environment."
}

# Final message
echo "Day CLI initialized for project '$PROJECT' in region '$region'."
echo -e "The Daylily CLI is now available."
echo -e "Available commands :"
echo -e "\t                           "
echo -e "\t                           "
echo -e "\t  day-build-env : initialize the DAY-EC conda environment (same as dy-b), MUST be run from ~/projects/daylily-ephemeral-cluster/"
echo -e "\t "
echo -e "\t  day-clone --help : clone an analysis repo"
echo -e "\t  day-clone --list : list available analysis repos"
echo -e "\t "
echo -e "\t  AWS_PROFILE=<<your_aws_profile_used_to_create_cluster>> daylily-stage-analysis-samples-headnode <analysis_samples.tsv> [stage_target] : generate /fsx staging directories and config/{samples.tsv,units.tsv} from an analysis_samples TSV !! NOTE, you will have to run aws configure on the headnode to run this script, this is the only command this is needed for. Use the same profile as the one used to create the cluster."
echo -e "\t                           "
echo -e "\t  daylily-export-fsx-to-s3 : export FSx data to S3, if this fails, use the AWS Fsx Filesystem console! *and* this might only be runnable from your local workstation not the headnode if aws configure was not run."
echo -e "\t                           "
echo -e '\t  sq : an alias to `squeue` with a useful default output format for monitoring running jobs. '
echo -e "\t "
echo -e "\t\t  $(sqq | head -n 8 ) "
echo -e "\t\t\t... first 7 jobs                           "
echo -e "\t  "
echo -e "\t The reference S3 bucket is set to: \$reference_bucket ($reference_bucket)"
echo -e "\t  "

return 0