#!/bin/bash
# Ask Joshua installer: a private AI assistant that runs entirely on your computer (macOS or Linux).
#
#   curl -fsSL https://ask.joshua.earth/install.sh | bash
#
# What it does, in order (nothing is installed outside ~/AskJoshua, except Ollama and the Joshua launcher):
#   1. checks your computer: chip, graphics card, memory, free disk
#   2. estimates how fast each model would run on it
#   3. asks a few questions (deep thinking? images? fine-tuning?) and picks the models that fit
#   4. installs Ollama (the engine that runs the models), Python, and Joshua
#   5. downloads the models
#   6. adds Joshua to your Applications folder (Mac) or applications menu (Linux)
# Run it again at any time to change your choices; it skips what is already there.

set -euo pipefail

JOSHUA_URL="${JOSHUA_URL:-https://ask.joshua.earth/joshua.tar.gz}"
HOME_DIR="$HOME/AskJoshua"
APP="$HOME_DIR/joshua"
SRC_DIR="$(cd "$(dirname "${BASH_SOURCE[0]:-$0}")" 2>/dev/null && pwd || echo "")"

bold=$'\e[1m'; dim=$'\e[2m'; green=$'\e[32m'; amber=$'\e[33m'; red=$'\e[31m'; off=$'\e[0m'
say()  { printf "%s\n" "$*"; }
step() { printf "\n${bold}▸ %s${off}\n" "$*"; }
ok()   { printf "  ${green}✓${off} %s\n" "$*"; }
warn() { printf "  ${amber}!${off} %s\n" "$*"; }
die()  { printf "\n${red}✗ %s${off}\n" "$*"; exit 1; }
ask()  {  # ask "question" default(y/n) -> returns 0 for yes
  local q="$1" d="$2" a
  printf "  %s %s " "$q" "$([ "$d" = y ] && echo "[Y/n]" || echo "[y/N]")"
  { read -r a </dev/tty; } 2>/dev/null || read -r a || a=""
  a="${a:-$d}"; [[ "$a" =~ ^[YySs] ]]
}

cat <<'EOF'

     ██╗ ██████╗ ███████╗██╗  ██╗██╗   ██╗ █████╗
     ██║██╔═══██╗██╔════╝██║  ██║██║   ██║██╔══██╗
     ██║██║   ██║███████╗███████║██║   ██║███████║
██   ██║██║   ██║╚════██║██╔══██║██║   ██║██╔══██║
╚█████╔╝╚██████╔╝███████║██║  ██║╚██████╔╝██║  ██║
 ╚════╝  ╚═════╝ ╚══════╝╚═╝  ╚═╝ ╚═════╝ ╚═╝  ╚═╝
        your AI, on your computer, under your control
EOF

# ---------- 1. the computer ----------
step "Checking your computer"
OS=$(uname -s)
mkdir -p "$HOME_DIR"
FREE=$(( $(df -k "$HOME_DIR" | awk 'NR==2 {print $4}') / 1048576 ))
VRAM=0; GPU=""
if [ "$OS" = Darwin ]; then
  [ "$(uname -m)" = arm64 ] || die "Joshua needs a Mac with Apple Silicon (M1 or later). Intel Macs are too slow for local models."
  MACOS=$(sw_vers -productVersion); MAJOR=${MACOS%%.*}
  [ "$MAJOR" -ge 13 ] || die "macOS $MACOS is too old: Joshua needs macOS 13 (Ventura) or later."
  CHIP=$(sysctl -n machdep.cpu.brand_string)
  RAM=$(( $(sysctl -n hw.memsize) / 1073741824 ))
  ok "$CHIP · ${RAM} GB memory · ${FREE} GB free disk · macOS $MACOS"
