#!/bin/bash
# First-boot: packages, Chromium kiosk, rpd-labwc + touch, reTerminal DM drivers.
# Run by cloud-init (user-data-remote-gnss.example). Run as root.

set -e
export DEBIAN_FRONTEND=noninteractive

# --- Constants ---
# All first-boot assets live in portal-files/first-boot/ on the file server.
FILE_SERVER="http://10.130.60.141:5000/files/first-boot"
HOSTNAME="gnss.guard"
PI_USER="pi"
PI_HOME="/home/$PI_USER"
AUTOSTART="$PI_HOME/.config/autostart"
LOGFILE="/var/log/first-boot.log"
PLYMOUTH_DIR="/usr/share/plymouth/themes/custom"
WALLPAPER_PATH="/usr/share/rpd-wallpaper/splash.png"

# --- Logging ---
log() { echo "[$(date -Iseconds)] $*"; }
exec > >(tee -a "$LOGFILE") 2>&1
log "=== first-boot.sh started ==="
log "FILE_SERVER=$FILE_SERVER PI_USER=$PI_USER LOGFILE=$LOGFILE"

# --- 0. Hostname and /etc/hosts (avoids "unable to resolve host" with sudo) ---
log "--- Hostname: $HOSTNAME ---"
echo "$HOSTNAME" > /etc/hostname
hostnamectl set-hostname "$HOSTNAME" 2>/dev/null || true
# Ensure hostname resolves so sudo and other tools don't warn
if ! grep -q "127.0.1.1[[:space:]]*$HOSTNAME" /etc/hosts 2>/dev/null; then
  sed -i "/127.0.1.1[[:space:]].*$/d" /etc/hosts
  echo "127.0.1.1	$HOSTNAME" >> /etc/hosts
fi
log "Hostname set to $HOSTNAME; /etc/hosts updated"

# --- Helpers ---
# Download script + .desktop from FILE_SERVER and install as one-shot autostart (runs once at pi's first login, then deletes itself).
install_oneshot() {
  local name="$1"
  log "--- Installing one-shot: $name ---"
  if curl -fsSL "${FILE_SERVER}/${name}.sh" -o "$PI_HOME/${name}.sh"; then
    log "Downloaded ${name}.sh to $PI_HOME/${name}.sh"
  else
    log "WARNING: Could not download ${name}.sh"; return 1
  fi
  if curl -fsSL "${FILE_SERVER}/${name}.desktop" -o "$AUTOSTART/${name}.desktop"; then
    log "Downloaded ${name}.desktop to $AUTOSTART/${name}.desktop"
  else
    log "WARNING: Could not download ${name}.desktop"; return 1
  fi
  chmod 755 "$PI_HOME/${name}.sh" && chmod 644 "$AUTOSTART/${name}.desktop"
  chown "$PI_USER:$PI_USER" "$PI_HOME/${name}.sh" "$AUTOSTART/${name}.desktop"
  log "One-shot $name installed (will run at first login and then remove itself)"
}

# --- 1. Packages ---
log "--- Installing packages ---"
log "Running apt-get update ..."
apt-get update -qq
log "Installing: git chromium wmctrl openssh-server swaybg wlr-randr maliit-keyboard xinput-calibrator rpi-eeprom"
apt-get install -y -qq git chromium wmctrl openssh-server \
  swaybg wlr-randr maliit-keyboard xinput-calibrator rpi-eeprom
log "Packages installed successfully"

# --- 2. Dirs and kiosk files from file server ---
log "--- Kiosk files ---"
log "Creating $AUTOSTART"
mkdir -p "$AUTOSTART"
log "Downloading start-chromium.sh from ${FILE_SERVER}/start-chromium.sh"
curl -fsSL "${FILE_SERVER}/start-chromium.sh" -o "$PI_HOME/start-chromium.sh"
log "Downloading chromium-kiosk.desktop from ${FILE_SERVER}/chromium-kiosk.desktop"
curl -fsSL "${FILE_SERVER}/chromium-kiosk.desktop" -o "$AUTOSTART/chromium-kiosk.desktop"
chmod 755 "$PI_HOME/start-chromium.sh" && chmod 644 "$AUTOSTART/chromium-kiosk.desktop"
chown -R "$PI_USER:$PI_USER" "$PI_HOME/start-chromium.sh" "$AUTOSTART/chromium-kiosk.desktop"
log "Kiosk files installed under $PI_HOME and $AUTOSTART"

