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

# ============================================
# PeopleCounter - Auto install (Mender + Connect)
# For Raspberry Pi 4 / Raspberry Pi 5
# ============================================

# ----------------------------
# USER CONFIG (edit only if needed)
# ----------------------------
SERVER_URL="${SERVER_URL:-https://visionai.bbctech.it/}"
SERVER_CERT_PATH="${SERVER_CERT_PATH:-}"   # leave empty for public CA (Let's Encrypt etc.)

# Where Mender reads device type (Debian package installs)
DEVICE_TYPE_FILE="${DEVICE_TYPE_FILE:-/var/lib/mender/device_type}"

# VisionAI Wi-Fi provisioning
VISIONAI_PROVISIONING_ENV="${VISIONAI_PROVISIONING_ENV:-/data/visionai-provisioning.env}"
WIFI_CONNECT_DEFAULTS_FILE="${WIFI_CONNECT_DEFAULTS_FILE:-/etc/default/wifi-connect-manager}"
WIFI_CONNECT_LAUNCHER_PATH="${WIFI_CONNECT_LAUNCHER_PATH:-/usr/local/bin/wifi-connect-launcher}"
VISIONAI_INVENTORY_SCRIPT="${VISIONAI_INVENTORY_SCRIPT:-/usr/share/mender/inventory/mender-inventory-visionai-provisioning}"
VISIONAI_WIFI_QR_PNG="${VISIONAI_WIFI_QR_PNG:-/data/visionai-wifi-qr.png}"
VISIONAI_PORTAL_GATEWAY="${VISIONAI_PORTAL_GATEWAY:-192.168.42.1}"
VISIONAI_PORTAL_INTERFACE="${VISIONAI_PORTAL_INTERFACE:-wlan0}"
VISIONAI_UI_PORT="${VISIONAI_UI_PORT:-2001}"

# ----------------------------
# Helpers
# ----------------------------
log() {
  echo
  echo ">>> $*"
}

warn() {
  echo
  echo "WARNING: $*" >&2
}

die() {
  echo
  echo "ERROR: $*" >&2
  exit 1
}

need_root() {
  if [[ "${EUID:-$(id -u)}" -ne 0 ]]; then
    die "Run this script with sudo (example: curl ... | sudo bash)"
  fi
}

detect_pi_model() {
  local model=""
  model="$(tr -d '\0' </proc/device-tree/model 2>/dev/null || true)"

  if [[ -z "$model" && -r /sys/firmware/devicetree/base/model ]]; then
    model="$(tr -d '\0' </sys/firmware/devicetree/base/model 2>/dev/null || true)"
  fi

  echo "$model"
}

detect_device_type() {
  local model="$1"

  case "$model" in
    *"Raspberry Pi 4"*)
      echo "raspberrypi4"
      ;;
    *"Raspberry Pi 5"*)
      echo "raspberrypi5"
      ;;
    *)
      return 1
      ;;
  esac
}

json_escape() {
  # Minimal JSON string escaper (safe for paths/URLs)
  local s="$1"
  s="${s//\\/\\\\}"
  s="${s//\"/\\\"}"
  printf '%s' "$s"
}

env_double_quote() {
  local s="$1"
  s="${s//\\/\\\\}"
  s="${s//\"/\\\"}"
  s="${s//\$/\\\$}"
  s="${s//\`/\\\`}"
  printf '"%s"' "$s"
}

random_lower_alnum() {
  local length="$1"
  local out="" chunk=""

  [[ "$length" =~ ^[0-9]+$ && "$length" -gt 0 ]] || return 1

  while [[ "${#out}" -lt "$length" ]]; do
    chunk="$(
      LC_ALL=C tr -dc 'abcdefghijklmnopqrstuvwxyz0123456789' </dev/urandom \
        | head -c "$((length - ${#out}))" \
        || true
    )"
    out+="$chunk"
  done

  printf '%s\n' "$out"
}