elif [ "$OS" = Linux ]; then
  case "$(uname -m)" in x86_64|aarch64) ;; *) die "Joshua needs a 64-bit PC (x86_64 or ARM64).";; esac
  CHIP=$(awk -F: '/model name/ {sub(/^ +/, "", $2); print $2; exit}' /proc/cpuinfo); CHIP=${CHIP:-$(uname -m)}
  RAM=$(( $(awk '/MemTotal/ {print $2}' /proc/meminfo) / 1048576 + 1 ))
  if command -v nvidia-smi >/dev/null && LINE=$(nvidia-smi --query-gpu=name,memory.total --format=csv,noheader,nounits 2>/dev/null | head -1) && [ -n "$LINE" ]; then
    GPU=${LINE%,*}; VRAM=$(( (${LINE##*,} + 512) / 1024 ))
  fi
  ok "$CHIP · ${RAM} GB memory · ${FREE} GB free disk"
  if [ -n "$GPU" ]; then ok "Graphics: $GPU with ${VRAM} GB"
  else warn "No NVIDIA graphics card found: models will run on the processor, which is much slower. Images and fine-tuning need an NVIDIA card."; fi
else
  die "Joshua runs on macOS and Linux. On Windows, install it inside WSL (Windows Subsystem for Linux)."
fi
[ "$RAM" -ge 8 ] || die "Joshua needs at least 8 GB of memory."

# ---------- 2. how fast is it? ----------
# How fast a model writes depends mostly on memory bandwidth (GB/s): every word means reading the whole model once.
# How fast it draws pictures or trains depends on the graphics cores (G: M4 Mac = 10).
if [ "$OS" = Darwin ]; then
  GEN=$(echo "$CHIP" | sed -nE 's/.*M([0-9]+).*/\1/p'); GEN=${GEN:-9}
  case "$CHIP" in *Ultra*) TIER=Ultra;; *Max*) TIER=Max;; *Pro*) TIER=Pro;; *) TIER="";; esac
  case "$GEN$TIER" in
    1) BW=68;; 1Pro|2Pro) BW=200;; 1Max|2Max|3Max|4Max) BW=400;; 1Ultra|2Ultra|3Ultra) BW=800;;
    2|3) BW=100;; 3Pro) BW=150;; 4) BW=120;; 4Pro) BW=273;; 5) BW=153;; 5Pro) BW=300;; 5Max) BW=550;;
    *Ultra) BW=800;; *Max) BW=500;; *Pro) BW=270;; *) BW=150;;
  esac
  case "$GEN" in 1) G=6;; 2) G=7;; 3) G=9;; 4) G=10;; 5) G=13;; *) G=13;; esac
  case "$TIER" in Pro) G=$(( G * 18 / 10 ));; Max) G=$(( G * 35 / 10 ));; Ultra) G=$(( G * 7 ));; esac
  GPU_BW=$BW
else
  BW=50                                  # ordinary PC memory, used when a model doesn't fit on the card
  GPU_BW=$(( VRAM >= 16 ? 700 : VRAM >= 10 ? 450 : 300 ))   # typical NVIDIA cards of that size
  G=$(( VRAM >= 16 ? 30 : VRAM >= 10 ? 20 : 12 ))
fi
fits() { [ "$OS" = Darwin ] || [ "$VRAM" -ge "$1" ]; }   # does a model needing $1 GB fit on the graphics card?
words() { local bw=$BW; fits "$2" && bw=$GPU_BW; echo $(( bw * 35 * 75 / 10000 / $1 )); }  # words/s, reading $1 GB per word (calibrated: M4 Mac + 8B = 5)
SPEED4=$(words 3 4); SPEED8=$(words 6 7); SPEED30=$(words 4 22)   # the 30B model only reads ~3B parameters per word
PIC_MIN=$(( (15 * 10 / G + 5) / 10 )); [ "$PIC_MIN" -lt 1 ] && PIC_MIN=1   # minutes per picture
TUNE_MIN=$(( 17 * 10 / G ))                                               # minutes to train on 300 examples
if [ "$OS" = Darwin ]; then ok "Estimated: ${BW} GB/s memory bandwidth"; fi

# ---------- 3. questions and model choice ----------
step "A few questions"
# The memory the model must fit in: on a Mac the whole memory; on a PC the graphics card's (7 GB+ runs the 8B, 22 GB+ the 30B)
if [ "$OS" = Darwin ]; then BIG=$RAM; else BIG=$(( VRAM >= 22 ? 32 : VRAM >= 7 ? 16 : 8 )); fi
if [ "$BIG" -lt 16 ]; then
  SIZE=4b
  say "  ${dim}On this computer Joshua uses the 4-billion-parameter model: it sees images and handles everyday questions, writing about ${SPEED4} words a second.${off}"