# --- 3. Boot splash and wallpaper (splash.png + Plymouth theme from file server) ---
log "--- Boot splash and wallpaper ---"
log "Creating $PLYMOUTH_DIR and /usr/share/rpd-wallpaper"
mkdir -p "$PLYMOUTH_DIR" /usr/share/rpd-wallpaper
if curl -fsSL "${FILE_SERVER}/splash.png" -o "$PLYMOUTH_DIR/splash.png"; then
  log "Downloaded splash.png; copying to $WALLPAPER_PATH"
  cp "$PLYMOUTH_DIR/splash.png" "$WALLPAPER_PATH"
  chmod 644 "$PLYMOUTH_DIR/splash.png" "$WALLPAPER_PATH"
  if curl -fsSL "${FILE_SERVER}/custom.plymouth" -o "$PLYMOUTH_DIR/custom.plymouth" \
     && curl -fsSL "${FILE_SERVER}/custom.script" -o "$PLYMOUTH_DIR/custom.script"; then
    chmod 644 "$PLYMOUTH_DIR/custom.plymouth" "$PLYMOUTH_DIR/custom.script"
    log "Plymouth theme files (custom.plymouth, custom.script) installed"
  else
    log "WARNING: Could not download custom.plymouth/custom.script; boot splash theme may be incomplete"
  fi
  grep -q '^Theme=custom' /etc/plymouth/plymouthd.conf 2>/dev/null || printf '%s\n' '[Daemon]' 'Theme=custom' >> /etc/plymouth/plymouthd.conf
  log "Running update-initramfs (may take a moment) ..."
  update-initramfs -u -k all 2>/dev/null || true
  mkdir -p /etc/lightdm/lightdm.conf.d
  curl -fsSL "${FILE_SERVER}/99-wallpaper.conf" -o /etc/lightdm/lightdm.conf.d/99-wallpaper.conf 2>/dev/null || log "WARNING: Could not download 99-wallpaper.conf"
  # Set desktop wallpaper once via pcmanfm config (rpd-labwc uses pcmanfm-pi; profile LXDE-pi or default)
  for PROFILE in LXDE-pi default; do
    PCMANFM_DESKTOP="$PI_HOME/.config/pcmanfm/$PROFILE/desktop-items-0.conf"
    mkdir -p "$(dirname "$PCMANFM_DESKTOP")"
    if [[ ! -f "$PCMANFM_DESKTOP" ]]; then
      printf '%s\n' '[*]' "wallpaper=$WALLPAPER_PATH" 'wallpaper_mode=crop' 'wallpaper_common=1' > "$PCMANFM_DESKTOP"
    else
      grep -q '^wallpaper=' "$PCMANFM_DESKTOP" && sed -i "s|^wallpaper=.*|wallpaper=$WALLPAPER_PATH|" "$PCMANFM_DESKTOP" || echo "wallpaper=$WALLPAPER_PATH" >> "$PCMANFM_DESKTOP"
      grep -q '^wallpaper_mode=' "$PCMANFM_DESKTOP" && sed -i 's/^wallpaper_mode=.*/wallpaper_mode=crop/' "$PCMANFM_DESKTOP" || echo 'wallpaper_mode=crop' >> "$PCMANFM_DESKTOP"
    fi
    chown -R "$PI_USER:$PI_USER" "$(dirname "$PCMANFM_DESKTOP")"
  done
  log "Set desktop wallpaper via pcmanfm config (LXDE-pi and default)"
  log "Splash and wallpaper set from file server"
else
  log "WARNING: Could not download splash.png"
fi

# --- 4. LightDM: rpd-labwc session + configs from file server ---
log "--- LightDM session (rpd-labwc) ---"
mkdir -p /etc/lightdm/lightdm.conf.d
if curl -fsSL "${FILE_SERVER}/99-default-session.conf" -o /etc/lightdm/lightdm.conf.d/99-default-session.conf 2>/dev/null; then
  log "99-default-session.conf installed"
else
  log "WARNING: Could not download 99-default-session.conf"
