#!/usr/bin/env bash
# Interactive launcher for SoulGlitch donate-compute hosting (Linux bundle).
# Bundled next to the `soulglitch` binary by build_linux.sh.
#
# Usage:
#   ./run_donate_host.sh           # menu
#   ./run_donate_host.sh --start   # launch with saved/current settings (no menu)
set -euo pipefail

ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
BIN="$ROOT/soulglitch"
CONF="$ROOT/.donate_host.conf"

MODEL_IDS=(
  "HuggingFaceTB/SmolLM2-135M-Instruct"
  "HuggingFaceTB/SmolLM2-360M-Instruct"
  "HuggingFaceTB/SmolLM2-1.7B-Instruct"
  "microsoft/bitnet-b1.58-2B-4T"
  "Qwen/Qwen3-0.6B"
  "Qwen/Qwen3-1.7B"
  "Qwen/Qwen3-4B"
)
MODEL_TITLES=(
  "LITE BRAIN"
  "BALANCED BRAIN"
  "DEEP BRAIN"
  "BITNET BRAIN"
  "QWEN3 LITE"
  "QWEN3 BALANCED"
  "QWEN3 HEAVY"
)

PORT=17001
PODS="auto"
USE_GPU=0
OPEN_UI=0
MODEL_IDX=0

load_conf() {
  [[ -f "$CONF" ]] || return 0
  # shellcheck disable=SC1090
  source "$CONF"
}

save_conf() {
  cat >"$CONF" <<EOF
PORT=$PORT
PODS=$PODS
USE_GPU=$USE_GPU
OPEN_UI=$OPEN_UI
MODEL_IDX=$MODEL_IDX
EOF
}

hub_dirs() {
  printf '%s\n' \
    "${HOME}/.cache/huggingface/hub" \
    "${HOME}/Documents/huggingface/hub"
}

model_installed() {
  local id="$1"
  local repo="models--${id//\//--}"
  local hub
  while IFS= read -r hub; do
    [[ -d "$hub/$repo/snapshots" ]] || continue
    if find "$hub/$repo/snapshots" -name config.json 2>/dev/null | grep -q .; then
      return 0
    fi
  done < <(hub_dirs)
  return 1
}

