#!/usr/bin/env bash
set -euo pipefail

# === CONFIG ===
STYX_USER="styx"
STYX_PASS="TempPass#123"       # change at first login
STYX_HOSTNAME="styx"
SSH_PUBKEY=""                  # e.g., 'ssh-ed25519 AAAA... your@host' (optional)

# === helpers ===
require_root() {
  if [[ $EUID -ne 0 ]]; then exec sudo -E "$0" "$@"; fi
}

ensure_binfmt_arm64() {
  # If binfmt for arm64 isn't registered, register it via Docker (idempotent).
  if [[ ! -e /proc/sys/fs/binfmt_misc/qemu-aarch64 ]]; then
    command -v docker >/dev/null || { echo "Docker required to register binfmt (sudo pacman -S docker)"; exit 1; }
    sudo systemctl enable --now docker >/dev/null 2>&1 || true
    sudo docker run --rm --privileged tonistiigi/binfmt --install arm64
  fi
}

find_parts() {
  BOOT=$(lsblk -o LABEL,PATH -nr | awk '$1=="system-boot"{print $2}' | head -n1)
  ROOT=$(lsblk -o LABEL,PATH -nr | awk '$1=="writable"{print $2}' | head -n1)
  if [[ -z "${BOOT:-}" || -z "${ROOT:-}" ]]; then
    echo "Could not find 'system-boot'/'writable' on any device."
    lsblk -o NAME,SIZE,FSTYPE,LABEL,PATH -nr
    exit 1
  fi
}

mount_parts() {
  mkdir -p /mnt/pi-boot /mnt/pi-root
  mount "$ROOT" /mnt/pi-root
  mount "$BOOT" /mnt/pi-boot

  # Bind only what we need (avoid /run to prevent postinst fights)
  for d in dev dev/pts proc sys; do mount --bind "/$d" "/mnt/pi-root/$d"; done

  # Ubuntu images use a resolv.conf symlink—replace with a real file
  if [[ -L /mnt/pi-root/etc/resolv.conf || ! -e /mnt/pi-root/etc/resolv.conf ]]; then
    rm -f /mnt/pi-root/etc/resolv.conf
    cat /etc/resolv.conf > /mnt/pi-root/etc/resolv.conf
  fi
}

prep_chroot() {
  # Block service starts inside chroot (no systemd there)
  cat >/mnt/pi-root/usr/sbin/policy-rc.d <<'EOF'
#!/bin/sh
exit 101
EOF
  chmod +x /mnt/pi-root/usr/sbin/policy-rc.d

  # All the work happens inside the ARM64 rootfs
  CHCMD=$(cat <<'EOS'
set -euo pipefail
export DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC
# Ensure sbin is in PATH so user/group tools work
export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin"

apt-get update
apt-get -y full-upgrade

# Remove snaps and keep them gone (Ubuntu for Pi ships with snaps)
apt-get -y purge snapd || true
rm -rf /snap /var/snap /var/lib/snapd /home/*/snap || true
mkdir -p /etc/apt/preferences.d
printf 'Package: snapd\nPin: release *\nPin-Priority: -10\n' > /etc/apt/preferences.d/nosnap.pref

# Ensure user/group tools exist
apt-get install -y passwd adduser || true
getent group i2c >/dev/null || /usr/sbin/groupadd i2c

# Base packages
BASE_PKGS="openssh-server git i2c-tools python3-smbus python3-pil zbar-tools qrencode lm-sensors"
apt-get install -y $BASE_PKGS

# ------- OLED (Luma) -------
# Prefer distro package; fall back to pip if not present in this release
if ! dpkg -s python3-luma.oled >/dev/null 2>&1; then
  apt-get update
  if ! apt-get install -y python3-luma.oled; then
    apt-get install -y python3-pip
    pip3 install --no-input --break-system-packages luma.oled
  fi
fi

# ------- Camera apps -------
# Ubuntu renamed libcamera-apps -> rpicam-apps for Raspberry Pi.
# Try in order; tolerate absence (the box might be display-only).
apt-get update
if ! apt-get install -y rpicam-apps; then
  apt-get install -y libcamera-apps || apt-get install -y libcamera-tools || true
fi

# Enable SSH on boot (no systemctl in chroot)
mkdir -p /etc/systemd/system/multi-user.target.wants
ln -sf /lib/systemd/system/ssh.service /etc/systemd/system/multi-user.target.wants/ssh.service

# Create user and set password
if ! id -u STYX_USER >/dev/null 2>&1; then
  /usr/sbin/useradd -m -s /bin/bash -G sudo,video,i2c STYX_USER
fi
echo 'STYX_USER:STYX_PASS' | /usr/sbin/chpasswd

# Optional: preload SSH key
if [ -n 'SSH_PUBKEY' ] && echo 'SSH_PUBKEY' | grep -q 'ssh-'; then
  install -d -m700 /home/STYX_USER/.ssh
  echo 'SSH_PUBKEY' >> /home/STYX_USER/.ssh/authorized_keys
  chmod 600 /home/STYX_USER/.ssh/authorized_keys
  chown -R STYX_USER:STYX_USER /home/STYX_USER/.ssh
fi

# Freenove code
git clone https://github.com/Freenove/Freenove_Computer_Case_Kit_for_Raspberry_Pi.git /opt/freenove || true

# Hostname
echo 'STYX_HOSTNAME' > /etc/hostname
if grep -q '^127\.0\.1\.1' /etc/hosts; then
  sed -i 's/^127\.0\.1\.1.*/127.0.1.1\tSTYX_HOSTNAME/' /etc/hosts
else
  echo -e '127.0.1.1\tSTYX_HOSTNAME' >> /etc/hosts
fi

apt-get clean
EOS
)
  # Inject config values safely
  CHCMD="${CHCMD//STYX_USER/${STYX_USER}}"
  CHCMD="${CHCMD//STYX_PASS/${STYX_PASS}}"
  CHCMD="${CHCMD//STYX_HOSTNAME/${STYX_HOSTNAME}}"
  CHCMD="${CHCMD//SSH_PUBKEY/${SSH_PUBKEY}}"

  chroot /mnt/pi-root /bin/bash -lc "$CHCMD"
}