valid_ssid() {
  [[ "${1:-}" =~ ^VisionAI-[a-z0-9]{4}$ ]]
}

valid_password() {
  [[ "${1:-}" =~ ^[a-z0-9]{8}$ ]]
}

read_env_value() {
  local file="$1"
  local key="$2"
  local value=""

  [[ -r "$file" ]] || return 1
  value="$(sed -n "s/^[[:space:]]*${key}=//p" "$file" 2>/dev/null | tail -n 1 || true)"
  [[ -n "$value" ]] || return 1

  value="${value%$'\r'}"
  if [[ "${value:0:1}" == '"' && "${value: -1}" == '"' ]]; then
    value="${value:1:${#value}-2}"
  elif [[ "${value:0:1}" == "'" && "${value: -1}" == "'" ]]; then
    value="${value:1:${#value}-2}"
  fi

  printf '%s\n' "$value"
}

write_visionai_provisioning_env() {
  local ssid="$1"
  local password="$2"
  local file="${VISIONAI_PROVISIONING_ENV}"
  local tmp=""

  install -d -m 0755 "$(dirname "$file")"
  tmp="$(mktemp "${file}.tmp.XXXXXX")"
  chmod 0600 "$tmp"

  {
    printf 'PORTAL_SSID=%s\n' "$(env_double_quote "$ssid")"
    printf 'PORTAL_PASSPHRASE=%s\n' "$(env_double_quote "$password")"
  } > "$tmp"

  chmod 0600 "$tmp"
  chown root:root "$tmp" 2>/dev/null || true
  mv "$tmp" "$file"
  chmod 0600 "$file"
}

setup_visionai_wifi_credentials() {
  local existing_ssid="" existing_password=""
  local ssid="" password=""

  existing_ssid="$(read_env_value "$VISIONAI_PROVISIONING_ENV" PORTAL_SSID || true)"
  existing_password="$(read_env_value "$VISIONAI_PROVISIONING_ENV" PORTAL_PASSPHRASE || true)"

  if valid_ssid "$existing_ssid"; then
    ssid="$existing_ssid"
  else
    ssid="VisionAI-$(random_lower_alnum 4)"
  fi

  if valid_password "$existing_password"; then
    password="$existing_password"
  else
    password="$(random_lower_alnum 8)"
  fi

  write_visionai_provisioning_env "$ssid" "$password"

  PORTAL_SSID="$ssid"
  PORTAL_PASSPHRASE="$password"
  export PORTAL_SSID PORTAL_PASSPHRASE
}

append_default_if_missing() {
  local file="$1"
  local key="$2"
  local value="$3"

  if ! grep -Eq "^[[:space:]]*(export[[:space:]]+)?${key}=" "$file" 2>/dev/null; then
    printf '%s=%s\n' "$key" "$value" >> "$file"
  fi
}

