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

MODULE_FILES=(
  max9296a.ko
  max96717.ko
  tevs.ko
)

DTBO_FILES=(
  tegra234-p3767-camera-p3768-x6orn-vls-gm2-overlay.dtbo
  tegra234-p3767-camera-p3768-x6orn-vls-gm2-fsync-overlay.dtbo
)

SERVICE_FILE="tn_cam.service"
TGT_SERVICE_FILE="/etc/systemd/system/$SERVICE_FILE"
EXTLINUX_CONF="/boot/extlinux/extlinux.conf"
DEFAULT_MODULE="vls-gm2"

ARG_AUTO_YES=0
ARG_REBOOT=""

print_usage() {
  cat <<'USAGE'
Usage: ./tn_install.sh [--yes] [--reboot yes|no]
Installs X6-ORN VLS-GM2 camera modules and makes tn-vls-gm2 the default boot entry.

Without arguments, this script asks before installing and rebooting.
USAGE
}

die() {
  printf 'ERROR: %s\n' "$*" >&2
  exit 1
}

require_cmd() {
  command -v "$1" >/dev/null 2>&1 || die "missing required command: $1"
}

sudo_cmd() {
  sudo "$@"
}

require_arg_value() {
  case "${2:-}" in
    ""|--*)
      print_usage >&2
      die "missing value for $1"
      ;;
  esac
}

normalize_yes_no() {
  case "$1" in
    y|Y|yes|YES|true|TRUE|1)
      printf 'y\n'
      ;;
    n|N|no|NO|false|FALSE|0)
      printf 'n\n'
      ;;
    *)
      die "cannot support this yes/no option: $1"
      ;;
  esac
}

parse_args() {
  while [ "$#" -gt 0 ]; do
    case "$1" in
      --yes|-y)
        ARG_AUTO_YES=1
        shift
        ;;
      --reboot)
        require_arg_value "$1" "${2:-}"
        ARG_REBOOT=$(normalize_yes_no "$2")
        shift 2
        ;;
      --help|-h)
        print_usage
        exit 0
        ;;
      *)
        print_usage >&2
        die "unknown argument: $1"
        ;;
    esac
  done
}

read_continue_choice() {
  if [ "$ARG_AUTO_YES" = "1" ]; then
    choice="y"
  else
    read -r -p "Do you want to continue?[Y/n] " choice
    choice=${choice:-y}
  fi
}

read_reboot_choice() {
  if [ -n "$ARG_REBOOT" ]; then
    choice="$ARG_REBOOT"
  else
    read -r -p "Do you want to reboot now?[Y/n] " choice
    choice=${choice:-y}
    choice=$(normalize_yes_no "$choice")
  fi
}

require_file_nonempty() {
  local path=$1

  [ -f "$path" ] || die "missing file: $path"
  [ -s "$path" ] || die "zero-byte file: $path"
}

backup_if_exists() {
  local path=$1
  local ts=$2

  if [ -e "$path" ]; then
    sudo_cmd cp -a "$path" "$path.bak-$ts"
  fi
}

label_exists() {
  local label=$1
  local conf=${2:-$EXTLINUX_CONF}

  awk -v label="$label" '
    toupper($1) == "LABEL" && $2 == label { found = 1 }
    END { exit found ? 0 : 1 }
  ' "$conf"
}

