#!/bin/bash

# Daylily Headnode Configuration Script

set -euo pipefail

usage() {
    cat <<'EOF'
Usage: daylily-cfg-headnode [options] <pem_file> <region> <profile> [cluster_name] [repo_overrides_file]

Configures a ParallelCluster head node (installs config files, optional repo overrides, etc).

Options:
  --pem FILE                SSH PEM file
  --region REGION           AWS region
  --profile PROFILE         AWS CLI profile (exported as AWS_PROFILE)
  --cluster NAME            Cluster name (prompted if omitted)
  --repo-overrides FILE     File containing repo overrides (format: repo-key:git-ref per line)
  -h, --help                Show this help message and exit

Notes:
  This script is intended to be run from any working directory. It will
  locate packaged assets via `daylily-ec resources-dir` when installed.
EOF
}

resolve_daylily_res_dir() {
    if [[ -n "${DAYLILY_EC_RESOURCES_DIR:-}" ]]; then
        echo "${DAYLILY_EC_RESOURCES_DIR}"
        return 0
    fi
    if command -v daylily-ec >/dev/null 2>&1; then
        daylily-ec resources-dir
        return 0
    fi
    # Dev fallback: repo checkout.
    local script_dir repo_root
    script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
    repo_root="$(cd "${script_dir}/.." && pwd)"
    if [[ -d "${repo_root}/config" && -d "${repo_root}/bin" ]]; then
        echo "${repo_root}"
        return 0
    fi
    echo "Error: could not resolve Daylily resources dir. Set DAYLILY_EC_RESOURCES_DIR or install daylily-ephemeral-cluster." >&2
    return 1
}

yaml_get() {
    local yaml_path="$1"
    local dotted_key="$2"
    python3 - "$yaml_path" "$dotted_key" <<'PY'
import sys
import yaml

path, key = sys.argv[1], sys.argv[2]
cfg = yaml.safe_load(open(path, encoding="utf-8")) or {}
cur = cfg
for part in key.split("."):
    if not isinstance(cur, dict):
        cur = None
        break
    cur = cur.get(part)
if cur is None:
    print("")
else:
    print(cur)
PY
}

pem_file=""
region=""
aws_profile=""
requested_cluster_name=""
repo_overrides_file=""

positional_count=0
while [[ $# -gt 0 ]]; do
    case "$1" in
        -h|--help)
            usage
            exit 0
            ;;
        --pem)
            pem_file="$2"
            shift 2
            ;;
        --region)
            region="$2"
            shift 2
            ;;
        --profile)
            aws_profile="$2"
            shift 2
            ;;
        --cluster|--cluster-name)
            requested_cluster_name="$2"
            shift 2
            ;;
        --repo-overrides)
            repo_overrides_file="$2"
            shift 2
            ;;
        --)
            shift
            break
            ;;
        -*)
            echo "Unknown option: $1" >&2
            usage
            exit 1
            ;;
        *)
            case "$positional_count" in
                0) pem_file="$1" ;;
                1) region="$1" ;;
                2) aws_profile="$1" ;;
                3) requested_cluster_name="$1" ;;
                4) repo_overrides_file="$1" ;;
                *) echo "Unexpected argument: $1" >&2; usage; exit 1 ;;
            esac
            positional_count=$((positional_count + 1))
            shift
            ;;
    esac
done

if [[ -z "$pem_file" || -z "$region" || -z "$aws_profile" ]]; then
    echo "Error: pem_file, region, and profile are required." >&2
    usage
    exit 1
fi

RES_DIR="$(resolve_daylily_res_dir)" || exit 1
duser="ubuntu"

dayec_name="daylily-ephemeral-cluster"

CONFIG_FILE="${RES_DIR}/config/daylily_cli_global.yaml"
AVAILABLE_REPOS_CONFIG="${RES_DIR}/config/daylily_available_repositories.yaml"

# Create modified repository config if overrides are provided
MODIFIED_REPOS_CONFIG=""
if [[ -n "$repo_overrides_file" && -f "$repo_overrides_file" ]]; then
    echo "📦 Applying repository overrides from: $repo_overrides_file"
    MODIFIED_REPOS_CONFIG=$(mktemp)
    python3 - "$AVAILABLE_REPOS_CONFIG" "$repo_overrides_file" "$MODIFIED_REPOS_CONFIG" <<'PY'