update_wifi_connect_manager_defaults() {
  local file="${WIFI_CONNECT_DEFAULTS_FILE}"
  local tmp=""
  local ssid="${PORTAL_SSID:-}"
  local password="${PORTAL_PASSPHRASE:-}"

  valid_ssid "$ssid" || die "Internal error: invalid generated PORTAL_SSID"
  valid_password "$password" || die "Internal error: invalid generated PORTAL_PASSPHRASE"

  install -d -m 0755 "$(dirname "$file")"
  tmp="$(mktemp "${file}.tmp.XXXXXX")"

  if [[ -f "$file" ]]; then
    awk '
      !/^[[:space:]]*(export[[:space:]]+)?PORTAL_SSID=/ &&
      !/^[[:space:]]*(export[[:space:]]+)?PORTAL_PASSPHRASE=/ { print }
    ' "$file" > "$tmp"
  else
    : > "$tmp"
  fi

  append_default_if_missing "$tmp" WIFI_CONNECT_BIN "/usr/local/bin/wifi-connect-launcher"
  append_default_if_missing "$tmp" CHECK_INTERVAL "5"
  append_default_if_missing "$tmp" BOOT_PORTAL_WINDOW "180"
  append_default_if_missing "$tmp" PORTAL_READY_TIMEOUT "15"
  append_default_if_missing "$tmp" PORTAL_RESTART_GRACE "20"
  append_default_if_missing "$tmp" PORTAL_INTERFACE "$VISIONAI_PORTAL_INTERFACE"
  append_default_if_missing "$tmp" PORTAL_GATEWAY "$VISIONAI_PORTAL_GATEWAY"
  append_default_if_missing "$tmp" PORTAL_DHCP_RANGE "192.168.42.2,192.168.42.254"
  append_default_if_missing "$tmp" PORTAL_LISTENING_PORT "80"
  append_default_if_missing "$tmp" WIFI_COUNTRY "RO"
  append_default_if_missing "$tmp" WIFI_CONNECT_ARGS '""'
  append_default_if_missing "$tmp" RESET_FLAG "/data/wifi-reset.flag"
  append_default_if_missing "$tmp" PID_FILE "/run/wifi-connect-manager/wifi-connect.pid"
  append_default_if_missing "$tmp" NM_DNS_ADJUST_FLAG "/run/wifi-connect-manager/nm-dns-adjusted.flag"

  printf 'PORTAL_SSID=%s\n' "$(env_double_quote "$ssid")" >> "$tmp"
  printf 'PORTAL_PASSPHRASE=%s\n' "$(env_double_quote "$password")" >> "$tmp"

  chown root:root "$tmp" 2>/dev/null || true
  chmod 0600 "$tmp"
  mv "$tmp" "$file"
  chmod 0600 "$file"
}

wifi_connect_launcher_is_safe() {
  local launcher="${1:-$WIFI_CONNECT_LAUNCHER_PATH}"
  local old_shared_ssid="VisionAI"" Connect"
  local old_shared_pass="mender""06"

  [[ -r "$launcher" ]] || return 1
  grep -q '/etc/default/wifi-connect-manager' "$launcher" || return 1
  grep -q -- '--portal-ssid' "$launcher" || return 1
  grep -Eq -- '--portal-passphrase|export[[:space:]]+PORTAL_PASSPHRASE' "$launcher" || return 1
  ! grep -Eq "${old_shared_ssid}|${old_shared_pass}" "$launcher"
}

install_wifi_connect_launcher_wrapper() {
  local launcher="${WIFI_CONNECT_LAUNCHER_PATH}"

  install -d -m 0755 "$(dirname "$launcher")"
  cat > "$launcher" <<'LAUNCHER'
#!/usr/bin/env bash
set -euo pipefail

ENV=/etc/default/wifi-connect-manager
if [[ -f "$ENV" ]]; then
  # shellcheck disable=SC1091
  source "$ENV"
fi

: "${PORTAL_SSID:?PORTAL_SSID is required in /etc/default/wifi-connect-manager}"
: "${PORTAL_PASSPHRASE:?PORTAL_PASSPHRASE is required in /etc/default/wifi-connect-manager}"

export PORTAL_INTERFACE="${PORTAL_INTERFACE:-wlan0}"
export PORTAL_SSID
export PORTAL_PASSPHRASE
export PORTAL_GATEWAY="${PORTAL_GATEWAY:-192.168.42.1}"
export PORTAL_DHCP_RANGE="${PORTAL_DHCP_RANGE:-192.168.42.2,192.168.42.254}"
export PORTAL_LISTENING_PORT="${PORTAL_LISTENING_PORT:-80}"
export WIFI_COUNTRY="${WIFI_COUNTRY:-RO}"

/usr/local/bin/wifi-prep.sh || true

ARGS=(
  --portal-interface "$PORTAL_INTERFACE"
  --portal-ssid "$PORTAL_SSID"
  --portal-gateway "$PORTAL_GATEWAY"
  --portal-dhcp-range "$PORTAL_DHCP_RANGE"
  --portal-listening-port "$PORTAL_LISTENING_PORT"
)

if [[ -n "${WIFI_CONNECT_ARGS:-}" ]]; then
  # shellcheck disable=SC2206
  EXTRA_ARGS=($WIFI_CONNECT_ARGS)
  ARGS+=("${EXTRA_ARGS[@]}")
fi

exec /usr/local/bin/wifi-connect "${ARGS[@]}" "$@"
LAUNCHER
  chmod 0755 "$launcher"
  chown root:root "$launcher" 2>/dev/null || true
}

