#!/usr/bin/env bash

CONFIG_FILE="$HOME/.talos/config"

COMMAND="$(basename "$0")"

# Check yq
if ! command -v yq &>/dev/null; then
    echo "Error: yq (Go version v4+) not found."
    exit 1
fi

mapfile -t CONTEXTS < <(yq e '.contexts | keys | .[]' "$CONFIG_FILE")

contains() {
    local seeking=$1; shift
    for element; do
        [[ "$element" == "$seeking" ]] && return 0
    done
    return 1
}

print_help() {
    cat <<EOF
Usage:
  $COMMAND <context> [node] <command>

Description:
  Wrapper for talosctl with context and optional node selection.

Examples:
  $COMMAND home dashboard
  $COMMAND home talos-nuc-1 dashboard
  $COMMAND zl-erm talos-geekom-1 kubeconfig

Arguments:
  context   Required. Must match a configured Talos context.
  node      Optional. Node name or short name.
  command   Any talosctl command.

Commands:
  help, -h, --help   Show this help message
EOF
}

# Help handling
if [[ "$1" == "help" || "$1" == "-h" || "$1" == "--help" ]]; then
    print_help
    exit 0
fi

# Require context
if [ -z "$1" ]; then
    echo "Error: Context is required."
    print_help
    exit 1
fi

# Validate context
if ! contains "$1" "${CONTEXTS[@]}"; then
    echo "Error: Invalid context '$1'"
    echo "Available contexts:"
    printf '  - %s\n' "${CONTEXTS[@]}"
    exit 1
fi

CONTEXT="$1"
shift

mapfile -t NODES < <(yq e ".contexts.${CONTEXT}.nodes[]" "$CONFIG_FILE")

NODE=""
NODE_INPUT=""

if [ -n "$1" ]; then
    NODE_INPUT="$1"

    for n in "${NODES[@]}"; do
        if [[ "$n" == "$NODE_INPUT" || "${n%%.*}" == "$NODE_INPUT" ]]; then
            NODE="$n"
            shift
            break
        fi
    done

    if [ -n "$NODE_INPUT" ] && [ -z "$NODE" ]; then
        echo "Error: Node '$NODE_INPUT' not found in context '$CONTEXT'"
        echo "Available nodes:"
        printf '  - %s\n' "${NODES[@]}"
        exit 1
    fi
fi

if [ -n "$NODE" ]; then
    talosctl --context "$CONTEXT" -n "$NODE" -e "$NODE" "$@"
else
    talosctl --context "$CONTEXT" "$@"
fi