install_service_host() {
  # Systemd unit for the Freenove example app
  mkdir -p /mnt/pi-root/etc/systemd/system/multi-user.target.wants
  cat >/mnt/pi-root/etc/systemd/system/freenove-case.service <<'SERVICE'
[Unit]
Description=Freenove Case OLED/Fans/LEDs
After=multi-user.target

[Service]
Type=simple
ExecStart=/usr/bin/python3 /opt/freenove/Code/application.py
Restart=on-failure

[Install]
WantedBy=multi-user.target
SERVICE
  ln -sf /etc/systemd/system/freenove-case.service \
        /mnt/pi-root/etc/systemd/system/multi-user.target.wants/freenove-case.service || true
}

boot_tweaks() {
  # Enable I2C and set DSI panel on the BOOT partition
  grep -q 'dtparam=i2c_arm=on' /mnt/pi-boot/config.txt || echo 'dtparam=i2c_arm=on' >> /mnt/pi-boot/config.txt
  # Append kernel cmdline only once
  if ! grep -q 'DSI-1:800x480@60D' /mnt/pi-boot/cmdline.txt 2>/dev/null; then
    sed -i '1 s#$# video=DSI-1:800x480@60D video=HDMI-A-1:off video=HDMI-A-2:off#' /mnt/pi-boot/cmdline.txt || true
  fi
}

cleanup() {
  rm -f /mnt/pi-root/usr/sbin/policy-rc.d || true
  for d in dev/pts dev proc sys; do umount -lf "/mnt/pi-root/$d" 2>/dev/null || true; done
  umount -lf /mnt/pi-boot 2>/dev/null || true
  umount -lf /mnt/pi-root 2>/dev/null || true
  sync || true
}

main() {
  require_root
  ensure_binfmt_arm64
  find_parts
  trap 'echo "ERROR at line $LINENO" >&2; cleanup' ERR INT
  mount_parts
  prep_chroot
  install_service_host
  boot_tweaks
  cleanup
  echo "✅ Done. Move the NVMe to the Pi and boot."
  echo "   Login: user '${STYX_USER}'  pass '${STYX_PASS}'  (change with 'passwd')."
  echo "   Quick checks on the Pi:"
  echo "     sudo i2cdetect -y 1"
  echo "     rpicam-still -n -o test.jpg   # (if rpicam-apps installed)"
  echo "     libcamera-still -n -o test.jpg # (if legacy libcamera-apps installed)"
  echo "     systemctl status freenove-case"
}
main "$@"