#!/usr/bin/env bash
# Install as /var/lib/dietpi/postboot.d/10-pikit-firstboot and chmod +x.
# Runs once on first boot to finalize device-unique setup.

set -euo pipefail

FIRSTBOOT_DIR="/var/lib/pikit/firstboot"
STATE_FILE="$FIRSTBOOT_DIR/state.json"
LOG_FILE="$FIRSTBOOT_DIR/firstboot.log"
ERROR_FILE="$FIRSTBOOT_DIR/firstboot.error"
DONE_FILE="$FIRSTBOOT_DIR/firstboot.done"
LOCK_FILE="$FIRSTBOOT_DIR/firstboot.lock"
CERT_DIR="/etc/pikit/certs"
WEB_ASSETS="/var/www/pikit-web/assets"
PROFILE_FILE="/etc/pikit/profile.json"
MOTD_FILE="/etc/motd"
FIRSTBOOT_CONF="/etc/pikit/firstboot.conf"

STEPS=(
  "Preparing system"
  "Generating security keys"
  "Securing the dashboard"
  "Updating software (this can take a while)"
  "Final checks"
  "Starting Pi-Kit"
)
STEP_STATUS=(pending pending pending pending pending pending)
CURRENT_STEP=""
CURRENT_INDEX=-1
PIKIT_FIRSTBOOT_UPDATES="${PIKIT_FIRSTBOOT_UPDATES:-1}"

log() {
  printf '[%s] %s\n' "$(date '+%Y-%m-%dT%H:%M:%S%z')" "$*"
}

load_config() {
  if [ -f "$FIRSTBOOT_CONF" ]; then
    # shellcheck disable=SC1090
    . "$FIRSTBOOT_CONF"
  fi
  PIKIT_FIRSTBOOT_UPDATES="${PIKIT_FIRSTBOOT_UPDATES:-1}"
}

skip_updates() {
  case "${PIKIT_FIRSTBOOT_UPDATES,,}" in
    0|false|no|off) return 0 ;;
  esac
  return 1
}

write_state() {
  local state="$1"
  local current="$2"
  local steps_joined
  local status_joined
  steps_joined=$(IFS='|'; echo "${STEPS[*]}")
  status_joined=$(IFS='|'; echo "${STEP_STATUS[*]}")
  PIKIT_STATE_FILE="$STATE_FILE" PIKIT_STATE="$state" PIKIT_CURRENT_STEP="$current" PIKIT_STEPS="$steps_joined" PIKIT_STEP_STATUSES="$status_joined" \
    python3 - <<'PY'
import json
import os
import pathlib
from datetime import datetime, timezone

state_path = pathlib.Path(os.environ["PIKIT_STATE_FILE"]) if "PIKIT_STATE_FILE" in os.environ else pathlib.Path("/var/lib/pikit/firstboot/state.json")
state = os.environ.get("PIKIT_STATE", "running")
current = os.environ.get("PIKIT_CURRENT_STEP") or None
steps = (os.environ.get("PIKIT_STEPS") or "").split("|")
statuses = (os.environ.get("PIKIT_STEP_STATUSES") or "").split("|")
if len(statuses) < len(steps):
    statuses += ["pending"] * (len(steps) - len(statuses))
updated_at = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
state_path.parent.mkdir(parents=True, exist_ok=True)
state_path.write_text(json.dumps({
    "state": state,
    "current_step": current,
    "steps": [{"label": label, "status": status} for label, status in zip(steps, statuses)],
    "updated_at": updated_at,
}, indent=2))
PY
}

begin_step() {
  local idx="$1"
  if [ "$CURRENT_INDEX" -ge 0 ]; then
    STEP_STATUS[$CURRENT_INDEX]="done"
  fi
  CURRENT_INDEX="$idx"
  CURRENT_STEP="${STEPS[$idx]}"
  STEP_STATUS[$idx]="current"
  write_state "running" "$CURRENT_STEP"
  log "--- $CURRENT_STEP ---"
}

