#!/bin/bash

# Daylily Headnode Remote Tests Script

set -euo pipefail

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
    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
}

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

# ---------------------------------------------------------------------------
# Flag parsing  (backward-compatible: positional args still work)
# ---------------------------------------------------------------------------
_rt_usage() {
    cat <<'EOF'
Usage: daylily-run-ephemeral-cluster-remote-tests <pem_file> <region> <profile> [options]

Positional (required):
  pem_file       Path to SSH PEM key
  region         AWS region
  profile        AWS CLI profile name

Options:
  --cluster NAME   Cluster name (skip interactive cluster selection)
  --yes            Auto-answer yes to the workflow launch prompt
  --no-launch      Auto-answer no to the workflow launch prompt
  -h, --help       Show this help message and exit

When --cluster is omitted the script falls back to interactive selection.
When neither --yes nor --no-launch is provided, the script prompts interactively.
EOF
}

flag_cluster=""
flag_launch=""  # "", "yes", or "no"

# Separate positional args from flags
_positionals=()
while [[ $# -gt 0 ]]; do
    case "$1" in
        --cluster)    flag_cluster="$2"; shift 2 ;;
        --yes)        flag_launch="yes"; shift   ;;
        --no-launch)  flag_launch="no";  shift   ;;
        -h|--help)    _rt_usage; exit 0          ;;
        -*)           echo "Unknown option: $1" >&2; _rt_usage; exit 1 ;;
        *)            _positionals+=("$1"); shift ;;
    esac
done

pem_file="${_positionals[0]:-}"
region="${_positionals[1]:-}"
aws_profile="${_positionals[2]:-}"
duser="ubuntu"

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

default_repo=$(python - "$CONFIG_FILE" <<'PY'
import sys
import yaml

cfg = yaml.safe_load(open(sys.argv[1], encoding="utf-8")) or {}
print(cfg.get("default_repository", "") or "daylily-omics-analysis")
PY
)

git_tag=$(python - "$CONFIG_FILE" "$default_repo" <<'PY'
import sys
import yaml

cfg = yaml.safe_load(open(sys.argv[1], encoding="utf-8")) or {}
repo_key = sys.argv[2]
repos = cfg.get("repositories", {}) or {}
repo = repos.get(repo_key, {}) or {}
print(repo.get("default_ref", "main"))
PY
)

git_repo=$(python - "$CONFIG_FILE" "$default_repo" <<'PY'
import sys
import yaml

cfg = yaml.safe_load(open(sys.argv[1], encoding="utf-8")) or {}
repo_key = sys.argv[2]
repos = cfg.get("repositories", {}) or {}
repo = repos.get(repo_key, {}) or {}
print(repo.get("https_url", ""))
PY
)

# Ensure both PEM file and region are provided
if [[ -z "$pem_file" || -z "$region" || -z "$aws_profile" ]]; then
    echo "Error: You must provide both the PEM file path and the AWS region."
    echo "Usage: $0 /path/to/pem_file region profile [--cluster NAME] [--yes|--no-launch]"
    exit 1
fi
 
export AWS_PROFILE=$aws_profile

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

# 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 from flag, single-match, or prompt interactively
if [[ -n "$flag_cluster" ]]; then
    selected_cluster="$flag_cluster"
    echo "Using cluster from --cluster flag: $selected_cluster"
elif [[ ${#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
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 || true)

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 2>/dev/null || true

# 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


# Determine whether to launch the remote workflow
if [[ "$flag_launch" == "yes" ]]; then
    REPLY="y"
elif [[ "$flag_launch" == "no" ]]; then
    REPLY="n"
else
    echo "Would you like to remotely launch a workflow w/in a tmux session on the headnode? [y/n]"
    read REPLY
fi
if [[ "$REPLY" == "y" ]]; then

    ssh -i "$pem_file" ubuntu@"$cluster_ip_address"  -o StrictHostKeyChecking=no   -o UserKnownHostsFile=/dev/null \
        "sudo su - $duser -c 'mkdir -p /fsx/analysis_results/ubuntu/daylily_remote_test && cd /fsx/analysis_results/ubuntu/daylily_remote_test &&  git clone -b ${git_tag} ${git_repo} daylily-omics-analysis'"

    session_name="cluster_test_$(date +%s)"  
    ssh -t -i "$pem_file" ubuntu@"$cluster_ip_address" -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null \
        "tmux new-session -d -s \"$session_name\" \
        \"bash -c 'source ~/.bashrc && cd /fsx/analysis_results/ubuntu/daylily_remote_test/daylily-omics-analysis && source dyinit && source bin/day_activate slurm hg38 remote && DAY_CONTAINERIZED=true ./bin/day_run produce_snv_concordances -p -k -j 2 --config aligners=[\\\"strobe\\\",\\\"bwa2a\\\"] dedupers=[\\\"dppl\\\"] genome_build=\\\"hg38\\\" snv_callers=[\\\"deep\\\"]; bash'\"" && \

    echo "Tmux session >>> '$session_name' <<< started on the remote server."
    echo "You can monitor the session using the following command:"
    echo "ssh -i \"$pem_file\" ubuntu@\"$cluster_ip_address\" -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null 'tmux attach-session -t $session_name'"
    echo "... or alternatively, you can SSH into the head node ( ssh -i $pem_file ubuntu@$cluster_ip_address )a, nd run ( tmux attach-session -t $session_name ) to monitor the session."
fi



echo ""
echo " >>>>>> 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 "        ssh -i $pem_file ubuntu@$cluster_ip_address"
echo "        tmux ls"
echo "        tmux a -t c"
echo " "
echo " "
echo "...fin"
