#!/bin/bash

CONFIG_FILE=".tk1_loader_config"

# --- Internal Drivers ---
run_shofel2_t124() { sudo "$1" PAYLOAD "$2"; }
run_fusee_glee() { sudo "$1" "$2"; }
run_tegrarcm() { sudo "$1" --bootloader "$2" --loadaddr 0x80108000; }

load_config() { [ -f "$CONFIG_FILE" ] && source "$CONFIG_FILE"; }

save_config() {
  {
    echo "TOOL_TYPE=\"$TOOL_TYPE\""
    echo "TOOL_PATH=\"$TOOL_PATH\""
    echo "DEFAULT_PAYLOAD=\"$DEFAULT_PAYLOAD\""
    echo "CUSTOM_CMD=\"$CUSTOM_CMD\""
  } >"$CONFIG_FILE"
}

select_payload_file() {
  echo "--> Scanning all subdirectories for .bin files..."

  mapfile -t candidates < <(find . -type f -name "*.bin" -not -path "*/.*")

  if [ ${#candidates[@]} -eq 0 ]; then
    echo " [!] No .bin files found in $(pwd) or its subfolders."
    read -e -p "Please enter the path manually: " selected
  else
    echo "-------------------------------------"
    echo "Found ${#candidates[@]} bin files:"
    for i in "${!candidates[@]}"; do
      display_path=${candidates[$i]#./}
      echo " [$i] $display_path"
    done
    echo " [c] Custom Path"
    echo "-------------------------------------"

    read -p "Selection: " choice
    if [[ "$choice" == "c" ]]; then
      read -e -p "Enter path: " selected
    elif [[ "$choice" =~ ^[0-9]+$ ]] && [ "$choice" -lt "${#candidates[@]}" ]; then
      selected="${candidates[$choice]}"
    else
      echo "Invalid selection."
      exit 1
    fi
  fi
  echo "$selected"
}

do_setup() {
  echo "=== TK1 Loader Setup ==="
  echo "1) shofel2_T124 | 2) fusee_glee | 3) tegrarcm | 4) Custom"
  read -p "Tool Choice: " t_choice

  case $t_choice in
  1) TOOL_TYPE="shofel2_T124" ;;
  2) TOOL_TYPE="fusee_glee" ;;
  3) TOOL_TYPE="tegrarcm" ;;
  4) TOOL_TYPE="custom" ;;
  *) exit 1 ;;
  esac

  read -e -p "Path to tool executable: " TOOL_PATH

  if [ "$TOOL_TYPE" == "custom" ]; then
    echo "Define your command (use \$TOOL_PATH and \$PAYLOAD_PATH):"
    read -p "Example: sudo \$TOOL_PATH --flash \$PAYLOAD_PATH: " CUSTOM_CMD
  fi

  read -p "Set a default payload file now? (y/n): " set_def
  if [[ "$set_def" == "y" ]]; then
    DEFAULT_PAYLOAD=$(select_payload_file | tail -n 1)
  fi

  save_config
  echo "Setup complete."
}

do_push() {
  load_config
  if [ -z "$TOOL_TYPE" ]; then
    echo "Not configured. Run: $0 setup"
    exit 1
  fi

  # Determine which file to use
  if [ -n "$DEFAULT_PAYLOAD" ] && [ -f "$DEFAULT_PAYLOAD" ]; then
    FINAL_FILE="$DEFAULT_PAYLOAD"
    echo "--> Using default: $FINAL_FILE"
  else
    # We need to capture the output but let the user interact with the menu
    # This trick redirects the menu to stderr so it shows up, while the result goes to the variable
    exec 3>&1
    FINAL_FILE=$(select_payload_file 1>&2 2>&3)
    exec 3>&-
  fi

  if [ ! -f "$FINAL_FILE" ]; then
    echo "Error: File '$FINAL_FILE' not found."
    exit 1
  fi

  echo "--> Launching $TOOL_TYPE..."
  case $TOOL_TYPE in
  shofel2_T124) run_shofel2_t124 "$TOOL_PATH" "$FINAL_FILE" ;;
  fusee_glee) run_fusee_glee "$TOOL_PATH" "$FINAL_FILE" ;;
  tegrarcm) run_tegrarcm "$TOOL_PATH" "$FINAL_FILE" ;;
  custom)
    export TOOL_PATH="$TOOL_PATH"
    export PAYLOAD_PATH="$FINAL_FILE"
    eval "$CUSTOM_CMD"
    ;;
  esac
}

case "$1" in
setup) do_setup ;;
*) do_push ;;
esac