finish_step() {
  local idx="$1"
  local state="${2:-running}"
  local current="${3:-$CURRENT_STEP}"
  STEP_STATUS[$idx]="done"
  write_state "$state" "$current"
}

clear_motd_block() {
  if [ -f "$MOTD_FILE" ]; then
    sed -i '/^\[Pi-Kit firstboot\]/,/^\[\/Pi-Kit firstboot\]/d' "$MOTD_FILE" || true
  fi
}

write_motd_error() {
  clear_motd_block
  cat >> "$MOTD_FILE" <<'TXT'
[Pi-Kit firstboot]
Pi-Kit setup needs attention.
Error log: sudo cat /var/lib/pikit/firstboot/firstboot.error
Full log:  sudo cat /var/lib/pikit/firstboot/firstboot.log
If needed: sudo systemctl restart nginx pikit-api
[/Pi-Kit firstboot]
TXT
}

handle_error() {
  local line="$1"
  local msg="Firstboot failed at step: ${CURRENT_STEP:-unknown} (line $line)"
  log "$msg"
  STEP_STATUS[$CURRENT_INDEX]="error"
  write_state "error" "${CURRENT_STEP:-}" || true
  echo "$msg" > "$ERROR_FILE"
  if [ -f "$LOG_FILE" ]; then
    echo "--- recent log ---" >> "$ERROR_FILE"
    tail -n 120 "$LOG_FILE" >> "$ERROR_FILE"
  fi
  write_motd_error
  exit 1
}

mkdir -p "$FIRSTBOOT_DIR"
:> "$LOG_FILE"
exec >>"$LOG_FILE" 2>&1

log "Pi-Kit firstboot starting"
load_config

if [ -f "$DONE_FILE" ]; then
  log "Firstboot already completed; exiting."
  exit 0
fi

rm -f "$ERROR_FILE"
clear_motd_block

exec 9>"$LOCK_FILE"
if ! flock -n 9; then
  log "Another firstboot run is in progress; exiting."
  exit 0
fi

trap 'handle_error $LINENO' ERR

begin_step 0
mkdir -p "$CERT_DIR" "$WEB_ASSETS"
if getent group pikit-cert >/dev/null 2>&1; then
  chgrp pikit-cert "$CERT_DIR" || true
fi
chmod 750 "$CERT_DIR"
finish_step 0

begin_step 1
if [ -x /usr/local/bin/pikit-certgen.sh ]; then
  /usr/local/bin/pikit-certgen.sh
else
  if [ -s "$CERT_DIR/pikit-ca.crt" ] && [ -s "$CERT_DIR/pikit-ca.key" ] && [ -s "$CERT_DIR/pikit.local.crt" ] && [ -s "$CERT_DIR/pikit.local.key" ]; then
    log "TLS certs already present; skipping generation."
  else
    if ! command -v openssl >/dev/null 2>&1; then
      echo "openssl not installed" >&2
      exit 1
    fi
    rm -f "$CERT_DIR/pikit-ca.key" "$CERT_DIR/pikit-ca.crt" "$CERT_DIR/pikit-ca.srl" || true
    rm -f "$CERT_DIR/pikit.local.key" "$CERT_DIR/pikit.local.crt" "$CERT_DIR/pikit.local.csr" || true

    openssl genrsa -out "$CERT_DIR/pikit-ca.key" 2048
    openssl req -x509 -new -nodes -key "$CERT_DIR/pikit-ca.key" -sha256 -days 3650 \
      -out "$CERT_DIR/pikit-ca.crt" -subj "/CN=Pi-Kit CA"

    openssl genrsa -out "$CERT_DIR/pikit.local.key" 2048
    openssl req -new -key "$CERT_DIR/pikit.local.key" -out "$CERT_DIR/pikit.local.csr" -subj "/CN=pikit.local"

    SAN_CFG=$(mktemp)
    cat > "$SAN_CFG" <<'CFG'