ensure_wifi_connect_launcher_uses_defaults() {
  if wifi_connect_launcher_is_safe "$WIFI_CONNECT_LAUNCHER_PATH"; then
    return 0
  fi

  log "Installing wifi-connect launcher wrapper with generated credential support"
  install_wifi_connect_launcher_wrapper
}

install_wifi_connect_runtime_from_repo_if_available() {
  local script_dir="" deploy_dir="" arch="" wifi_binary=""

  script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd || true)"
  [[ -n "$script_dir" ]] || return 0
  deploy_dir="$script_dir/deploy"
  [[ -s "$deploy_dir/bin/wifi-connect-manager.sh" ]] || return 0

  log "Installing wifi-connect runtime files from repository deploy assets"
  install -d -m 0755 /usr/local/bin /etc/systemd/system
  install -m 0755 "$deploy_dir/bin/wifi-connect-manager.sh" /usr/local/bin/wifi-connect-manager.sh
  install -m 0755 "$deploy_dir/bin/wifi-prep.sh" /usr/local/bin/wifi-prep.sh
  install -m 0755 "$deploy_dir/bin/wifi-connect-launcher" /usr/local/bin/wifi-connect-launcher
  install -m 0755 "$deploy_dir/bin/wifi-reset.sh" /usr/local/bin/wifi-reset.sh
  install -m 0644 "$deploy_dir/systemd/wifi-connect-manager.service" /etc/systemd/system/wifi-connect-manager.service

  arch="$(uname -m 2>/dev/null || true)"
  case "$arch" in
    aarch64|arm64) wifi_binary="$deploy_dir/bin/wifi-connect-aarch" ;;
    armv7l|armv7*|armhf) wifi_binary="$deploy_dir/bin/wifi-connect-armv7" ;;
    *) wifi_binary="$deploy_dir/bin/wifi-connect" ;;
  esac
  if [[ -s "$wifi_binary" ]]; then
    install -m 0755 "$wifi_binary" /usr/local/bin/wifi-connect
  elif [[ -s "$deploy_dir/bin/wifi-connect" ]]; then
    install -m 0755 "$deploy_dir/bin/wifi-connect" /usr/local/bin/wifi-connect
  else
    warn "wifi-connect binary not found in repository deploy assets"
  fi

  sed -i 's/\r$//' \
    /usr/local/bin/wifi-connect-manager.sh \
    /usr/local/bin/wifi-prep.sh \
    /usr/local/bin/wifi-connect-launcher \
    /usr/local/bin/wifi-reset.sh \
    /etc/systemd/system/wifi-connect-manager.service 2>/dev/null || true

  if command -v systemctl >/dev/null 2>&1; then
    systemctl daemon-reload || true
    systemctl enable wifi-connect-manager.service || warn "Could not enable wifi-connect-manager.service"
  fi
}

