dyoainit

#!/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]"
    exit 3
fi

region=""
PROJECT="${PROJECT:-}"
DAYOA_PLATFORM="$(uname -s 2>/dev/null || echo unknown)"
DAYOA_MAC_LOCAL=false
if [[ "$DAYOA_PLATFORM" == "Darwin" ]]; then
    DAYOA_MAC_LOCAL=true
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:-}"
elif [[ "$DAYOA_MAC_LOCAL" == true ]]; then
    region="local-mac"
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."
    return 3
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/references/runtime_assets/budget_tags/pcluster-project-budget-tags.tsv" ]]; then
        cat "/fsx/references/runtime_assets/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
}

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

cluster_config_tag_value() {
    local tag_key="$1"
    local config_file="${2:-/opt/parallelcluster/shared/cluster-config.yaml}"
    if [[ ! -f "$config_file" ]]; then
        return 1
    fi
    awk -v tag_key="$tag_key" '
        $1 == "-" && $2 == "Key:" && $3 == tag_key {
            getline
            if ($1 == "Value:") {
                print $2
                found = 1
                exit
            }
        }
        END { exit(found ? 0 : 1) }
    ' "$config_file"
}

day_current_user() {
    local current_user="${USER:-}"
    if [[ -z "$current_user" ]]; then
        current_user="$(id -un 2>/dev/null || true)"
    fi
    if [[ -z "$current_user" ]]; then
        return 1
    fi
    printf "%s\n" "$current_user"
}

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

# Ensure conda is available
if ! command_exists conda; then
    echo "Error: conda is not available in PATH. Please install or load conda before continuing."
    return 3
fi

# Parse input arguments
SKIP_PROJECT_CHECK=false
DEACTIVATE=false
while [[ "$#" -gt 0 ]]; do
    case "$1" in
        --project)
            if [[ -z "${2:-}" ]]; then
                echo "Error: --project requires a project name."
                return 3
            fi
            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

if [[ "$DAYOA_MAC_LOCAL" == true ]]; then
    if [[ "${CONDA_DEFAULT_ENV:-}" != "DAY-EC" && "${CONDA_DEFAULT_ENV:-}" != "DAYOA" ]]; then
        echo "Error: macOS local mode requires the DAY-EC or DAYOA conda environment."
        echo "Run: conda activate DAY-EC"
        return 3
    fi
    SKIP_PROJECT_CHECK=true
fi

# 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
cluster_config_file="/opt/parallelcluster/shared/cluster-config.yaml"
BUDGET_ENFORCEMENT=""
if [[ -f "$cluster_config_file" ]]; then
    BUDGET_ENFORCEMENT="$(cluster_config_tag_value "aws-parallelcluster-enforce-budget" "$cluster_config_file" || true)"
fi
if [[ -z "$PROJECT" ]]; then
    if [[ "$DAYOA_MAC_LOCAL" == true ]]; then
        PROJECT="dayoa-local-mac"
        echo "Notice: macOS local mode active. Using project name: $PROJECT"
    elif [[ -f "$cluster_config_file" ]]; then
        PROJECT="$(cluster_config_tag_value "aws-parallelcluster-project" "$cluster_config_file" || true)"
        if [[ -z "$PROJECT" ]]; then
            echo "Error: Project tag aws-parallelcluster-project was not found in $cluster_config_file"
            return 3
        fi
        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"
        return 3
    fi
fi

if [[ "$BUDGET_ENFORCEMENT" == "skip" && "$SKIP_PROJECT_CHECK" == false ]]; then
    echo "Notice: Cluster config sets aws-parallelcluster-enforce-budget=skip; skipping project validation."
    SKIP_PROJECT_CHECK=true
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"
export AWS_ACCOUNT_ID="NA"

# Skip project check if the flag is provided
if [[ "$SKIP_PROJECT_CHECK" == false ]]; then
    CURRENT_USER="$(day_current_user)"
    if [[ -z "$CURRENT_USER" ]]; then
        echo "Error: USER is not set and id -un failed; cannot validate project budget authorization."
        return 3
    fi

    # Path to the budget file
    budget_file="/fsx/references/runtime_assets/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="$CURRENT_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 '$CURRENT_USER'."
        echo "Valid projects for '$CURRENT_USER': $USER_PROJECTS"
        return 3
    fi
else
    echo "Skipping project validation as --skip-project-check was passed."