elif [ "$BIG" -lt 32 ]; then
  SIZE=8b
  if [ "$SPEED8" -lt 4 ]; then
    say "  ${dim}You could run the 8-billion-parameter model, but here it would write only about ${SPEED8} words a second."
    say "  The 4-billion one writes about ${SPEED4} a second and is a little less capable.${off}"
    ask "Use the faster 4B model? (recommended)" y && SIZE=4b
  else
    say "  ${dim}Joshua uses the 8-billion-parameter model, the best fit for this computer: about ${SPEED8} words a second.${off}"
  fi
else
  SIZE=8b
  say "  ${dim}You can run a much larger model: 30 billion parameters, about 20 GB to download."
  say "  It only uses 3 billion of them for each word, so it stays quick: about ${SPEED30} words a second (the 8B: about ${SPEED8}).${off}"
  if [ "$FREE" -ge 30 ]; then ask "Use the larger model? (recommended)" y && SIZE=30b
  else say "  ${dim}You have ${FREE} GB free, not enough for it with room to spare, so Joshua will use the 8B.${off}"; fi
fi
case "$SIZE" in 4b) LABEL="qwen3-vl 4b"; MODEL_GB=3; SPEED=$SPEED4;; 8b) LABEL="qwen3-vl 8b"; MODEL_GB=6; SPEED=$SPEED8;; 30b) LABEL="qwen3-vl 30b"; MODEL_GB=20; SPEED=$SPEED30;; esac
FAST_CANDIDATES="qwen3-vl:${SIZE}-instruct"
THINK_CANDIDATES="qwen3-vl:${SIZE}-thinking qwen3-vl:${SIZE}"
[ "$SIZE" = 30b ] && FAST_CANDIDATES="qwen3-vl:30b-a3b-instruct qwen3-vl:30b-instruct" && THINK_CANDIDATES="qwen3-vl:30b-a3b-thinking qwen3-vl:30b-thinking qwen3-vl:30b"

say ""
say "  ${bold}Deep thinking${off}: a second version of the model that reasons step by step before answering."
say "  ${dim}Better on hard questions, slower. You switch it on per message. Another ~${MODEL_GB} GB.${off}"
DEF=y
if [ "$SPEED" -lt 4 ]; then DEF=n; say "  ${dim}On this computer its reasoning can take a minute or more before the answer starts.${off}"; fi
THINK=n; ask "Install the deep-thinking version?" $DEF && THINK=y

IMAGES=n; IMG_GB=7
say ""
if { [ "$OS" = Darwin ] && [ "$RAM" -ge 16 ]; } || { [ "$OS" = Linux ] && [ "$VRAM" -ge 12 ]; }; then
  [ "$OS" = Linux ] && IMG_GB=20   # PyTorch with CUDA, plus the full-precision image model
  say "  ${bold}Images${off}: Joshua can create pictures and edit your photos, on your computer (FLUX.2 klein)."
  say "  ${dim}About ${IMG_GB} GB more. A picture takes about ${PIC_MIN} min here, editing a photo several times longer.${off}"
  DEF=y; [ "$PIC_MIN" -gt 3 ] && DEF=n
  ask "Install image generation?" $DEF && IMAGES=y
elif [ "$OS" = Darwin ]; then
  say "  ${dim}Image generation needs 16 GB of memory, so it's not offered on this Mac. Joshua can still look at images.${off}"
else
  say "  ${dim}Image generation needs an NVIDIA card with 12 GB or more, so it's not offered here. Joshua can still look at images.${off}"
fi

TUNE=n; TUNE_GB=5
say ""
if [ "$OS" = Darwin ] || [ "$VRAM" -ge 8 ]; then
  [ "$OS" = Linux ] && TUNE_GB=12   # PyTorch with CUDA, plus the full-precision base model
  say "  ${bold}Fine-tuning${off}: teach Joshua your style with your own examples, guided step by step."
  say "  ${dim}About ${TUNE_GB} GB of tools now, and ~18 GB free while you install a trained model."
  say "  Training on 300 examples takes about ${TUNE_MIN} minutes here.${off}"
  ask "Install the fine-tuning tools?" n && TUNE=y
else
  say "  ${dim}Fine-tuning needs an NVIDIA card with 8 GB or more on Linux, so it's not offered here.${off}"
