#!/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  daylily-ec --help            : show all CLI subcommands"
echo -e "\t  daylily-ec cluster-info      : list clusters and their status"
echo -e "\t  daylily-ec version           : show installed version"
echo -e "\t  daylily-ec info              : show system info"
echo -e "\t "
echo -e "\t  day-build-env                : initialize the DAY-EC conda environment (same as dy-b)"
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  sq                           : alias to `squeue` with a useful default output format'
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