fi

# Function to create a new AWS budget
create_budget() {
    echo "To create a budget, please exit and run the following command:"
    echo "bash bin/create_budget.sh --help"
    echo "Then append the new project/budget name to /fsx/references/runtime_assets/budget_tags/pcluster-project-budget-tags.tsv."
    echo "Note: /fsx/references/runtime_assets is read-only; update the file via S3."

}

is_nonnegative_number() {
    [[ "$1" =~ ^[0-9]+([.][0-9]+)?$ ]]
}

export TOTAL_BUDGET="NA"
export USED_BUDGET="NA"
export PERCENT_USED="NA"

if [[ "$SKIP_PROJECT_CHECK" == true ]]; then
    echo ""
    echo "________________________________________________________"
    echo "AWS Budget lookup skipped for project '$PROJECT' in region '$region'."
    echo "  Total: $TOTAL_BUDGET"
    echo "  Used: $USED_BUDGET"
    echo "  Percent Used: $PERCENT_USED"
    echo "________________________________________________________"
else
    # 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."
        return 3
    fi

    if ! command_exists jq; then
        echo "Error: jq is required to parse AWS budget JSON."
        return 3
    fi

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

    if [[ "$budgets_ret" -ne 0 || -z "$BUDGETS" ]]; then
        echo "Error: Unable to retrieve AWS budgets. Please check your AWS permissions or configuration."
        return 3
    fi

    if ! echo "$BUDGETS" | jq -e '.Budgets | type == "array"' >/dev/null 2>&1; then
        echo "Error: AWS budget response was missing a Budgets array."
        return 3
    fi

    # Check if the specified project budget exists
    MATCHING_BUDGET=$(echo "$BUDGETS" | jq -c --arg project "$PROJECT" '[.Budgets[] | select(.BudgetName==$project)] | if length == 1 then .[0] else empty end')

    if [[ -z "$MATCHING_BUDGET" ]]; 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
            create_budget
        else
            echo "Exiting without creating a budget."
        fi
        return 3
    fi

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

    if ! is_nonnegative_number "$TOTAL_BUDGET" || ! is_nonnegative_number "$USED_BUDGET"; then
        echo "Error: AWS budget response did not include numeric BudgetLimit.Amount and CalculatedSpend.ActualSpend.Amount values for project '$PROJECT'."
        return 3
    fi

    if ! awk -v total="$TOTAL_BUDGET" 'BEGIN { exit (total > 0 ? 0 : 1) }'; then
        echo "Error: AWS budget total for project '$PROJECT' must be greater than zero."
        return 3
    fi

    # Calculate usage percentage
    PERCENT_USED=$(awk -v used="$USED_BUDGET" -v total="$TOTAL_BUDGET" 'BEGIN { printf "%.2f", (used / total) * 100 }')
    percent_ret=$?
    if [[ "$percent_ret" -ne 0 || -z "$PERCENT_USED" ]]; then
        echo "Error: Unable to calculate budget usage percentage for project '$PROJECT'."
        return 3
    fi

    export TOTAL_BUDGET
    export USED_BUDGET
    export PERCENT_USED

    # 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 awk -v pct="$PERCENT_USED" 'BEGIN { exit (pct >= 100 ? 0 : 1) }'; then
        echo "Warning: Budget for project '$PROJECT' is exhausted!"
    fi
fi

# Set environment variables
DAY_USER="$(day_current_user)"
if [[ -z "$DAY_USER" ]]; then
    echo "Error: USER is not set and id -un failed; cannot initialize APPTAINER_HOME."
    return 3
fi
export USER="${USER:-$DAY_USER}"
export DAY_ROOT="$PWD"
if [[ "$DAYOA_MAC_LOCAL" == true ]]; then
    export DAY_BIOME="MAC"
    export DAYOA_MAC_STATE_DIR="${DAYOA_MAC_STATE_DIR:-$DAY_ROOT/.dayoa-local-mac}"
    export APPTAINER_HOME="$DAYOA_MAC_STATE_DIR/apptainer_home/$DAY_USER/$(hostname)/"
else
    export DAY_BIOME="AWSPC"
    export APPTAINER_HOME="/fsx/resources/environments/containers/$DAY_USER/$(hostname)/"
fi

# Increase file descriptor limit
ulimit -n 16384