fi

NEED=$(( 2 + MODEL_GB ))
[ "$THINK" = y ] && NEED=$(( NEED + MODEL_GB ))
[ "$IMAGES" = y ] && NEED=$(( NEED + IMG_GB ))
[ "$TUNE" = y ] && NEED=$(( NEED + TUNE_GB ))
say ""
if [ "$FREE" -lt $(( NEED + 3 )) ]; then
  die "This needs about ${NEED} GB of disk (plus a little room to breathe) and you have ${FREE} GB free. Free some space, or run again and say no to some options."
fi
ok "Plan: ${LABEL} (~${SPEED} words/s)$([ "$THINK" = y ] && echo " + deep thinking")$([ "$IMAGES" = y ] && echo " + images")$([ "$TUNE" = y ] && echo " + fine-tuning") · about ${NEED} GB"
ask "Go ahead?" y || die "Nothing was installed."

# ---------- 4. Ollama, Python, Joshua ----------
step "Ollama (runs the models)"
OLLAMA=""
for c in "$(command -v ollama || true)" /Applications/Ollama.app/Contents/Resources/ollama "$HOME/Applications/Ollama.app/Contents/Resources/ollama"; do
  [ -n "$c" ] && [ -x "$c" ] && OLLAMA="$c" && break
done
if [ -n "$OLLAMA" ]; then
  ok "already installed"
elif [ "$OS" = Darwin ]; then
  TMP=$(mktemp -d)
  curl -fL --progress-bar https://ollama.com/download/Ollama-darwin.zip -o "$TMP/Ollama.zip"
  unzip -q "$TMP/Ollama.zip" -d "$TMP"
  DEST=/Applications; [ -w "$DEST" ] || { DEST="$HOME/Applications"; mkdir -p "$DEST"; }
  mv "$TMP/Ollama.app" "$DEST/"
  rm -rf "$TMP"
  OLLAMA="$DEST/Ollama.app/Contents/Resources/ollama"
  ok "installed in $DEST"
else
  say "  ${dim}Ollama's official Linux installer will ask for your password (it installs a system service).${off}"
  curl -fsSL https://ollama.com/install.sh | sh
  OLLAMA=$(command -v ollama) || die "Ollama didn't install. See https://ollama.com/download/linux"
fi
start_ollama() {
  curl -fs http://127.0.0.1:11434/api/version >/dev/null && return 0
  if [ "$OS" = Darwin ]; then open -a Ollama 2>/dev/null || ("$OLLAMA" serve >/dev/null 2>&1 &)
  else (nohup "$OLLAMA" serve >/dev/null 2>&1 &); fi
  for _ in $(seq 60); do curl -fs http://127.0.0.1:11434/api/version >/dev/null && return 0; sleep 1; done
  return 1
}
start_ollama || die "Ollama didn't start. Start it once by hand ('ollama serve'), then run this installer again."
ok "running"

if [ "$OS" = Linux ]; then
  step "Sandbox for Joshua's code"
  if command -v bwrap >/dev/null; then ok "bubblewrap is installed"
  else
    say "  ${dim}Joshua runs Python code in a sandbox (bubblewrap) so it can only write in the chat's folder.${off}"
    PKG=""; for p in apt-get dnf pacman zypper; do command -v $p >/dev/null && PKG=$p && break; done
    if [ -n "$PKG" ] && ask "Install bubblewrap now? (asks for your password)" y; then
      case $PKG in apt-get) sudo apt-get install -y bubblewrap;; dnf) sudo dnf install -y bubblewrap;; pacman) sudo pacman -S --noconfirm bubblewrap;; zypper) sudo zypper install -y bubblewrap;; esac
      ok "bubblewrap installed"
    else warn "without bubblewrap, Joshua's code runs without a sandbox; install the 'bubblewrap' package to add it"; fi
  fi
fi

step "Python"
UV="$(command -v uv || echo "$HOME/.local/bin/uv")"
if [ ! -x "$UV" ]; then
  curl -LsSf https://astral.sh/uv/install.sh | env UV_NO_MODIFY_PATH=1 sh >/dev/null
  UV="$HOME/.local/bin/uv"