authorityKeyIdentifier=keyid,issuer
basicConstraints=CA:FALSE
keyUsage = digitalSignature, nonRepudiation, keyEncipherment, dataEncipherment
subjectAltName = @alt_names

[alt_names]
DNS.1 = pikit.local
DNS.2 = pikit
CFG

    openssl x509 -req -in "$CERT_DIR/pikit.local.csr" -CA "$CERT_DIR/pikit-ca.crt" -CAkey "$CERT_DIR/pikit-ca.key" \
      -CAcreateserial -out "$CERT_DIR/pikit.local.crt" -days 825 -sha256 -extfile "$SAN_CFG"
    rm -f "$SAN_CFG" "$CERT_DIR/pikit.local.csr"

    chmod 600 "$CERT_DIR/pikit-ca.key" "$CERT_DIR/pikit.local.key"
    chmod 644 "$CERT_DIR/pikit-ca.crt" "$CERT_DIR/pikit.local.crt"
  fi
fi
finish_step 1

begin_step 2
cp "$CERT_DIR/pikit-ca.crt" "$WEB_ASSETS/pikit-ca.crt"
chmod 644 "$WEB_ASSETS/pikit-ca.crt"
if command -v sha256sum >/dev/null 2>&1; then
  sha256sum "$WEB_ASSETS/pikit-ca.crt" | awk '{print $1}' > "$WEB_ASSETS/pikit-ca.sha256"
elif command -v openssl >/dev/null 2>&1; then
  openssl dgst -sha256 "$WEB_ASSETS/pikit-ca.crt" | awk '{print $2}' > "$WEB_ASSETS/pikit-ca.sha256"
fi
if [ -s "$WEB_ASSETS/pikit-ca.sha256" ]; then
  chmod 644 "$WEB_ASSETS/pikit-ca.sha256"
fi
if command -v systemctl >/dev/null 2>&1; then
  systemctl reload nginx || systemctl restart nginx
fi
finish_step 2

begin_step 3
if skip_updates; then
  log "Skipping software updates (PIKIT_FIRSTBOOT_UPDATES=$PIKIT_FIRSTBOOT_UPDATES)."
else
  export DEBIAN_FRONTEND=noninteractive
  mkdir -p /var/cache/apt/archives/partial /var/lib/apt/lists/partial
  chmod 755 /var/cache/apt/archives /var/cache/apt/archives/partial /var/lib/apt/lists /var/lib/apt/lists/partial
  apt-get update
  apt-get -y -o Dpkg::Options::=--force-confold full-upgrade
fi
finish_step 3

begin_step 4
if [ -f "$PROFILE_FILE" ] && command -v ufw >/dev/null 2>&1; then
  python3 - <<'PY' > /tmp/pikit-profile-ports.txt
import json
import sys
from pathlib import Path
profile = Path("/etc/pikit/profile.json")
try:
    data = json.loads(profile.read_text())
except Exception:
    data = {}
ports = data.get("firewall_ports") or []
for port in ports:
    try:
        port_int = int(port)
    except Exception:
        continue
    print(port_int)
PY
  while read -r port; do
    [ -z "$port" ] && continue
    for subnet in 10.0.0.0/8 172.16.0.0/12 192.168.0.0/16 100.64.0.0/10 169.254.0.0/16; do
      ufw allow from "$subnet" to any port "$port" || true
    done
  done < /tmp/pikit-profile-ports.txt
  rm -f /tmp/pikit-profile-ports.txt
else
  log "Profile firewall step skipped (no profile.json or ufw missing)"
fi

if [ ! -f "$WEB_ASSETS/pikit-ca.crt" ]; then
  echo "CA bundle missing in web assets" >&2
  exit 1
fi
finish_step 4

begin_step 5
touch "$DONE_FILE"
touch /var/run/pikit-ready
finish_step 5 "done" "${STEPS[5]}"

log "Pi-Kit firstboot complete"
exit 0