install_wifi_connect_ui_if_possible() {
  local version="${WIFI_CONNECT_VERSION:-v4.11.84}"
  local tgz="/tmp/wifi-connect-ui-${version}.tar.gz"
  local ui_root="/opt/wifi-connect/ui"
  local ui_dir=""

  if find "$ui_root" -maxdepth 8 -type f -iname index.html 2>/dev/null | grep -q .; then
    return 0
  fi
  command -v curl >/dev/null 2>&1 || return 0
  command -v tar >/dev/null 2>&1 || return 0

  log "Installing optional wifi-connect captive portal UI"
  install -d -m 0755 "$ui_root"
  if ! curl -fL --retry 3 --retry-delay 2 \
    -o "$tgz" \
    "https://github.com/balena-os/wifi-connect/releases/download/${version}/wifi-connect-ui.tar.gz"; then
    warn "Could not download wifi-connect UI; captive portal may return 404 until UI assets are installed"
    return 0
  fi

  find "$ui_root" -mindepth 1 -maxdepth 1 -exec rm -rf {} + 2>/dev/null || true
  tar -xzf "$tgz" -C "$ui_root" || {
    warn "Could not extract wifi-connect UI"
    return 0
  }
  ui_dir="$(dirname "$(find "$ui_root" -maxdepth 8 -type f -iname index.html 2>/dev/null | head -n 1)")"
  [[ -n "$ui_dir" ]] || {
    warn "wifi-connect UI archive did not contain index.html"
    return 0
  }

  if grep -q '^WIFI_CONNECT_ARGS=' "$WIFI_CONNECT_DEFAULTS_FILE" 2>/dev/null; then
    sed -i "s|^WIFI_CONNECT_ARGS=.*|WIFI_CONNECT_ARGS=\"--ui-directory $ui_dir\"|g" "$WIFI_CONNECT_DEFAULTS_FILE"
  else
    printf 'WIFI_CONNECT_ARGS="--ui-directory %s"\n' "$ui_dir" >> "$WIFI_CONNECT_DEFAULTS_FILE"
  fi
}

restart_wifi_connect_manager_if_present() {
  if command -v systemctl >/dev/null 2>&1 \
    && systemctl list-unit-files 2>/dev/null | grep -q '^wifi-connect-manager\.service'; then
    systemctl daemon-reload || true
    systemctl restart wifi-connect-manager.service || warn "Could not restart wifi-connect-manager.service"
  fi
}

install_visionai_inventory_script() {
  local script="${VISIONAI_INVENTORY_SCRIPT}"

  install -d -m 0755 "$(dirname "$script")"
  cat > "$script" <<'INVENTORY'
#!/bin/sh

ENV_FILE="${VISIONAI_PROVISIONING_ENV:-/data/visionai-provisioning.env}"
PORTAL_GATEWAY="${VISIONAI_PORTAL_GATEWAY:-192.168.42.1}"
PORTAL_INTERFACE="${VISIONAI_PORTAL_INTERFACE:-wlan0}"
VISIONAI_UI_PORT="${VISIONAI_UI_PORT:-2001}"

read_env_key() {
  key="$1"
  [ -r "$ENV_FILE" ] || return 1

  value="$(
    awk -v key="$key" '
      BEGIN { prefix = key "=" }
      index($0, prefix) == 1 { value = substr($0, length(prefix) + 1) }
      END { if (value != "") print value }
    ' "$ENV_FILE" 2>/dev/null || true
  )"
  [ -n "$value" ] || return 1

  value="$(printf '%s' "$value" | tr -d '\r')"
  case "$value" in
    \"*\") value=${value#\"}; value=${value%\"} ;;
  esac
  case "$value" in
    \'*\') value=${value#\'}; value=${value%\'} ;;
  esac

  printf '%s' "$value" | tr -d '\n'
}

valid_ssid() {
  value="${1:-}"
  case "$value" in
    VisionAI-????)
      suffix=${value#VisionAI-}
      case "$suffix" in
        *[!abcdefghijklmnopqrstuvwxyz0123456789]*) return 1 ;;
      esac
      [ "${#suffix}" -eq 4 ]
      ;;
    *) return 1 ;;
  esac
}

valid_password() {
  value="${1:-}"
  case "$value" in
    ????????)
      case "$value" in
        *[!abcdefghijklmnopqrstuvwxyz0123456789]*) return 1 ;;
      esac
      ;;
    *) return 1 ;;
  esac
}

