#!/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.
# Optional: copy first-boot.conf.example to first-boot.conf and edit; it is loaded
# from the script dir, /tmp/first-boot.conf, or /etc/cm4-provisioning/first-boot.conf.
# Set ENABLE_STEP_01=0 .. ENABLE_STEP_13=0 in config to disable a step (default: all enabled).

set -e
export DEBIAN_FRONTEND=noninteractive

# --- Defaults (overridden by first-boot.conf if present) ---
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]:-.}")" 2>/dev/null && pwd)"
FILE_SERVER="${FILE_SERVER:-http://10.20.50.1:5000/files/first-boot}"
HOSTNAME="${HOSTNAME:-guard}"
PI_USER="${PI_USER:-pi}"
LOGFILE="${LOGFILE:-/var/log/first-boot.log}"
PLYMOUTH_DIR="${PLYMOUTH_DIR:-/usr/share/plymouth/themes/custom}"
WALLPAPER_PATH="${WALLPAPER_PATH:-/usr/share/rpd-wallpaper/splash.png}"
PACKAGES="${PACKAGES:-git chromium wmctrl openssh-server swaybg wlr-randr maliit-keyboard xinput-calibrator}"
LIGHTDM_SESSION="${LIGHTDM_SESSION:-rpd-labwc}"
GTK_THEME_NAME="${GTK_THEME_NAME:-PiXnoir}"
WALLPAPER_MODE="${WALLPAPER_MODE:-crop}"
DSI_ROTATE="${DSI_ROTATE:-90}"
SWIOTLB_SIZE="${SWIOTLB_SIZE:-65536}"
RETERMINAL_DEVICE="${RETERMINAL_DEVICE:-reTerminal-DM}"
RETERMINAL_REPO_URL="${RETERMINAL_REPO_URL:-https://github.com/Seeed-Studio/seeed-linux-dtoverlays}"
ONESHOT_SCRIPTS="${ONESHOT_SCRIPTS:-}"

# --- Load config file first (first found); then defaults apply only to unset vars ---
FIRST_BOOT_CONF=""
for _f in "$SCRIPT_DIR/first-boot.conf" /tmp/first-boot.conf /etc/cm4-provisioning/first-boot.conf; do
  if [[ -f "$_f" ]]; then
    # shellcheck source=first-boot.conf.example
    set -a && source "$_f" && set +a
    FIRST_BOOT_CONF="$_f"
    break
  fi
done

# Step enable flags (1 = run, 0 = skip). Default all enabled; only set if not already set by config.
ENABLE_STEP_01="${ENABLE_STEP_01:-1}"
ENABLE_STEP_02="${ENABLE_STEP_02:-1}"
ENABLE_STEP_03="${ENABLE_STEP_03:-1}"
ENABLE_STEP_04="${ENABLE_STEP_04:-1}"
ENABLE_STEP_05="${ENABLE_STEP_05:-1}"
ENABLE_STEP_06="${ENABLE_STEP_06:-1}"
ENABLE_STEP_07="${ENABLE_STEP_07:-1}"
ENABLE_STEP_08="${ENABLE_STEP_08:-1}"
ENABLE_STEP_09="${ENABLE_STEP_09:-1}"
ENABLE_STEP_10="${ENABLE_STEP_10:-1}"
ENABLE_STEP_11="${ENABLE_STEP_11:-1}"
ENABLE_STEP_12="${ENABLE_STEP_12:-1}"
ENABLE_STEP_13="${ENABLE_STEP_13:-1}"

# --- Derived paths ---
PI_HOME="/home/$PI_USER"
AUTOSTART="$PI_HOME/.config/autostart"