fi
[ -x "$UV" ] || die "Couldn't install uv (the Python installer)."
[ -x "$HOME_DIR/venv/bin/python" ] || "$UV" venv -q --seed --python 3.12 "$HOME_DIR/venv"
PY="$HOME_DIR/venv/bin/python"
"$UV" pip install -q --python "$PY" pip   # lets Joshua install packages for its own code
ok "Python $("$PY" -c 'import platform; print(platform.python_version())') in $HOME_DIR/venv"

step "Joshua"
mkdir -p "$APP"
if [ -n "$SRC_DIR" ] && [ -f "$SRC_DIR/joshua/web.py" ]; then
  cp "$SRC_DIR"/joshua/*.py "$SRC_DIR"/joshua/requirements.txt "$APP/"
else
  curl -fsSL "$JOSHUA_URL" | tar -xz -C "$APP" --strip-components=1
fi
"$UV" pip install -q --python "$PY" -r "$APP/requirements.txt"
ok "installed in $APP"

# ---------- 5. models ----------
pull_first() {  # pull the first tag that exists; print its name
  for t in "$@"; do
    if "$OLLAMA" pull "$t" >&2; then echo "$t"; return 0; fi
    warn "$t is not available, trying the next one" >&2
  done
  return 1
}
step "Downloading the model (this is the long part)"
FAST=$(pull_first $FAST_CANDIDATES) || die "Couldn't download the model. Check your connection and run again."
ok "$FAST"
MODEL="$FAST"
if [ "$THINK" = y ]; then
  MODEL=$(pull_first $THINK_CANDIDATES) || { warn "deep-thinking version not available; skipping"; MODEL="$FAST"; }
  ok "$MODEL"
fi

if [ "$IMAGES" = y ]; then
  step "Image generation"
  if [ "$OS" = Darwin ]; then
    [ -x "$APP/mflux-venv/bin/python" ] || "$UV" venv -q --python 3.12 "$APP/mflux-venv"
    "$UV" pip install -q --python "$APP/mflux-venv/bin/python" "mflux==0.24.0"
    TEST=("$APP/mflux-venv/bin/mflux-generate-flux2" --model "RunPod/FLUX.2-klein-4B-mflux-4bit" --prompt "a small lighthouse at dusk" --width 256 --height 256 --steps 2 --output "$HOME_DIR/test-image.png")
  else
    [ -x "$APP/image-venv/bin/python" ] || "$UV" venv -q --python 3.12 "$APP/image-venv"
    "$UV" pip install -q --python "$APP/image-venv/bin/python" torch diffusers transformers accelerate safetensors pillow
    TEST=("$APP/image-venv/bin/python" "$APP/imagegen.py" --prompt "a small lighthouse at dusk" --width 256 --height 256 --out "$HOME_DIR/test-image.png")
  fi
  say "  ${dim}Downloading the image model and drawing a test picture…${off}"
  "${TEST[@]}" >/dev/null 2>&1 \
    && ok "working (test picture: $HOME_DIR/test-image.png)" \
    || warn "the test picture failed; Joshua will retry the download the first time you ask for an image"
fi

if [ "$OS" = Darwin ]; then TUNE_BASE="mlx-community/Qwen3-4B-Instruct-2507-4bit"; else TUNE_BASE="Qwen/Qwen3-4B-Instruct-2507"; fi
if [ "$TUNE" = y ]; then
  step "Fine-tuning tools"
  if [ "$OS" = Darwin ]; then
    [ "$MAJOR" -ge 14 ] || warn "macOS 14 or later is recommended for training (you have $MACOS)."
    "$UV" pip install -q --python "$PY" mlx-lm huggingface_hub
  else
    "$UV" pip install -q --python "$PY" torch transformers peft bitsandbytes accelerate huggingface_hub
  fi
  if [ ! -f "$HOME_DIR/llama.cpp/convert_hf_to_gguf.py" ]; then
    TMP=$(mktemp -d)
    curl -fsSL https://github.com/ggml-org/llama.cpp/archive/refs/heads/master.tar.gz | tar -xz -C "$TMP"
    rm -rf "$HOME_DIR/llama.cpp"
    mv "$TMP"/llama.cpp-master "$HOME_DIR/llama.cpp"   # we only use its Python converter; nothing is compiled
    rm -rf "$TMP"
  fi
  (cd "$HOME_DIR/llama.cpp/requirements" && "$UV" pip install -q --python "$PY" -r requirements-convert_hf_to_gguf.txt) \
    || "$UV" pip install -q --python "$PY" torch transformers sentencepiece protobuf numpy
  "$UV" pip install -q --python "$PY" "$HOME_DIR/llama.cpp/gguf-py"
  say "  ${dim}Downloading the model you will teach…${off}"
  "$PY" -c "from huggingface_hub import snapshot_download; snapshot_download('$TUNE_BASE')" >/dev/null 2>&1 \
    || warn "couldn't download it now; Joshua will download it when you first train"
  ok "ready: open Joshua and press 🧪 Teach Joshua"
fi

cat > "$APP/config.json" <<EOF
{"model": "$MODEL", "fast": "$FAST", "label": "$LABEL", "tune_base": "$TUNE_BASE", "tune_template": "qwen3:4b-instruct-2507-q4_K_M"}
EOF

# ---------- 6. launcher ----------
step "Joshua launcher"
if [ "$OS" = Darwin ]; then OPEN=open; START_OLLAMA="open -a Ollama 2>/dev/null || (\"$OLLAMA\" serve >/dev/null 2>&1 &)"
else OPEN=xdg-open; START_OLLAMA="(nohup \"$OLLAMA\" serve >/dev/null 2>&1 &)"; fi
cat > "$HOME_DIR/start.sh" <<EOF
#!/bin/bash
# Starts Joshua (and Ollama if needed), then opens it in your browser.
curl -fs http://127.0.0.1:11434/api/version >/dev/null || { $START_OLLAMA; sleep 4; }
if curl -fs http://127.0.0.1:7860 >/dev/null; then $OPEN http://127.0.0.1:7860; exit 0; fi
cd "$APP" && exec "$PY" "$APP/web.py" >> "$HOME_DIR/joshua.log" 2>&1
EOF
cat > "$HOME_DIR/stop.sh" <<EOF
#!/bin/bash
pkill -f "$APP/web.py" && echo "Joshua stopped." || echo "Joshua wasn't running."
EOF
if [ "$OS" = Darwin ]; then LAUNCHER="$HOME/Applications/Joshua.app"; else LAUNCHER="$HOME/.local/share/applications/joshua.desktop"; fi
cat > "$HOME_DIR/uninstall.sh" <<EOF
#!/bin/bash
# Removes Joshua. Your chats are in $APP/conversations: copy them first if you want to keep them.
read -r -p "Remove Joshua, its chats and its Python? [y/N] " a </dev/tty; [[ "\$a" =~ ^[Yy] ]] || exit 0
pkill -f "$APP/web.py" || true
rm -rf "$HOME_DIR" "$LAUNCHER"
read -r -p "Also delete the models from Ollama ($FAST $MODEL)? [y/N] " a </dev/tty
[[ "\$a" =~ ^[Yy] ]] && "$OLLAMA" rm $FAST $MODEL
echo "Done. Ollama itself is still installed; remove it separately if you don't need it."
EOF
chmod +x "$HOME_DIR"/{start,stop,uninstall}.sh
if [ "$OS" = Darwin ]; then
  mkdir -p "$HOME/Applications"
  rm -rf "$LAUNCHER"
  osacompile -o "$LAUNCHER" -e "do shell script \"nohup '$HOME_DIR/start.sh' >/dev/null 2>&1 &\"" >/dev/null 2>&1
  ok "Joshua is in your Applications folder (~/Applications/Joshua.app)"
else
  mkdir -p "$(dirname "$LAUNCHER")"
  cat > "$LAUNCHER" <<EOF
[Desktop Entry]
Type=Application
Name=Joshua
Comment=Your AI, on your computer
Exec=$HOME_DIR/start.sh
Icon=utilities-terminal
Terminal=false
Categories=Utility;
EOF
  ok "Joshua is in your applications menu"
fi

step "Done"
say "  Open ${bold}Joshua${off} from your applications. It opens in your browser at http://127.0.0.1:7860"
say "  ${dim}Everything lives in $HOME_DIR · stop it with stop.sh · remove it with uninstall.sh${off}"
if ask "Start Joshua now?" y; then (nohup "$HOME_DIR/start.sh" >/dev/null 2>&1 &); fi