safe_value() {
  value="${1:-}"
  value="$(printf '%s' "$value" | tr -d '\r\n')"
  [ -n "$value" ] || value="unknown"
  printf '%s' "$value"
}

print_kv() {
  key="${1:-}"
  value="$(safe_value "${2:-}")"

  case "$key" in
    ''|*[!abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_-]*) return 0 ;;
  esac

  printf '%s=%s\n' "$key" "$value"
}

get_mac() {
  iface="$1"
  path="/sys/class/net/$iface/address"
  mac="unknown"

  if [ -r "$path" ]; then
    IFS= read -r mac < "$path" || mac="unknown"
  fi

  mac="$(printf '%s' "$mac" | tr 'ABCDEFGHIJKLMNOPQRSTUVWXYZ' 'abcdefghijklmnopqrstuvwxyz' | tr -d '\r\n')"
  [ -n "$mac" ] || mac="unknown"
  printf '%s' "$mac"
}

ssid="$(read_env_key PORTAL_SSID || true)"
password="$(read_env_key PORTAL_PASSPHRASE || true)"

if valid_ssid "$ssid"; then
  print_kv visionai_setup_ssid "$ssid"
fi
if valid_password "$password"; then
  print_kv visionai_setup_password "$password"
fi

print_kv visionai_setup_gateway "$PORTAL_GATEWAY"
print_kv visionai_setup_url "http://$PORTAL_GATEWAY"
print_kv visionai_setup_ui_url "http://$PORTAL_GATEWAY:$VISIONAI_UI_PORT"
print_kv visionai_mac_wlan0 "$(get_mac wlan0)"
print_kv visionai_mac_eth0 "$(get_mac eth0)"
print_kv visionai_hostname "$(hostname 2>/dev/null || printf 'unknown')"
print_kv visionai_portal_interface "$PORTAL_INTERFACE"
INVENTORY
  chmod 0755 "$script"
  chown root:root "$script" 2>/dev/null || true
}

ensure_qrencode_available() {
  if command -v qrencode >/dev/null 2>&1; then
    return 0
  fi

  if [[ "${VISIONAI_DISABLE_QRENCODE_INSTALL:-0}" == "1" ]]; then
    return 1
  fi

  if command -v apt-get >/dev/null 2>&1; then
    log "Installing optional qrencode package"
    apt-get install -y --no-install-recommends qrencode || return 1
    command -v qrencode >/dev/null 2>&1
    return $?
  fi

  return 1
}

generate_visionai_wifi_qr_if_possible() {
  local ssid="$1"
  local password="$2"
  local qr_content="WIFI:T:WPA;S:${ssid};P:${password};;"

  if ! ensure_qrencode_available; then
    echo "qrencode not installed; skipping QR code generation." >&2
    return 0
  fi

  install -d -m 0755 "$(dirname "$VISIONAI_WIFI_QR_PNG")"
  if qrencode -o "$VISIONAI_WIFI_QR_PNG" -t PNG "$qr_content"; then
    chmod 0600 "$VISIONAI_WIFI_QR_PNG" 2>/dev/null || true
    echo
    echo "VisionAI Wi-Fi QR PNG generated: $VISIONAI_WIFI_QR_PNG"
    echo "Copy examples:"
    echo "  sudo cp $VISIONAI_WIFI_QR_PNG /boot/firmware/visionai-wifi-qr.png"
    echo "  scp mender@<device-ip>:$VISIONAI_WIFI_QR_PNG ."
  else
    warn "Could not generate QR PNG at $VISIONAI_WIFI_QR_PNG"
  fi

  echo
  echo "VisionAI Wi-Fi QR code:"
  qrencode -t ANSIUTF8 "$qr_content" || true
}