fi
# Raspberry Pi OS may apply main lightdm.conf after .conf.d; force session in main config too
if [[ -f /etc/lightdm/lightdm.conf ]]; then
  sed -i 's/^user-session=.*/user-session=rpd-labwc/' /etc/lightdm/lightdm.conf
  sed -i 's/^autologin-session=.*/autologin-session=rpd-labwc/' /etc/lightdm/lightdm.conf
  log "Patched /etc/lightdm/lightdm.conf to use rpd-labwc"
fi

# --- 5. Maliit on-screen keyboard (from file server) ---
log "--- Maliit ---"
mkdir -p "$AUTOSTART" "$PI_HOME/.config"
curl -fsSL "${FILE_SERVER}/maliit-keyboard.desktop" -o "$AUTOSTART/maliit-keyboard.desktop" 2>/dev/null && log "maliit-keyboard.desktop installed" || log "WARNING: Could not download maliit-keyboard.desktop"

# --- 5b. Dark theme (GTK + prefer dark for apps) ---
log "--- Dark theme ---"
GTK_SETTINGS="$PI_HOME/.config/gtk-3.0/settings.ini"
mkdir -p "$(dirname "$GTK_SETTINGS")"
if [[ ! -f "$GTK_SETTINGS" ]]; then
  printf '%s\n' '[Settings]' 'gtk-application-prefer-dark-theme=1' 'gtk-theme-name=PiXnoir' > "$GTK_SETTINGS"
else
  grep -q '^gtk-application-prefer-dark-theme=' "$GTK_SETTINGS" && sed -i 's/^gtk-application-prefer-dark-theme=.*/gtk-application-prefer-dark-theme=1/' "$GTK_SETTINGS" || echo 'gtk-application-prefer-dark-theme=1' >> "$GTK_SETTINGS"
  grep -q '^gtk-theme-name=' "$GTK_SETTINGS" && sed -i 's/^gtk-theme-name=.*/gtk-theme-name=PiXnoir/' "$GTK_SETTINGS" || echo 'gtk-theme-name=PiXnoir' >> "$GTK_SETTINGS"
fi
# Fallback if PiXnoir not installed (e.g. older image): Adwaita-dark
log "Set dark theme (PiXnoir) in gtk-3.0/settings.ini"
chown -R "$PI_USER:$PI_USER" "$PI_HOME/.config"

# --- 6. reTerminal DM drivers (Seeed) ---
log "--- reTerminal DM drivers ---"
REPO_DIR="/tmp/seeed-linux-dtoverlays"
log "Cloning seeed-linux-dtoverlays to $REPO_DIR ..."
git clone --depth 1 https://github.com/Seeed-Studio/seeed-linux-dtoverlays "$REPO_DIR"
# Script must run from repo root (it uses pwd for MOD_PATH). On bookworm+ --compat-kernel is not supported.
log "Running reTerminal.sh --device reTerminal-DM from $REPO_DIR ..."
if ( cd "$REPO_DIR" && "$REPO_DIR/scripts/reTerminal.sh" --device reTerminal-DM ); then
  log "reTerminal DM drivers installed (reboot will apply)"
else
  log "WARNING: reTerminal.sh failed (see log above). Display/touch may still work; you can retry later with: cd $REPO_DIR && sudo ./scripts/reTerminal.sh --device reTerminal-DM"
fi
log "Removing $REPO_DIR"
rm -rf "$REPO_DIR"

# --- 6b. Re-apply splash and display (Seeed script sets disable_splash=1 and can duplicate Plymouth theme) ---
log "--- Re-applying boot splash and Plymouth theme ---"
CFG_PATH="/boot/firmware/config.txt"
[[ -f /boot/firmware/config.txt ]] || CFG_PATH="/boot/config.txt"
if [[ -f "$CFG_PATH" ]]; then
  if grep -q '^disable_splash=1' "$CFG_PATH"; then
    sed -i 's/^disable_splash=1$/disable_splash=0/' "$CFG_PATH"
    log "Set disable_splash=0 so Plymouth splash is shown"
  fi
fi
# Ensure Plymouth uses our custom theme only (single [Daemon], Theme=custom)
if [[ -f /etc/plymouth/plymouthd.conf ]]; then
  sed -i '/^Theme=/d' /etc/plymouth/plymouthd.conf
  sed -i '/^\[Daemon\]$/d' /etc/plymouth/plymouthd.conf
  grep -q '^\[Daemon\]' /etc/plymouth/plymouthd.conf || echo '[Daemon]' >> /etc/plymouth/plymouthd.conf
  echo 'Theme=custom' >> /etc/plymouth/plymouthd.conf
  log "Plymouth theme set to custom only"