# Portal base URL for first-boot status API (derived from FILE_SERVER, e.g. http://10.20.50.1:5000)
if [[ "$FILE_SERVER" =~ ^(https?://[^/]+) ]]; then
  PORTAL_BASE="${BASH_REMATCH[1]}"
else
  PORTAL_BASE=""
fi

# --- Logging ---
log() { echo "[$(date -Iseconds)] $*"; }
exec > >(tee -a "$LOGFILE") 2>&1
log "=== first-boot.sh started ==="
[[ -n "$FIRST_BOOT_CONF" ]] && log "Config loaded from $FIRST_BOOT_CONF" || log "Using built-in defaults (no config file found)"
log "FILE_SERVER=$FILE_SERVER PI_USER=$PI_USER HOSTNAME=$HOSTNAME LOGFILE=$LOGFILE"

# --- Report status to portal (no-op if PORTAL_BASE empty; failures ignored) ---
report_status() {
  local phase="$1" message="$2" step="$3" step_name="$4" ip="$5"
  [[ -z "$PORTAL_BASE" ]] && return 0
  local json="{\"phase\":\"${phase}\",\"message\":\"${message}\",\"step\":\"${step}\",\"step_name\":\"${step_name}\",\"hostname\":\"${HOSTNAME}\",\"ip\":\"${ip}\"}"
  curl -s -X POST -H "Content-Type: application/json" -d "$json" "${PORTAL_BASE}/api/first-boot-status" || true
}
report_status "started" "First-boot started" "" "" ""

# --- Helper: run step if enabled (accepts "1" or "0"; strips CR/LF/whitespace) ---
run_step() {
  local n="$1" name="$2"
  local enable_var="ENABLE_STEP_${n}" val
  val="${!enable_var}"
  val="${val//[^01]/}"
  val="${val:0:1}"
  if [[ "$val" == "1" ]]; then
    log "--- Step $n: $name ---"
    "step_${n}_${name}" || return $?
    report_status "running" "Step $n: $name completed" "$n" "$name" ""
  else
    log "--- Step $n: $name (disabled) ---"
  fi
}

# --- Helper: install one-shot from FILE_SERVER ---
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)"
}

# --- Step 01: Hostname and /etc/hosts ---
step_01_hostname() {
  echo "$HOSTNAME" > /etc/hostname
  hostnamectl set-hostname "$HOSTNAME" 2>/dev/null || true
  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"
}

# --- Step 02: Packages ---
step_02_packages() {
  log "Running apt-get update ..."
  apt-get update -qq
  log "Installing: $PACKAGES"
  apt-get install -y -qq $PACKAGES
  log "Packages installed successfully"
}

# --- Step 03: Kiosk files from file server ---
step_03_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"
}

# --- Step 04: Boot splash and wallpaper ---
step_04_splash_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"
    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=$WALLPAPER_MODE" '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=$WALLPAPER_MODE/" "$PCMANFM_DESKTOP" || echo "wallpaper_mode=$WALLPAPER_MODE" >> "$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
}

# --- Step 05: LightDM session ---
step_05_lightdm() {
  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
  if [[ -f /etc/lightdm/lightdm.conf ]]; then
    sed -i "s/^user-session=.*/user-session=$LIGHTDM_SESSION/" /etc/lightdm/lightdm.conf
    sed -i "s/^autologin-session=.*/autologin-session=$LIGHTDM_SESSION/" /etc/lightdm/lightdm.conf
    log "Patched /etc/lightdm/lightdm.conf to use $LIGHTDM_SESSION"
  fi
}

# --- Step 06: Maliit on-screen keyboard ---
step_06_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"
}

# --- Step 07: Dark theme (GTK) ---
step_07_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=$GTK_THEME_NAME" > "$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=$GTK_THEME_NAME/" "$GTK_SETTINGS" || echo "gtk-theme-name=$GTK_THEME_NAME" >> "$GTK_SETTINGS"
  fi
  log "Set dark theme ($GTK_THEME_NAME) in gtk-3.0/settings.ini"
  chown -R "$PI_USER:$PI_USER" "$PI_HOME/.config"
}

# --- Step 08: reTerminal DM drivers (Seeed) ---
step_08_reterminal_drivers() {
  if [[ -z "$RETERMINAL_REPO_URL" ]]; then
    log "Skipping reTerminal drivers (RETERMINAL_REPO_URL not set)"
    return 0
  fi
  REPO_DIR="/tmp/seeed-linux-dtoverlays"
  log "Cloning seeed-linux-dtoverlays to $REPO_DIR ..."
  git clone --depth 1 "$RETERMINAL_REPO_URL" "$REPO_DIR"
  log "Running reTerminal.sh --device $RETERMINAL_DEVICE from $REPO_DIR ..."
  if ( cd "$REPO_DIR" && "$REPO_DIR/scripts/reTerminal.sh" --device "$RETERMINAL_DEVICE" ); 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_DEVICE"
  fi
  rm -rf "$REPO_DIR"
}