import sys
import yaml

src, overrides_path, dst = sys.argv[1], sys.argv[2], sys.argv[3]
cfg = yaml.safe_load(open(src, encoding="utf-8")) or {}
repos = cfg.get("repositories", {}) or {}

with open(overrides_path, encoding="utf-8") as fh:
    for line in fh:
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        if ":" not in line:
            continue
        repo_key, git_ref = line.split(":", 1)
        repo_key, git_ref = repo_key.strip(), git_ref.strip()
        if repo_key in repos and isinstance(repos[repo_key], dict):
            repos[repo_key]["default_ref"] = git_ref

with open(dst, "w", encoding="utf-8") as out:
    yaml.safe_dump(cfg, out, default_flow_style=False, sort_keys=False)
PY

    echo "   Modified config saved to: $MODIFIED_REPOS_CONFIG"
fi

git_ephemeral_cluster_repo_tag="$(yaml_get "$CONFIG_FILE" "daylily.git_ephemeral_cluster_repo_tag")"
git_ephemeral_cluster_repo="$(yaml_get "$CONFIG_FILE" "daylily.git_ephemeral_cluster_repo")"

export AWS_PROFILE=$aws_profile

# List available clusters in the specified region
echo "Clusters detected in region $region:"
cluster_names=$(pcluster list-clusters --region $region | grep clusterName | awk '{print $2}' | cut -d '"' -f 2)

# Check if there are any clusters detected
if [[ -z "$cluster_names" ]]; then
    echo "Error: No clusters found in region $region."
    exit 1
fi
# Convert detected cluster names into an array
cluster_array=()
while IFS= read -r cluster_name; do
    cluster_array+=("$cluster_name")
done <<< "$cluster_names"

# Auto-select the requested cluster when possible.
selected_cluster=""
if [[ -n "$requested_cluster_name" ]]; then
    for cluster in "${cluster_array[@]}"; do
        if [[ "$cluster" == "$requested_cluster_name" ]]; then
            selected_cluster="$cluster"
            echo "Auto-selecting user-requested cluster: $selected_cluster"
            break
        fi
    done
    if [[ -z "$selected_cluster" ]]; then
        echo "⚠️ Requested cluster '$requested_cluster_name' not found in region $region."
    fi
fi