# Set Sentieon variables
if [[ "$DAYOA_MAC_LOCAL" == true ]]; then
    export SENTIEON_TMPDIR="$DAYOA_MAC_STATE_DIR/sentieon_tmp"
    export SENTIEON_INSTALL_DIR=""
    export SENTIEON_LICENSE="${SENTIEON_LICENSE:-}"
else
    export SENTIEON_TMPDIR='/dev/shm/'
    export SENTIEON_INSTALL_DIR='/fsx/references/runtime_assets/cached_envs/sentieon-genomics-202503.02'
    export PATH="$PATH:$SENTIEON_INSTALL_DIR/bin/"
    sentieon_config_file="${DAYLILY_CLI_GLOBAL_CONFIG:-$HOME/.config/daylily/daylily_cli_global.yaml}"

    if [[ -n "${SENTIEON_LICENSE:-}" && -f "$SENTIEON_LICENSE" ]]; then
        echo "Using explicit SENTIEON_LICENSE: $SENTIEON_LICENSE"
    else
        if [[ -n "${SENTIEON_LICENSE:-}" ]]; then
            echo "ERROR: SENTIEON_LICENSE is set but the file does not exist: $SENTIEON_LICENSE"
            return 3
        fi

        if [[ ! -f "$sentieon_config_file" ]]; then
            echo "ERROR: Daylily config not found: $sentieon_config_file"
            echo "Set SENTIEON_LICENSE to an explicit license file or configure daylily.sentieon_lic_path."
            return 3
        fi

        sentieon_configured_license="$(awk '
            $1 == "sentieon_lic_path:" {
                $1 = ""
                sub(/^[[:space:]]+/, "")
                gsub(/^["'\'']|["'\'']$/, "")
                print
                found = 1
                exit
            }
            END { exit(found ? 0 : 1) }
        ' "$sentieon_config_file" 2>/dev/null || true)"
        if [[ -z "$sentieon_configured_license" || "$sentieon_configured_license" == "null" ]]; then
            echo "ERROR: daylily.sentieon_lic_path is missing from $sentieon_config_file"
            echo "Set SENTIEON_LICENSE to an explicit license file or configure daylily.sentieon_lic_path."
            return 3
        fi
        if [[ ! -f "$sentieon_configured_license" ]]; then
            echo "ERROR: daylily.sentieon_lic_path does not exist: $sentieon_configured_license"
            return 3
        fi
        export SENTIEON_LICENSE="$sentieon_configured_license"
        echo "Using configured Sentieon license file: $SENTIEON_LICENSE"
    fi
fi

# 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
if ! command_exists colr; then
    colr() {
        if [[ "${1:-}" == "-n" ]]; then
            shift
            printf "%s" "${1:-}"
        else
            printf "%s\n" "${1:-}"
        fi
    }
fi

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

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

# Source tab completion script
if [[ -f "bin/tabcomp.bash" ]]; then
    if command_exists complete; then
        source "bin/tabcomp.bash"
    elif [[ "$DAYOA_MAC_LOCAL" == true ]]; then
        echo "Notice: Bash tab completion is not available in this shell; skipping bin/tabcomp.bash."
    else
        echo "Warning: Bash tab completion is not available in this shell; skipping bin/tabcomp.bash."
    fi
else
    echo "Warning: Tab completion script not found at bin/tabcomp.bash"
fi

# Define aliases
alias day-activate="source bin/day_activate"
alias day-deactivate="source bin/day_deactivate"
alias day-run="bin/day_run"
alias day-monitor="bin/day_monitor"
alias dy-a="source bin/day_activate"
alias dy-d="source bin/day_deactivate"
alias dy-r="bin/day_run"
alias dy-m="bin/day_monitor"
alias day-help="bin/day_run help"
alias dy-h="bin/day_run help"
alias day-set-genome-build="source bin/day_set_genome_build"
alias dy-g="source bin/day_set_genome_build"


# 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

if [[ "$DAYOA_MAC_LOCAL" == true ]]; then
    if ! conda env list | awk '{print $1}' | grep -Fx "DAYOA" >/dev/null 2>&1; then
        echo "Error: macOS local mode requires the DAYOA conda environment."
        echo "Install it with: source config/day/day_env_installer.sh DAYOA"
        return 1
    fi
    conda activate DAYOA || {
        echo "Error: Failed to activate 'DAYOA' conda environment."
        return 1
    }
    echo "macOS local mode: using conda environment DAYOA."
