#!/usr/bin/env bash
set -euo pipefail
DEVICE="${1:-}"
shift || true
DIST="${DIST:-bookworm}"
MODEL="${JACKAILOCAL_MODEL:-qwen3.5:4b}"
CONFIRM=""
for arg in "$@"; do
  case "$arg" in
    --confirm-erase) CONFIRM="ERASE" ;;
    *) echo "Unknown argument: $arg" >&2; exit 2 ;;
  esac
done
if [[ -z "$DEVICE" ]]; then
  echo "Usage: sudo $0 /dev/sdX [--confirm-erase]"
  exit 1
fi
if [[ $EUID -ne 0 ]]; then echo "Run this script as root." >&2; exit 1; fi
if [[ ! -b "$DEVICE" ]]; then echo "$DEVICE is not a block device"; exit 1; fi
if [[ "$CONFIRM" != "ERASE" ]]; then
  read -r -p "This will ERASE $DEVICE. Type ERASE to continue: " CONFIRM
fi
[[ "$CONFIRM" == "ERASE" ]] || exit 1

ROOT_SRC="$(cd "$(dirname "$0")/../.." && pwd)"
if findmnt -rn -S "${DEVICE}*" >/dev/null 2>&1; then
  echo "A partition on $DEVICE is mounted. Unmount it before continuing." >&2
  exit 1
fi

"$ROOT_SRC/factory/linux/build-linux-binary.sh"
"$ROOT_SRC/factory/linux/install-linux-backends.sh"
"$ROOT_SRC/factory/unix/build-whisper-backend.sh"
"$ROOT_SRC/unix/Install-Voice.sh" "$ROOT_SRC"
[ -x "$ROOT_SRC/bin/jackailocald" ] || { echo "Missing staged Linux runtime." >&2; exit 1; }
[ -x "$ROOT_SRC/backends/ollama/linux/ollama" ] || { echo "Missing staged Linux Ollama backend." >&2; exit 1; }

mkdir -p "$ROOT_SRC/models/ollama" "$ROOT_SRC/.jackailocal-builder/logs"
export OLLAMA_MODELS="$ROOT_SRC/models/ollama"
export OLLAMA_HOST="127.0.0.1:11435"
"$ROOT_SRC/backends/ollama/linux/ollama" serve > "$ROOT_SRC/.jackailocal-builder/logs/ollama-linux-builder.log" 2>&1 &
OLLAMA_PID=$!
cleanup_ollama() { kill "$OLLAMA_PID" >/dev/null 2>&1 || true; }
trap cleanup_ollama EXIT
for _ in $(seq 1 60); do
  "$ROOT_SRC/backends/ollama/linux/ollama" list >/dev/null 2>&1 && break
  sleep 0.5
done
"$ROOT_SRC/backends/ollama/linux/ollama" pull "$MODEL"
cleanup_ollama
trap - EXIT

apt-get update
apt-get install -y parted dosfstools e2fsprogs debootstrap grub-efi-amd64-bin rsync
parted -s "$DEVICE" mklabel gpt
parted -s "$DEVICE" mkpart ESP fat32 1MiB 513MiB
parted -s "$DEVICE" set 1 esp on
parted -s "$DEVICE" mkpart root ext4 513MiB 100%
partprobe "$DEVICE"
udevadm settle
EFI="${DEVICE}1"; ROOT="${DEVICE}2"
if [[ "$DEVICE" =~ nvme|mmcblk ]]; then EFI="${DEVICE}p1"; ROOT="${DEVICE}p2"; fi
mkfs.vfat -F32 -n JACKAI_EFI "$EFI"
mkfs.ext4 -F -L JACKAI_ROOT "$ROOT"
MOUNT_ROOT=/mnt/jackailocal-root
mkdir -p "$MOUNT_ROOT"
mount "$ROOT" "$MOUNT_ROOT"
cleanup_mounts() { umount -R "$MOUNT_ROOT" >/dev/null 2>&1 || true; }
trap cleanup_mounts EXIT
mkdir -p "$MOUNT_ROOT/boot/efi"
mount "$EFI" "$MOUNT_ROOT/boot/efi"
debootstrap --arch amd64 "$DIST" "$MOUNT_ROOT" http://deb.debian.org/debian
cat > "$MOUNT_ROOT/etc/fstab" <<FSTAB
LABEL=JACKAI_ROOT / ext4 defaults,noatime 0 1
LABEL=JACKAI_EFI /boot/efi vfat umask=0077 0 1
FSTAB
mount --bind /dev "$MOUNT_ROOT/dev"
mount --bind /proc "$MOUNT_ROOT/proc"
mount --bind /sys "$MOUNT_ROOT/sys"
chroot "$MOUNT_ROOT" /bin/bash -lc "apt-get update && apt-get install -y linux-image-amd64 grub-efi-amd64 network-manager dbus xserver-xorg xinit chromium curl ca-certificates jq python3 pciutils usbutils"
chroot "$MOUNT_ROOT" /bin/bash -lc "grub-install --target=x86_64-efi --efi-directory=/boot/efi --bootloader-id=JackAILocal --removable && update-grub"
chroot "$MOUNT_ROOT" /bin/bash -lc "useradd -m -s /bin/bash jackai && usermod -aG audio,video,render jackai"
mkdir -p "$MOUNT_ROOT/opt/jackailocal"
rsync -a \
  --exclude '.git' --exclude target --exclude downloads --exclude workspace --exclude diagnostics \
  --exclude '.jackailocal' --exclude '.jackailocal-builder' --exclude 'config/update-private-key.xml' \
  --exclude 'license-keys' --exclude '__pycache__' \
  "$ROOT_SRC/" "$MOUNT_ROOT/opt/jackailocal/"
chroot "$MOUNT_ROOT" /bin/bash -lc "bash /opt/jackailocal/linux/runtime/install-linux-runtime.sh"
cleanup_mounts
trap - EXIT
echo "Linux boot disk prepared: $DEVICE"