pick_first_installed_model() {
  local i
  for ((i = 0; i < ${#MODEL_IDS[@]}; i++)); do
    if model_installed "${MODEL_IDS[$i]}"; then
      MODEL_IDX=$i
      return 0
    fi
  done
  MODEL_IDX=0
}

gpu_label() {
  if [[ "$USE_GPU" == 1 ]]; then echo "ON"; else echo "OFF"; fi
}

ui_label() {
  if [[ "$OPEN_UI" == 1 ]]; then echo "open Donate hub"; else echo "headless (TCP only)"; fi
}

current_model_id() {
  echo "${MODEL_IDS[$MODEL_IDX]}"
}

current_model_title() {
  echo "${MODEL_TITLES[$MODEL_IDX]}"
}

print_status() {
  local id title status
  id="$(current_model_id)"
  title="$(current_model_title)"
  if model_installed "$id"; then status="installed"; else status="NOT on disk"; fi
  echo ""
  echo "Current settings:"
  echo "  Model : [$((MODEL_IDX + 1))] $title"
  echo "          $id ($status)"
  echo "  Port  : $PORT"
  echo "  Pods  : $PODS"
  echo "  GPU   : $(gpu_label)"
  echo "  UI    : $(ui_label)"
  echo ""
}

print_menu() {
  echo "Menu:"
  echo "  1) Pick model"
  echo "  2) Set TCP port"
  echo "  3) Set pod count (auto or number)"
  echo "  4) Toggle GPU ($(gpu_label))"
  echo "  5) Toggle UI ($(ui_label))"
  echo "  6) START donate hosting"
  echo "  7) Run normal SoulGlitch (no donate)"
  echo "  q) Quit"
  echo ""
}

menu_pick_model() {
  local i id title mark
  echo ""
  echo "Models (download missing ones from Model Brain in the app):"
  for ((i = 0; i < ${#MODEL_IDS[@]}; i++)); do
    id="${MODEL_IDS[$i]}"
    title="${MODEL_TITLES[$i]}"
    if model_installed "$id"; then mark="installed"; else mark="missing"; fi
    echo "  $((i + 1))) $title — $id [$mark]"
  done
  echo ""
  read -r -p "Model number [1-${#MODEL_IDS[@]}]: " pick
  if [[ ! "$pick" =~ ^[0-9]+$ ]] || (( pick < 1 || pick > ${#MODEL_IDS[@]} )); then
    echo "Invalid model number."
    return
  fi
  MODEL_IDX=$((pick - 1))
  save_conf
}

menu_set_port() {
  read -r -p "TCP port [$PORT]: " pick
  [[ -z "$pick" ]] && return
  if [[ ! "$pick" =~ ^[0-9]+$ ]] || (( pick < 1 || pick > 65535 )); then
    echo "Invalid port."
    return
  fi
  PORT=$pick
  save_conf
}

menu_set_pods() {
  read -r -p "Pod count (auto or 1-256) [$PODS]: " pick
  [[ -z "$pick" ]] && return
  pick="${pick,,}"
  if [[ "$pick" == "auto" ]]; then
    PODS=auto
    save_conf
    return
  fi
  if [[ ! "$pick" =~ ^[0-9]+$ ]] || (( pick < 1 || pick > 256 )); then
    echo "Invalid pod count."
    return
  fi
  PODS=$pick
  save_conf
}

menu_toggle_gpu() {
  if [[ "$USE_GPU" == 1 ]]; then USE_GPU=0; else USE_GPU=1; fi
  save_conf
}

menu_toggle_ui() {
  if [[ "$OPEN_UI" == 1 ]]; then OPEN_UI=0; else OPEN_UI=1; fi
  save_conf
}

launch_donate() {
  local id args=()
  id="$(current_model_id)"
  if ! model_installed "$id"; then
    echo ""
    echo "WARN: $id is not installed on this machine."
    echo "      Hosting may fail until you download it from Model Brain."
    read -r -p "Continue anyway? [y/N]: " ok
    [[ "${ok,,}" == "y" || "${ok,,}" == "yes" ]] || return 0
  fi

  args=(
    --donate
    --skip-splash
    --mount
    --model "$id"
    --port "$PORT"
  )
  if [[ "$OPEN_UI" == 1 ]]; then
    args+=(--open)
  else
    args+=(--no-open)
  fi
  if [[ "$PODS" != "auto" ]]; then
    args+=(--pods "$PODS")
  fi
  if [[ "$USE_GPU" == 1 ]]; then
    args+=(--gpu)
  else
    args+=(--no-gpu)
  fi

  save_conf
  echo ""
  echo "Starting SoulGlitch donate host..."
  echo "  model=$id port=$PORT pods=$PODS gpu=$(gpu_label) ui=$(ui_label)"
  echo "  Ctrl+C to stop"
  echo ""
  exec "$BIN" "${args[@]}"
}

launch_normal() {
  exec "$BIN"
}

[[ -x "$BIN" ]] || {
  echo "soulglitch binary not found: $BIN"
  echo "Run this script from the Linux release bundle folder."
  exit 1
}

load_conf
if [[ "$MODEL_IDX" -lt 0 || "$MODEL_IDX" -ge ${#MODEL_IDS[@]} ]]; then
  pick_first_installed_model
fi

if [[ "${1:-}" == "--start" ]]; then
  launch_donate
fi

while true; do
  clear || true
  echo "============================================================"
  echo " SoulGlitch — Donate Compute Host"
  echo "============================================================"
  print_status
  print_menu
  read -r -p "Choice: " choice
  case "${choice,,}" in
    1) menu_pick_model ;;
    2) menu_set_port ;;
    3) menu_set_pods ;;
    4) menu_toggle_gpu ;;
    5) menu_toggle_ui ;;
    6) launch_donate ;;
    7) launch_normal ;;
    q|quit|exit) exit 0 ;;
    *) echo "Unknown option: $choice" ; sleep 1 ;;
  esac
done