# Auto-select if there is only one cluster or if a requested cluster was matched.
if [[ -z "$selected_cluster" ]]; then
    if [[ ${#cluster_array[@]} -eq 1 ]]; then
        selected_cluster=${cluster_array[0]}
        echo "Only one cluster found: $selected_cluster. Auto-selecting it."
    else
        echo "Select a cluster name:"
        select selected_cluster in "${cluster_array[@]}"; do
            if [[ -n "$selected_cluster" ]]; then
                echo "You selected: $selected_cluster"
                break
            else
                echo "Invalid selection, please try again."
            fi
        done
    fi
fi
cluster_name=$selected_cluster


# Get the public IP address of the cluster's head node
cluster_ip_address=$(pcluster describe-cluster-instances -n "$cluster_name" --region "$region" \
    | grep publicIpAddress | perl -p -e 's/[ |"|,]//g;' | cut -d ':' -f 2)

if [[ -z "$cluster_ip_address" ]]; then
    echo "Error: Could not retrieve the public IP address of the cluster."
    exit 1
fi

echo "Cluster $cluster_name's public IP is $cluster_ip_address."
echo " "

# List available PEM files in the .ssh directory
echo "Detected PEM files in ~/.ssh:"
ls -1 ~/.ssh/*.pem

# If PEM file is not provided as an argument, prompt the user
if [[ -z "$pem_file" ]]; then
    echo "Enter the full absolute path to your PEM file:"
    read pem_file
fi

# Ensure the PEM file exists
if [[ ! -f "$pem_file" ]]; then
    echo "Error: PEM file '$pem_file' does not exist."
    exit 1
fi

# Generate SSH key for the head node user
echo "Generating SSH key on the head node..."
ssh -i "$pem_file" ubuntu@"$cluster_ip_address"   -o StrictHostKeyChecking=no   -o UserKnownHostsFile=/dev/null \
    "ssh-keygen -q -t rsa -f ~/.ssh/id_rsa -N '' <<< $'\ny' | sudo su - $duser"

# Display the public key and instruct the user to add it to GitHub
echo "You must have a GitHub account and access to the Daylily repository."
echo "Please add the following public SSH key to GitHub:"
git_key_log=$(ssh -i "$pem_file"  -o StrictHostKeyChecking=no   -o UserKnownHostsFile=/dev/null ubuntu@"$cluster_ip_address" "cat ~/.ssh/id_rsa.pub")
echo $git_key_log
mkdir -p logs
echo "$git_key_log" >> logs/${cluster_ip_address}_pubkey.txt

echo " "
echo "(optional) save this SSH key in GitHub to your settings->ssh/gpg keys, which will allow you to push changes back to github if you make any. You can always add this key later too, find it in you ~/.ssh/id_rsa.pub.\n\n\tSleeping for 15s then proceeding."
sleep 15

# Clone the $dayec_name repository to the head node
echo "Cloning $dayec_name repository using https to the head node ~/projects"
ssh -t -i "$pem_file" ubuntu@"$cluster_ip_address"  -o StrictHostKeyChecking=no   -o UserKnownHostsFile=/dev/null \
    "mkdir -p ~/projects && cd ~/projects &&  git clone -b ${git_ephemeral_cluster_repo_tag} ${git_ephemeral_cluster_repo} $dayec_name && cd $dayec_name && bin/install_miniconda || echo "failed to install miniconda" && bin/init_dayec || echo "failed to install day-ec" &&
. dyinit "

# Install shared tooling from the repo (day-clone, config files, etc.)
echo "Installing head node tooling from the cloned repository..."
ssh -t -i "$pem_file" ubuntu@"$cluster_ip_address" -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null \
    "source ~/.bashrc && cd ~/projects/$dayec_name && ./bin/install-daylily-headnode-tools"

# If we have a modified repository config with overrides, deploy it to the headnode
if [[ -n "$MODIFIED_REPOS_CONFIG" && -f "$MODIFIED_REPOS_CONFIG" ]]; then
    echo "📦 Deploying modified repository configuration with overrides to headnode..."
    scp -i "$pem_file" -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null \
        "$MODIFIED_REPOS_CONFIG" ubuntu@"$cluster_ip_address":~/.config/daylily/daylily_available_repositories.yaml
    echo "   ✅ Modified repository config deployed to ~/.config/daylily/daylily_available_repositories.yaml"

    # Clean up local temp file
    rm -f "$MODIFIED_REPOS_CONFIG"
fi

echo "Installing Miniconda and initializing the DAY-EC environment on the head node..."
ssh -t -i "$pem_file" ubuntu@"$cluster_ip_address" -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null \
    "source ~/.bashrc && cd ~/projects/$dayec_name && ./bin/install_miniconda && ./bin/init_dayec"

echo "day-clone is now available on the head node. Use it to clone analysis repositories such as daylily-omics-analysis."

# Provide final instructions for SSH access to the head node
echo "You can now SSH into the head node with the following command:"
echo "        ssh -i $pem_file ubuntu@$cluster_ip_address"
echo " "
echo "Once logged in, as the 'ubuntu' user, run the following commands:"
echo "  cd ~/projects/$dayec_name"
echo "  ./bin/install_miniconda"
echo "  ./bin/init_dayec"
echo "  ./bin/headnode_utils/day-clone --list"
echo "  day-clone -d <analysis_name> [--repository <repo-key>]"
echo " "
echo ".... this will clone the selected analysis repository into /fsx/analysis_results/<user>/<analysis_name>/"
echo " "
echo "Setup complete. "

echo " "
echo "And, you may now access the headnode via the PCUI, via 'source bin/daylily-ssh-into-headnode', or SSH into the head node with the following command:"
echo -e "       ssh -i $pem_file ubuntu@$cluster_ip_address\n"
echo " "
echo "If you wish to run some remote tests, you can do so with the following command:"
echo -e "   ./bin/daylily-run-ephemeral-cluster-remote-tests $pem_file $region $aws_profile\n"
echo " "
echo "...fin"