fi
log "Running update-initramfs to apply Plymouth theme ..."
update-initramfs -u -k all 2>/dev/null || true

# --- 6b2. Kernel cmdline: swiotlb for vc4-drm (avoids "swiotlb buffer is full" / blank DSI on CM4) ---
CMDLINE_PATH="/boot/firmware/cmdline.txt"
[[ -f "$CMDLINE_PATH" ]] || CMDLINE_PATH="/boot/cmdline.txt"
if [[ -f "$CMDLINE_PATH" ]] && ! grep -q 'swiotlb=' "$CMDLINE_PATH"; then
  sed -i 's/rootwait/rootwait swiotlb=65536/' "$CMDLINE_PATH"
  log "Added swiotlb=65536 to kernel cmdline (vc4-drm / DSI)"
fi

# --- 6c. CM4: enable rpi-eeprom-update (flashrom + config.txt [cm4] block) so boot order can be set ---
# On CM4, rpi-eeprom-update is disabled by default. Enable it so we can set BOOT_ORDER=0x21 (network first).
# See: https://github.com/raspberrypi/usbboot , /etc/default/rpi-eeprom-update
log "--- CM4 EEPROM update enable (for boot order) ---"
EEPROM_DEFAULT="/etc/default/rpi-eeprom-update"
if [[ ! -f "$EEPROM_DEFAULT" ]]; then
  touch "$EEPROM_DEFAULT"
  log "Created $EEPROM_DEFAULT"
fi
grep -q '^RPI_EEPROM_USE_FLASHROM=' "$EEPROM_DEFAULT" && sed -i 's/^RPI_EEPROM_USE_FLASHROM=.*/RPI_EEPROM_USE_FLASHROM=1/' "$EEPROM_DEFAULT" || echo 'RPI_EEPROM_USE_FLASHROM=1' >> "$EEPROM_DEFAULT"
grep -q '^CM4_ENABLE_RPI_EEPROM_UPDATE=' "$EEPROM_DEFAULT" && sed -i 's/^CM4_ENABLE_RPI_EEPROM_UPDATE=.*/CM4_ENABLE_RPI_EEPROM_UPDATE=1/' "$EEPROM_DEFAULT" || echo 'CM4_ENABLE_RPI_EEPROM_UPDATE=1' >> "$EEPROM_DEFAULT"
log "Set RPI_EEPROM_USE_FLASHROM=1 and CM4_ENABLE_RPI_EEPROM_UPDATE=1 in $EEPROM_DEFAULT"
# config.txt: [cm4] section with SPI for EEPROM (only applied when booting a CM4)
if [[ -f "$CFG_PATH" ]]; then
  if ! grep -q '^\[cm4\]' "$CFG_PATH"; then
    printf '%s\n' '' '[cm4]' 'dtparam=spi=on' 'dtoverlay=audremap' 'dtoverlay=spi-gpio40-45' >> "$CFG_PATH"
    log "Added [cm4] block (spi, audremap, spi-gpio40-45) to config.txt for EEPROM"
  else
    for ENTRY in 'dtparam=spi=on' 'dtoverlay=audremap' 'dtoverlay=spi-gpio40-45'; do
      grep -q "^$ENTRY" "$CFG_PATH" || echo "$ENTRY" >> "$CFG_PATH"
    done
    log "Ensured [cm4] EEPROM entries in config.txt"
  fi
fi