# --- Step 09: Re-apply splash and Plymouth theme ---
step_09_reapply_splash() {
  CFG_PATH="/boot/firmware/config.txt"
  [[ -f /boot/firmware/config.txt ]] || CFG_PATH="/boot/config.txt"
  if [[ -f "$CFG_PATH" ]] && 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
  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
}

# --- Step 10: Kernel cmdline (swiotlb + DSI rotation) ---
step_10_cmdline() {
  CMDLINE_PATH="/boot/firmware/cmdline.txt"
  [[ -f "$CMDLINE_PATH" ]] || CMDLINE_PATH="/boot/cmdline.txt"
  if [[ -f "$CMDLINE_PATH" ]]; then
    if [[ -n "$SWIOTLB_SIZE" ]] && ! grep -q 'swiotlb=' "$CMDLINE_PATH"; then
      sed -i "s/rootwait/rootwait swiotlb=$SWIOTLB_SIZE/" "$CMDLINE_PATH"
      log "Added swiotlb=$SWIOTLB_SIZE to kernel cmdline (vc4-drm / DSI)"
    fi
    if [[ -n "$DSI_ROTATE" ]] && ! grep -q 'video=DSI-1:rotate=' "$CMDLINE_PATH"; then
      sed -i "s/\$/ video=DSI-1:rotate=$DSI_ROTATE/" "$CMDLINE_PATH"
      log "Added video=DSI-1:rotate=$DSI_ROTATE to kernel cmdline (DSI rotation)"
    fi
  fi
}

# --- Step 11: One-shot scripts ---
step_11_oneshots() {
  if [[ -n "$DSI_ROTATE" ]]; then
    log "Rotation is set via kernel cmdline (video=DSI-1:rotate=$DSI_ROTATE)"
  fi
  if [[ -n "$ONESHOT_SCRIPTS" ]]; then
    for _name in $ONESHOT_SCRIPTS; do
      install_oneshot "$_name" || true
    done
  else
    log "No one-shot scripts configured (ONESHOT_SCRIPTS empty)"
  fi
}

# --- Step 12: Log file permissions ---
step_12_log_permissions() {
  chmod 666 "$LOGFILE"
  log "Log file $LOGFILE is now appendable by user $PI_USER for one-shot scripts"
}

# --- Step 13: Reboot ---
step_13_reboot() {
  DEVICE_IP="$(hostname -I 2>/dev/null | awk '{print $1}')"
  report_status "done" "First-boot complete" "13" "reboot" "$DEVICE_IP"
  log "Device IP: ${DEVICE_IP:-unknown}"
  log "=== first-boot.sh finished, rebooting ==="
  reboot
}

# --- Main: run steps in order ---
run_step 01 hostname
run_step 02 packages
run_step 03 kiosk_files
run_step 04 splash_wallpaper
run_step 05 lightdm
run_step 06 maliit
run_step 07 dark_theme
run_step 08 reterminal_drivers
run_step 09 reapply_splash
run_step 10 cmdline
run_step 11 oneshots
run_step 12 log_permissions
run_step 13 reboot

# If reboot was disabled, still report done and device IP so the portal shows completion
_step13_val="${ENABLE_STEP_13:-1}"
_step13_val="${_step13_val//[^01]/}"
_step13_val="${_step13_val:0:1}"
if [[ "$_step13_val" != "1" ]]; then
  DEVICE_IP="$(hostname -I 2>/dev/null | awk '{print $1}')"
  report_status "done" "First-boot complete (reboot disabled)" "13" "reboot" "$DEVICE_IP"
  log "Device IP: ${DEVICE_IP:-unknown}"
  log "=== first-boot.sh finished (reboot disabled) ==="
fi