get_default_label() {
  local conf=${1:-$EXTLINUX_CONF}

  awk '
    /^[[:space:]]*(#|$)/ { next }
    toupper($1) == "DEFAULT" { print $2; found = 1; exit }
    END { exit found ? 0 : 1 }
  ' "$conf" || die "cannot find DEFAULT label in $conf"
}

resolve_boot_path() {
  local boot_path=$1

  boot_path=${boot_path%\"}
  boot_path=${boot_path#\"}

  case "$boot_path" in
    /*)
      printf '%s\n' "$boot_path"
      ;;
    *)
      printf '/boot/%s\n' "${boot_path#./}"
      ;;
  esac
}

label_boot_paths() {
  local label=$1
  local conf=${2:-$EXTLINUX_CONF}

  awk -v label="$label" '
    toupper($1) == "LABEL" {
      if (capture) {
        exit
      }
      capture = ($2 == label)
    }

    capture {
      key = toupper($1)
      if (key == "FDT" || key == "OVERLAY" || key == "OVERLAYS" || key == "FDTOVERLAYS") {
        for (i = 2; i <= NF; i++) {
          if ($i ~ /^#/) {
            break
          }
          print $i
        }
      }
    }
  ' "$conf"
}

validate_label_boot_files() {
  local label=$1
  local conf=${2:-$EXTLINUX_CONF}
  local boot_path
  local fs_path

  label_exists "$label" "$conf" || die "extlinux label does not exist: $label"

  while IFS= read -r boot_path; do
    [ -n "$boot_path" ] || continue

    case "$boot_path" in
      *'$'*)
        die "cannot validate variable boot path in label $label: $boot_path"
        ;;
    esac

    fs_path=$(resolve_boot_path "$boot_path")
    require_file_nonempty "$fs_path"
  done < <(label_boot_paths "$label" "$conf")
}

validate_default_boot() {
  local conf=${1:-$EXTLINUX_CONF}
  local label

  label=$(get_default_label "$conf")
  validate_label_boot_files "$label" "$conf"
}

require_label_uses_path() {
  local label=$1
  local expected_path=$2
  local conf=${3:-$EXTLINUX_CONF}

  label_boot_paths "$label" "$conf" | grep -Fx -- "$expected_path" >/dev/null ||
    die "extlinux label $label does not reference expected path: $expected_path"
}

find_base_label() {
  local label

  if label_exists primary; then
    label=primary
  else
    label=$(get_default_label)
  fi

  validate_label_boot_files "$label" "$EXTLINUX_CONF"
  printf '%s\n' "$label"
}

write_extlinux_entry() {
  local base_label=$1
  local new_label=$2
  local overlay_path=$3
  local tmp_block

  tmp_block=$(mktemp)
  awk -v base="$base_label" -v new="$new_label" -v overlay="$overlay_path" '
    function emit_block(    i, line, trimmed, fields, key, inserted) {
      inserted = 0

      for (i = 1; i <= n; i++) {
        line = lines[i]
        trimmed = line
        sub(/^[ \t]+/, "", trimmed)
        split(trimmed, fields, /[ \t]+/)
        key = toupper(fields[1])

        if (key == "LABEL") {
          sub(/LABEL[ \t]+[^ \t]+/, "LABEL " new, line)
        } else if (key == "MENU" && toupper(fields[2]) == "LABEL") {
          line = "      MENU LABEL " new
        } else if (key == "OVERLAY" || key == "OVERLAYS" || key == "FDTOVERLAYS") {
          continue
        } else if (!inserted && key == "APPEND") {
          print "      OVERLAYS " overlay
          inserted = 1
        }

        print line
      }

      if (!inserted) {
        print "      OVERLAYS " overlay
      }
    }

    toupper($1) == "LABEL" {
      if (capture) {
        exit
      }
      capture = ($2 == base)
    }

    capture {
      lines[++n] = $0

      line = $0
      sub(/^[ \t]+/, "", line)
      split(line, fields, /[ \t]+/)
      if (toupper(fields[1]) == "APPEND") {
        exit
      }
    }

    END {
      if (n == 0) {
        exit 33
      }
      emit_block()
    }
  ' "$EXTLINUX_CONF" > "$tmp_block" || {
    rm -f "$tmp_block"
    die "cannot clone extlinux label: $base_label"
  }

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

append_extlinux_entry() {
  local base_label=$1
  local new_label=$2
  local dtbo_file=$3
  local overlay_path="/boot/$dtbo_file"
  local tmp_block
  local tmp_conf

  require_file_nonempty "/boot/$dtbo_file"

  if label_exists "$new_label"; then
    echo "Extlinux label already exists: $new_label"
    require_label_uses_path "$new_label" "$overlay_path"
    validate_label_boot_files "$new_label"
    return
  fi

  tmp_block=$(write_extlinux_entry "$base_label" "$new_label" "$overlay_path")
  tmp_conf=$(mktemp)
  {
    cat "$EXTLINUX_CONF"
    printf '\n\n'
    cat "$tmp_block"
  } > "$tmp_conf"

  label_exists "$new_label" "$tmp_conf" || die "cannot create extlinux label: $new_label"
  require_label_uses_path "$new_label" "$overlay_path" "$tmp_conf"
  validate_label_boot_files "$new_label" "$tmp_conf"

  sudo_cmd install -m 644 "$tmp_conf" "$EXTLINUX_CONF"
  rm -f "$tmp_block" "$tmp_conf"

  label_exists "$new_label" || die "extlinux label was not installed: $new_label"
  validate_label_boot_files "$new_label"
}

set_default_conf() {
  local label=$1
  local tmp_conf

  label_exists "$label" || die "refusing to set DEFAULT to missing extlinux label: $label"
  validate_label_boot_files "$label"

  tmp_conf=$(mktemp)
  if grep -Eq '^[[:space:]]*DEFAULT[[:space:]]+' "$EXTLINUX_CONF"; then
    sed "s/^[[:space:]]*DEFAULT[[:space:]].*/DEFAULT $label/" "$EXTLINUX_CONF" > "$tmp_conf"
  else
    {
      printf 'DEFAULT %s\n' "$label"
      cat "$EXTLINUX_CONF"
    } > "$tmp_conf"
  fi

  [ "$(get_default_label "$tmp_conf")" = "$label" ] || die "cannot set DEFAULT to: $label"
  validate_default_boot "$tmp_conf"

  sudo_cmd install -m 644 "$tmp_conf" "$EXTLINUX_CONF"
  rm -f "$tmp_conf"

  [ "$(get_default_label)" = "$label" ] || die "failed to set DEFAULT to: $label"
  validate_default_boot
}

install_tncam_modules() {
  local kernel_ver=$1
  local module_file
  local target_dir="/lib/modules/$kernel_ver/extra"

  sudo_cmd mkdir -p "$target_dir"

  for module_file in "${MODULE_FILES[@]}"; do
    require_file_nonempty "$module_file"
    backup_if_exists "$target_dir/$module_file" "$TS"
    sudo_cmd install -m 644 "$module_file" "$target_dir/$module_file"
    require_file_nonempty "$target_dir/$module_file"
  done

  sudo_cmd depmod -a "$kernel_ver"
}

install_tncam_dtbo() {
  local dtbo_file

  for dtbo_file in "${DTBO_FILES[@]}"; do
    require_file_nonempty "$dtbo_file"
    backup_if_exists "/boot/$dtbo_file" "$TS"
    sudo_cmd install -m 644 "$dtbo_file" "/boot/$dtbo_file"
    require_file_nonempty "/boot/$dtbo_file"
  done
}

install_tncam_extlinux() {
  local module=$1
  local base_label

  require_file_nonempty "$EXTLINUX_CONF"
  backup_if_exists "$EXTLINUX_CONF" "$TS"

  base_label=$(find_base_label)
  append_extlinux_entry "$base_label" "tn-vls-gm2" \
    "tegra234-p3767-camera-p3768-x6orn-vls-gm2-overlay.dtbo"
  append_extlinux_entry "$base_label" "tn-vls-gm2-fsync" \
    "tegra234-p3767-camera-p3768-x6orn-vls-gm2-fsync-overlay.dtbo"
  set_default_conf "tn-$module"
}

install_tncam_service() {
  require_file_nonempty "$SERVICE_FILE"

  sudo_cmd install -m 644 "$SERVICE_FILE" "$TGT_SERVICE_FILE"
  sudo_cmd systemctl daemon-reload
  sudo_cmd systemctl enable "$SERVICE_FILE"
  sudo_cmd systemctl restart "$SERVICE_FILE" 2>/dev/null || true
}

show_installed_summary() {
  local kernel_ver=$1

  validate_default_boot

  echo "Installed artifacts:"
  sha256sum \
    "/lib/modules/$kernel_ver/extra/max9296a.ko" \
    "/lib/modules/$kernel_ver/extra/max96717.ko" \
    "/lib/modules/$kernel_ver/extra/tevs.ko" \
    "/boot/tegra234-p3767-camera-p3768-x6orn-vls-gm2-overlay.dtbo" \
    "/boot/tegra234-p3767-camera-p3768-x6orn-vls-gm2-fsync-overlay.dtbo"

  echo "Extlinux camera entries:"
  awk '/^[[:space:]]*DEFAULT/ || /^[[:space:]]*LABEL/ || /FDT/ || /OVERLAYS/ { print }' "$EXTLINUX_CONF"
}

main() {
  local module=$DEFAULT_MODULE

  parse_args "$@"

  for cmd in awk cat cp date grep install mkdir mktemp rm sed sha256sum systemctl uname; do
    require_cmd "$cmd"
  done

  cat <<'BANNER'
****** TechNexion Camera Driver Installation ******
This installs X6-ORN VLS-GM2 camera modules and both X6-ORN DTBO overlays.
Back up the image before installing on a production system.
BANNER

  read_continue_choice
  if [ "$choice" != "y" ] && [ "$choice" != "Y" ]; then
    echo "Installation aborted."
    exit 0
  fi

  echo "Installing default module: $module"

  sudo_cmd true >/dev/null

  TS=$(date +%Y%m%d-%H%M%S)
  KERNEL_VER=$(uname -r)

  install_tncam_modules "$KERNEL_VER"
  install_tncam_dtbo
  install_tncam_extlinux "$module"
  install_tncam_service
  show_installed_summary "$KERNEL_VER"

  echo "Finish Camera Driver Installation."
  echo "You should reboot the device to enable $module cameras."
  read_reboot_choice
  if [ "$choice" = "y" ]; then
    echo "Rebooting..."
    sudo_cmd reboot
  else
    echo "Reboot manually later."
  fi
}

main "$@"