# --- 6d. Boot order: network first, then eMMC/SD (for future network boot / re-provisioning) ---
# BOOT_ORDER: 0x2 = network, 0x1 = SD/eMMC. 0x21 = try network first, then local storage.
# On CM4, rpi-eeprom-update -l only works after reboot (once 6c is applied). So we try now; if it fails, a one-shot runs after next boot.
log "--- Boot order (network first, then eMMC/SD) ---"
BOOTCONF="/tmp/first-boot-eeprom-conf.txt"
BOOT_ORDER_SET=0
if command -v rpi-eeprom-config >/dev/null 2>&1 && command -v rpi-eeprom-update >/dev/null 2>&1; then
  if PEE="$(rpi-eeprom-update -l 2>/dev/null)" && [[ -n "$PEE" ]] && [[ -f "$PEE" ]]; then
    rpi-eeprom-config "$PEE" > "$BOOTCONF" 2>/dev/null || true
  fi
  if [[ -s "$BOOTCONF" ]]; then
    sed -i 's/^BOOT_ORDER=.*/BOOT_ORDER=0x21/' "$BOOTCONF"
    grep -q '^BOOT_ORDER=' "$BOOTCONF" || echo 'BOOT_ORDER=0x21' >> "$BOOTCONF"
    if rpi-eeprom-config --apply "$BOOTCONF" 2>/dev/null; then
      log "Boot order set to 0x21 (network first, then eMMC/SD); EEPROM update scheduled for next reboot"
      BOOT_ORDER_SET=1
    else
      log "WARNING: rpi-eeprom-config --apply failed; boot order unchanged"
    fi
  else
    log "rpi-eeprom-update -l did not return a config (on CM4 this is normal until after reboot with 6c applied); scheduling one-shot to set boot order after next boot"
  fi
  rm -f "$BOOTCONF"
else
  log "rpi-eeprom-config/rpi-eeprom-update not found; skipping boot order (not a Pi4/CM4 or package missing)"
fi

# If boot order was not set (e.g. CM4 first boot), install a one-shot systemd service to set it after reboot
if [[ "$BOOT_ORDER_SET" -eq 0 ]] && command -v rpi-eeprom-config >/dev/null 2>&1 && command -v rpi-eeprom-update >/dev/null 2>&1; then
  ONCE_SCRIPT="/usr/local/bin/set-cm4-boot-order-once.sh"
  ONCE_SVC="/etc/systemd/system/set-cm4-boot-order-once.service"
  cat > "$ONCE_SCRIPT" << 'SETBOOTEOF'
#!/bin/bash
# One-shot: set BOOT_ORDER=0x21 (network first) when rpi-eeprom-update becomes available (e.g. after CM4 enable and reboot).
BOOTCONF="/tmp/eeprom-boot-order-once.txt"
if PEE="$(rpi-eeprom-update -l 2>/dev/null)" && [[ -n "$PEE" ]] && [[ -f "$PEE" ]]; then
  rpi-eeprom-config "$PEE" > "$BOOTCONF" 2>/dev/null
  if [[ -s "$BOOTCONF" ]]; then
    sed -i 's/^BOOT_ORDER=.*/BOOT_ORDER=0x21/' "$BOOTCONF"
    grep -q '^BOOT_ORDER=' "$BOOTCONF" || echo 'BOOT_ORDER=0x21' >> "$BOOTCONF"
    if rpi-eeprom-config --apply "$BOOTCONF" 2>/dev/null; then
      echo "Boot order set to 0x21 (network first, then eMMC/SD)"
    fi
  fi
  rm -f "$BOOTCONF"
fi
systemctl disable set-cm4-boot-order-once.service 2>/dev/null || true
rm -f /etc/systemd/system/set-cm4-boot-order-once.service
rm -f "$0"
SETBOOTEOF
  chmod 755 "$ONCE_SCRIPT"
  cat > "$ONCE_SVC" << 'SVCEOF'
[Unit]
Description=Set CM4 boot order once (network first)
After=multi-user.target

[Service]
Type=oneshot
ExecStart=/usr/local/bin/set-cm4-boot-order-once.sh
RemainAfterExit=no

[Install]
WantedBy=multi-user.target
SVCEOF
  systemctl enable set-cm4-boot-order-once.service 2>/dev/null && log "Enabled set-cm4-boot-order-once.service to set boot order after next boot"
fi

# --- 7. One-shots (rotation at first login; wallpaper already set in pcmanfm config above) ---
log "--- One-shot scripts (run at pi first login) ---"
install_oneshot set-rotation-once || true
log "One-shots will append to $LOGFILE when they run at first login"

# --- 8. Allow pi to append to first-boot.log (for one-shot scripts) ---
chmod 666 "$LOGFILE"
log "Log file $LOGFILE is now appendable by user $PI_USER for one-shot scripts"

# --- 9. Reboot ---
log "=== first-boot.sh finished, rebooting ==="
reboot