get_mac() {
  local iface="$1"
  local path="/sys/class/net/${iface}/address"
  local mac="unknown"

  if [[ -r "$path" ]]; then
    IFS= read -r mac < "$path" || mac="unknown"
  fi

  mac="$(printf '%s' "$mac" | tr '[:upper:]' '[:lower:]' | tr -d '\r\n')"
  [[ -n "$mac" ]] || mac="unknown"
  printf '%s\n' "$mac"
}

print_visionai_final_summary() {
  local ssid="$1"
  local password="$2"
  local device_type="$3"
  local hostname_value=""
  local wlan_mac="" eth_mac=""

  hostname_value="$(hostname 2>/dev/null || true)"
  [[ -n "$hostname_value" ]] || hostname_value="unknown"
  wlan_mac="$(get_mac wlan0)"
  eth_mac="$(get_mac eth0)"

  cat <<SUMMARY

==================================================
VISIONAI DEVICE SETUP DETAILS
==================================================
VisionAI Wi-Fi provisioning SSID:     $ssid
VisionAI Wi-Fi provisioning password: $password
Wi-Fi MAC address:                    $wlan_mac
Ethernet MAC address:                 $eth_mac
Gateway URL:                          http://$VISIONAI_PORTAL_GATEWAY
VisionAI UI URL while on hotspot:     http://$VISIONAI_PORTAL_GATEWAY:$VISIONAI_UI_PORT
Mender server URL:                    $SERVER_URL
device_type:                          $device_type
hostname:                             $hostname_value
Saved credentials file:               $VISIONAI_PROVISIONING_ENV
wifi-connect-manager defaults file:   $WIFI_CONNECT_DEFAULTS_FILE
Mender inventory script:              $VISIONAI_INVENTORY_SCRIPT
==================================================
SUMMARY
}