else
    conda env list
    # Ensure the DAYOA conda environment exists; build it if missing
    if ! conda env list | awk '{print $1}' | grep -Fx "DAYOA" >/dev/null 2>&1; then
        echo "!Conda environment 'DAYOA' not found. Building it now..."
        echo " ... you might need to start from a new shell."

        source config/day/day_env_installer.sh DAYOA
        if [[ "$?" != "0" ]]; then
            echo "Error: Failed to install the DAYOA environment. - $!"
            return 1
        else
            echo "DAYOA install success"
        fi
    fi

    conda activate DAYOA || {
        echo "Error: Failed to activate 'DAYOA' conda  environment."
    }
fi


# Final message
echo "Day CLI initialized for project '$PROJECT' in region '$region'."
echo -e "The Daylily CLI is now available."
echo -e "Available commands (tab completion is enabled for all):"
echo  " "
echo -e "\t(day-activate / dy-a) [slurm|local] ([hg38|hg38_broad|b37])- Activate a Slurm or local environment."
echo -e "\t                           "
echo -e "\t{day-set-genome-build / dy-g) [b37|hg38|hg38_broad] - Set the genome build for the current environment."
echo -e "\t                           "
echo -e "\t(day-run / dy-r)          - Run a command in the current environment."
echo -e "\t                           <tab> for exposed targets, -<tab> for all command line flags."
echo -e "\t                           "
echo -e "\t(day-monitor / dy-m)      - Monitor workflow status (Snakemake, SLURM, logs)."
echo -e "\t                           Use --block-and-poll to wait for completion."
echo -e "\t                           "
echo -e "\t(day-deactivate / dy-d)   - Deactivate the current environment."
echo -e "\t                           Use 'dy-d reset' to hard reset the environment."
echo -e "\t                           "
echo -e "\t                           "
echo -e "\t                           "
echo -e "\t                           "
echo -e "\t. To Stage Sample Data, see daylily-ephemeral-cluster docs for running ~/projects/daylily-ephemeral-cluster/bin/daylily-stage-analysis-samples-headnode "
echo -e "\t  ... which once run, you will copy the samples.tsv and units.tsv files to config/ in this directory."                 
echo -e "\t                           "
echo -e "\t(example): ACTIVATE AN ANALYSIS ENV\n"
echo -e "\n"
echo -e "\n"
echo -e "\t\tdy-a slurm hg38 # or hg38_broad or b37\n"

echo -e "\t(example): RUN ANALYSES\n"
echo -e "\t\tcp .test_data/data/0.01xwgs_HG002_hg38.samples.tsv config/samples.tsv\n"
echo -e "\t\tcp .test_data/data/0.01xwgs_HG002_hg38.units.tsv config/units.tsv\n"
echo -e "\n"
echo -e "\t\t# Use target names directly (tab-complete available):\n"
echo -e "\t\tdy-r produce_snv_concordances -p -k -j 20 -n   # Illumina short-read SNV concordance\n"
echo -e "\t\tdy-r produce_alignstats -p -k -j 20 -n         # Alignment statistics\n"
echo -e "\n"
echo -e "\t\t# Platform-specific targets:\n"
echo -e "\t\tdy-r produce_sentdont_vcf -p -k -j 20 -n       # ONT SNV calling\n"
echo -e "\t\tdy-r produce_sentdpb_vcf -p -k -j 20 -n        # PacBio SNV calling\n"
echo -e "\t\tdy-r produce_sentdug_vcf -p -k -j 20 -n        # Ultima SNV calling (use hg38_broad)\n"
echo -e "\n"
echo -e "\t\t# Hybrid workflow targets:\n"
echo -e "\t\tdy-r produce_sentdhio_vcf -p -k -j 20 -n       # Hybrid Illumina+ONT CLI\n"
echo -e "\t\tdy-r produce_sentdhuo_vcf -p -k -j 20 -n       # Hybrid Ultima+ONT CLI (use hg38_broad)\n"
echo -e "\t\tdy-r produce_sentdhiom_vcf -p -k -j 20 -n      # Hybrid Illumina+ONT Modular\n"
echo -e "\t\tdy-r produce_sentdhuom_vcf -p -k -j 20 -n      # Hybrid Ultima+ONT Modular (use hg38_broad)\n"
echo -e "\n"
echo -e "\t\t# Remove -n to execute (not dry-run)\n"

return 0