main() {
  need_root
  export DEBIAN_FRONTEND=noninteractive

  log "Detecting Raspberry Pi model"
  PI_MODEL="$(detect_pi_model || true)"
  if [[ -z "${PI_MODEL:-}" ]]; then
    die "Could not detect Raspberry Pi model. Is this running on a Raspberry Pi?"
  fi
  echo "Detected model: $PI_MODEL"

  if ! DEVICE_TYPE="$(detect_device_type "$PI_MODEL")"; then
    die "Unsupported Raspberry Pi model: $PI_MODEL (supported: Pi 4 / Pi 5)"
  fi
  echo "Using Mender device_type: $DEVICE_TYPE"

  log "Checking internet connectivity"
  if ! ping -c 1 -W 2 1.1.1.1 >/dev/null 2>&1; then
    warn "No response from 1.1.1.1. Continuing, but package install may fail if there is no internet."
  fi

  log "Installing prerequisites"
  apt-get update
  apt-get install -y --no-install-recommends \
    curl ca-certificates gnupg jq

  log "Setting up unique VisionAI Wi-Fi provisioning credentials"
  setup_visionai_wifi_credentials
  update_wifi_connect_manager_defaults
  install_wifi_connect_runtime_from_repo_if_available
  ensure_wifi_connect_launcher_uses_defaults
  install_wifi_connect_ui_if_possible
  generate_visionai_wifi_qr_if_possible "$PORTAL_SSID" "$PORTAL_PASSPHRASE"
  restart_wifi_connect_manager_if_present

  log "Downloading Mender installer helper"
  curl -fLsS https://get.mender.io -o /tmp/get-mender.sh
  chmod +x /tmp/get-mender.sh

  log "Installing Mender Client v4 + Mender Connect"
  # --force-mender-client4 ensures correct generation where supported
  bash /tmp/get-mender.sh --force-mender-client4 mender-client4 mender-connect

  log "Verifying Mender binaries"
  command -v mender-update >/dev/null || die "mender-update not found after install"
  if ! command -v mender-connect >/dev/null 2>&1; then
    warn "mender-connect binary not found in PATH (service may still exist depending on package layout)"
  fi

  log "Installing VisionAI Mender inventory script"
  install_visionai_inventory_script

  log "Writing device type files in correct Mender format"
  # IMPORTANT: file content must be exactly: device_type=<value>
  install -d -m 0755 "$(dirname "$DEVICE_TYPE_FILE")"
  printf 'device_type=%s\n' "$DEVICE_TYPE" > "$DEVICE_TYPE_FILE"
  chmod 0644 "$DEVICE_TYPE_FILE"

  # Fallback path for some layouts/tools
  install -d -m 0755 /etc/mender
  printf 'device_type=%s\n' "$DEVICE_TYPE" > /etc/mender/device_type
  chmod 0644 /etc/mender/device_type

  log "Writing /etc/mender/mender.conf (backup if exists)"
  if [[ -f /etc/mender/mender.conf ]]; then
    cp -a /etc/mender/mender.conf "/etc/mender/mender.conf.bak.$(date +%Y%m%d_%H%M%S)"
  fi

  SERVER_URL_JSON="$(json_escape "$SERVER_URL")"
  DEVICE_TYPE_FILE_JSON="$(json_escape "$DEVICE_TYPE_FILE")"

  if [[ -n "$SERVER_CERT_PATH" ]]; then
    SERVER_CERT_PATH_JSON="$(json_escape "$SERVER_CERT_PATH")"
    cat > /etc/mender/mender.conf <<JSON
{
  "ServerURL": "${SERVER_URL_JSON}",
  "DeviceTypeFile": "${DEVICE_TYPE_FILE_JSON}",
  "ServerCertificate": "${SERVER_CERT_PATH_JSON}"
}
JSON
  else
    cat > /etc/mender/mender.conf <<JSON
{
  "ServerURL": "${SERVER_URL_JSON}",
  "DeviceTypeFile": "${DEVICE_TYPE_FILE_JSON}"
}
JSON
  fi

  # Validate JSON
  jq -e . /etc/mender/mender.conf >/dev/null || die "Invalid JSON in /etc/mender/mender.conf"

  log "Enabling and restarting Mender services"
  # Some systems may not have mender-connect service installed yet if package failed partially
  systemctl enable --now mender-authd mender-updated || die "Failed to enable/start mender-authd or mender-updated"
  systemctl restart mender-authd mender-updated || true

  if systemctl list-unit-files | grep -q '^mender-connect\.service'; then
    systemctl enable --now mender-connect || warn "Could not enable/start mender-connect"
    systemctl restart mender-connect || warn "Could not restart mender-connect"
  else
    warn "mender-connect.service not found"
  fi

  log "Triggering initial Mender communication (best effort)"
  mender-update send-inventory || true
  mender-update check-update || true

  log "Debug info (important for device_type verification)"
  echo "--- /proc/device-tree/model ---"
  echo "$PI_MODEL"
  echo
  echo "--- ${DEVICE_TYPE_FILE} ---"
  cat "$DEVICE_TYPE_FILE" || true
  echo
  echo "--- /etc/mender/device_type ---"
  cat /etc/mender/device_type || true
  echo
  echo "--- /etc/mender/mender.conf ---"
  cat /etc/mender/mender.conf || true

  log "Service status"
  systemctl --no-pager --full status mender-authd mender-updated mender-connect || true

  print_visionai_final_summary "$PORTAL_SSID" "$PORTAL_PASSPHRASE" "$DEVICE_TYPE"

  log "DONE"
  echo "Expected Mender inventory device_type: ${DEVICE_TYPE}"
  echo "VisionAI provisioning inventory was sent with: mender-update send-inventory"
  echo "If device_type or VisionAI provisioning fields are still missing in UI, wait a few minutes and refresh inventory."
  echo "If needed, reboot once:"
  echo "  sudo reboot"
}

if [[ "${VISIONAI_INSTALL_AUTO_SOURCE_ONLY:-0}" != "1" ]]; then
  main "$@"
fi