#!/bin/bash
set -o pipefail

# TechNexion Camera Driver Modules Installation for
# Nvidia Jetson ORIN NANO Development Kits.
# Before installing, You should backup your image to avoid
# there are any files you lost.

APP_PRODUCT=2NOR0x
APP_VERSION=JP62
APP_DESCRIPTION='2NOR0x JP62'
EXTLINUX_CONF="/boot/extlinux/extlinux.conf"
BOOT_DTB_DIR="/boot/dtb"
REF_DTB_FILE=kernel_tegra234-p3768-0000+p3767-0003-nv.dtb
SUPPORTS_PHY=0
ARG_MODULE=""
ARG_AUTO_YES=0
ARG_REBOOT=""
ARG_PHY_MODE=single
INTERACTIVE=0
TS=""
KERNEL_VER=""

SUPPORTED_MODULES=(
  tevs
  vls-gm2
)

MODULE_FILES=(
  'technexion/maxim-gmsl/max96724.ko|updates/drivers/media/i2c'
  'technexion/maxim-gmsl/max96717.ko|updates/drivers/media/i2c'
  'technexion/tevs/tevs.ko|updates/drivers/media/i2c'
)

BOOT_FILES=(
  'dtbs/kernel_tegra234-p3768-0000+p3767-0003-nv.dtb|/boot/dtb/kernel_tegra234-p3768-0000+p3767-0003-nv.dtb'
  'dtbs/tegra234-p3767-camera-2nor0x-tevs-dual-overlay.dtbo|/boot/tegra234-p3767-camera-2nor0x-tevs-dual-overlay.dtbo'
  'dtbs/tegra234-p3767-camera-2nor0x-vls-gm2-overlay.dtbo|/boot/tegra234-p3767-camera-2nor0x-vls-gm2-overlay.dtbo'
  'dtbs/tegra234-p3767-camera-2nor0x-vls-gm2-fsync-overlay.dtbo|/boot/tegra234-p3767-camera-2nor0x-vls-gm2-fsync-overlay.dtbo'
  'dtbs/tegra234-p3767-camera-2nor0x-vls-gm2-tunnel-overlay.dtbo|/boot/tegra234-p3767-camera-2nor0x-vls-gm2-tunnel-overlay.dtbo'
  'dtbs/tegra234-p3767-camera-2nor0x-vls-gm2-tunnel-fsync-overlay.dtbo|/boot/tegra234-p3767-camera-2nor0x-vls-gm2-tunnel-fsync-overlay.dtbo'
)

EXTLINUX_ENTRIES=(
  'tevs|tevs|auto|OVERLAYS|/boot/tegra234-p3767-camera-2nor0x-tevs-dual-overlay.dtbo|'
  'vls-gm2|vls-gm2|auto|OVERLAYS|/boot/tegra234-p3767-camera-2nor0x-vls-gm2-overlay.dtbo|'
  'vls-gm2|vls-gm2-fsync|auto|OVERLAYS|/boot/tegra234-p3767-camera-2nor0x-vls-gm2-fsync-overlay.dtbo|'
  'vls-gm2|vls-gm2-tunnel|auto|OVERLAYS|/boot/tegra234-p3767-camera-2nor0x-vls-gm2-tunnel-overlay.dtbo|'
  'vls-gm2|vls-gm2-tunnel-fsync|auto|OVERLAYS|/boot/tegra234-p3767-camera-2nor0x-vls-gm2-tunnel-fsync-overlay.dtbo|'
)

print_usage() {
  cat <<USAGE
Usage: ./tn_install.sh [--yes] [--module $(IFS='|'; echo "${SUPPORTED_MODULES[*]}")] [--reboot yes|no]

Installs ${APP_DESCRIPTION} camera driver artifacts.
Without arguments, this script runs an interactive install flow.
With --module, it runs non-interactively for APP-driven installation.
USAGE
}

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

warn() {
  printf 'WARN: %s\n' "$*" >&2
}

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

require_root() {
  [ "$(id -u)" -eq 0 ] || die "run this script as root through the APP runner, pkexec, or 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
}

normalize_phy_mode() {
  case "$1" in
    ""|single|SINGLE|Single|s|S|1)
      printf 'single\n'
      ;;
    dual|DUAL|Dual|d|D|2)
      printf 'dual\n'
      ;;
    *)
      die "cannot support this PHY mode: $1"
      ;;
  esac
}

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

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

install_file_atomic() {
  local source=$1
  local target=$2
  local mode=${3:-644}
  local target_dir=${target%/*}
  local target_tmp="$target_dir/.${target##*/}.tmp-$$"

  require_file_nonempty "$source"
  mkdir -p "$target_dir"
  rm -f "$target_tmp"
  install -m "$mode" "$source" "$target_tmp"
  require_file_nonempty "$target_tmp"
  sync -f "$target_tmp" 2>/dev/null || sync
  mv "$target_tmp" "$target"
  require_file_nonempty "$target"
  sync -f "$target" 2>/dev/null || sync
}

normalize_module() {
  case "$1" in
    tevs|TEVS)
      printf '%s\n' tevs
      ;;
    vls-gm2|VLS-GM2|vls_gm2|VLS_GM2)
      printf '%s\n' vls-gm2
      ;;
    *)
      die "unsupported module for ${APP_PRODUCT} ${APP_VERSION}: $1"
      ;;
  esac
}

parse_args() {
  if [ "$#" -eq 0 ]; then
    INTERACTIVE=1
    return
  fi

  while [ "$#" -gt 0 ]; do
    case "$1" in
      --module)
        require_arg_value "$1" "${2:-}"
        ARG_MODULE=$(normalize_module "$2")
        shift 2
        ;;
      --reboot)
        require_arg_value "$1" "${2:-}"
        ARG_REBOOT=$(normalize_yes_no "$2")
        shift 2
        ;;
      --yes|-y)
        ARG_AUTO_YES=1
        shift
        ;;
      --help|-h)
        print_usage
        exit 0
        ;;
      --platform|--label|--phy)
        print_usage >&2
        die "$1 is not accepted by this package installer"
        ;;
      *)
        print_usage >&2
        die "unknown argument: $1"
        ;;
    esac
  done
}

select_module_interactive() {
  local choice
  local index

  if [ "${#SUPPORTED_MODULES[@]}" -eq 1 ]; then
    ARG_MODULE=$(normalize_module "${SUPPORTED_MODULES[0]}")
    return
  fi

  echo "Camera module:"
  index=1
  for module in "${SUPPORTED_MODULES[@]}"; do
    printf '  [%d] %s\n' "$index" "$module"
    index=$((index + 1))
  done
  read -r -p "Select module [1-${#SUPPORTED_MODULES[@]}]: " choice
  case "$choice" in
    ''|*[!0-9]*)
      ARG_MODULE=$(normalize_module "$choice")
      ;;
    *)
      [ "$choice" -ge 1 ] && [ "$choice" -le "${#SUPPORTED_MODULES[@]}" ] ||
        die "module selection is out of range: $choice"
      ARG_MODULE=$(normalize_module "${SUPPORTED_MODULES[$((choice - 1))]}")
      ;;
  esac
}

prepare_install_choices() {
  local choice

  if [ -z "$ARG_MODULE" ]; then
    if [ "$INTERACTIVE" = "1" ]; then
      select_module_interactive
    elif [ "$ARG_AUTO_YES" = "1" ] && [ "${#SUPPORTED_MODULES[@]}" -eq 1 ]; then
      ARG_MODULE=$(normalize_module "${SUPPORTED_MODULES[0]}")
    else
      print_usage >&2
      die "--module is required for non-interactive installation"
    fi
  fi

  if [ "$SUPPORTS_PHY" = "1" ]; then
    if [ "$INTERACTIVE" = "1" ]; then
      read -r -p "Camera PHY mode [single/dual] (default: single): " choice
      ARG_PHY_MODE=$(normalize_phy_mode "${choice:-single}")
    else
      ARG_PHY_MODE="single"
    fi
  fi

  if [ "$INTERACTIVE" = "1" ]; then
    read -r -p "Install ${APP_DESCRIPTION} for ${ARG_MODULE}? [Y/n] " choice
    choice=${choice:-y}
    choice=$(normalize_yes_no "$choice")
    if [ "$choice" != "y" ]; then
      echo "Installation aborted."
      exit 0
    fi
    read -r -p "Reboot after installation? [y/N] " choice
    ARG_REBOOT=$(normalize_yes_no "${choice:-n}")
  elif [ -z "$ARG_REBOOT" ]; then
    ARG_REBOOT="n"
  fi
}

default_label_for_module() {
  case "$1" in
    tevs)
      printf '%s\n' tevs
      ;;
    vls-gm2)
      printf '%s\n' vls-gm2
      ;;
    *)
      die "missing default extlinux label for module: $1"
      ;;
  esac
}

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_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"
}

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

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

find_base_label() {
  local default_label

  if label_exists primary; then
    printf 'primary\n'
    return
  fi

  default_label=$(get_default_label || true)
  if [ -n "$default_label" ] && [[ "$default_label" != tn-* ]] && label_exists "$default_label"; then
    printf '%s\n' "$default_label"
    return
  fi

  awk '
    toupper($1) == "LABEL" && $2 !~ /^tn-/ {
      print $2
      found = 1
      exit
    }
    END { exit found ? 0 : 1 }
  ' "$EXTLINUX_CONF" || die "cannot find a non-TN extlinux label to clone"
}

first_boot_dtb() {
  [ -d "$BOOT_DTB_DIR" ] || die "missing boot DTB directory: $BOOT_DTB_DIR"
  find "$BOOT_DTB_DIR" -maxdepth 1 \( -type f -o -type l \) -printf '%f\n' | sort | sed -n '1p'
}

resolve_auto_fdt() {
  local base_label=$1
  local label_fdt
  local resolved
  local first_dtb

  if [ -n "$REF_DTB_FILE" ] && [ -s "$BOOT_DTB_DIR/$REF_DTB_FILE" ]; then
    printf '%s/%s\n' "$BOOT_DTB_DIR" "$REF_DTB_FILE"
    return
  fi

  label_fdt=$(get_label_fdt "$base_label")
  if [ -n "$label_fdt" ]; then
    resolved=$(resolve_boot_path "$label_fdt")
    if [ -s "$resolved" ]; then
      printf '%s\n' "$label_fdt"
      return
    fi
  fi

  first_dtb=$(first_boot_dtb)
  [ -n "$first_dtb" ] || die "no DTB files found in $BOOT_DTB_DIR"
  require_file_nonempty "$BOOT_DTB_DIR/$first_dtb"
  printf '%s/%s\n' "$BOOT_DTB_DIR" "$first_dtb"
}

resolve_entry_fdt() {
  local requested=$1
  local base_label=$2
  local resolved

  case "$requested" in
    ""|auto)
      resolve_auto_fdt "$base_label"
      ;;
    /*)
      resolved=$(resolve_boot_path "$requested")
      require_file_nonempty "$resolved"
      printf '%s\n' "$requested"
      ;;
    *)
      resolved="$BOOT_DTB_DIR/$requested"
      require_file_nonempty "$resolved"
      printf '%s\n' "$resolved"
      ;;
  esac
}

remove_legacy_tncam_service() {
	TGT_SERVICE_FILE="/etc/systemd/system/tn_cam.service"
    if [ -f "$TGT_SERVICE_FILE" ]; then
        echo "Remove legacy TN-CAM service..."
        if command -v systemctl >/dev/null 2>&1; then
            sudo systemctl stop tn_cam.service 2>/dev/null || true
            sudo systemctl disable tn_cam.service 2>/dev/null || true
        fi
        sudo rm -f "$TGT_SERVICE_FILE"
        if command -v systemctl >/dev/null 2>&1; then
            sudo systemctl daemon-reload
        fi
    fi
}

remove_legacy_tncam_modules() {
	LEGACY_MODULES_FILE="/lib/modules/$KERNEL_VER/kernel/drivers/media/i2c"
    if [ -f "$LEGACY_MODULES_FILE/$1" ]; then
        echo "Remove legacy TN-CAM module $1..."
        sudo rm -f "$LEGACY_MODULES_FILE/$1"
    fi
}

remove_extlinux_label() {
  local label=$1
  local source_conf=$2

  awk -v label="$label" '
    toupper($1) == "LABEL" {
      skip = ($2 == label)
    }
    !skip { print }
  ' "$source_conf"
}

cleanup_tn_extlinux_labels() {
  local tmp_conf

  tmp_conf=$(mktemp)
  awk '
    function flush_blanks() {
      for (i = 1; i <= blank_count; i++) {
        print ""
      }
      blank_count = 0
    }

    /^[[:space:]]*$/ {
      if (skip) {
        next
      }
      blank_count++
      next
    }

    /^[[:space:]]*LABEL[[:space:]]+tn-/ {
      blank_count = 0
      skip = 1
      next
    }

    /^[[:space:]]*LABEL[[:space:]]+/ {
      skip = 0
    }

    {
      if (!skip) {
        flush_blanks()
        print
      }
    }

    END {
      if (!skip) {
        flush_blanks()
      }
    }
  ' "$EXTLINUX_CONF" > "$tmp_conf"

  require_file_nonempty "$tmp_conf"
  install_file_atomic "$tmp_conf" "$EXTLINUX_CONF"
  rm -f "$tmp_conf"
}

write_extlinux_entry_block() {
  local base_label=$1
  local new_label=$2
  local fdt_path=$3
  local overlay_keyword=$4
  local overlay_path=$5
  local tmp_block

  tmp_block=$(mktemp)
  awk \
    -v base="$base_label" \
    -v new="$new_label" \
    -v fdt="$fdt_path" \
    -v overlay_keyword="$overlay_keyword" \
    -v overlay="$overlay_path" '
    function emit_insertions() {
      if (!inserted_fdt) {
        print "      FDT " fdt
        inserted_fdt = 1
      }
      if (overlay != "" && !inserted_overlay) {
        print "      " overlay_keyword " " overlay
        inserted_overlay = 1
      }
    }
    function emit_block(    i, line, trimmed, fields, key) {
      inserted_fdt = 0
      inserted_overlay = 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 == "FDT" || key == "OVERLAY" || key == "OVERLAYS" || key == "FDTOVERLAYS") {
          continue
        } else if (key == "APPEND") {
          emit_insertions()
        }
        print line
      }
      emit_insertions()
    }
    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"
}

validate_extlinux_label() {
  local label=$1
  local conf=$2

  label_exists "$label" "$conf" || die "generated extlinux config is missing label: $label"
}

install_extlinux_entry() {
  local label=$1
  local requested_fdt=$2
  local overlay_keyword=$3
  local overlay_path=$4
  local base_label
  local fdt_path
  local tmp_block
  local tmp_conf

  base_label=$(find_base_label)
  fdt_path=$(resolve_entry_fdt "$requested_fdt" "$base_label")

  if [ -n "$overlay_path" ]; then
    require_file_nonempty "$(resolve_boot_path "$overlay_path")"
  fi

  tmp_block=$(write_extlinux_entry_block "$base_label" "tn-$label" "$fdt_path" "$overlay_keyword" "$overlay_path")
  tmp_conf=$(mktemp)
  {
    remove_extlinux_label "tn-$label" "$EXTLINUX_CONF"
    printf '\n\n'
    cat "$tmp_block"
  } > "$tmp_conf"

  validate_extlinux_label "tn-$label" "$tmp_conf"
  install_file_atomic "$tmp_conf" "$EXTLINUX_CONF"
  rm -f "$tmp_block" "$tmp_conf"
}

set_default_label() {
  local label=$1
  local tmp_conf

  label="tn-$label"
  label_exists "$label" || die "refusing to set DEFAULT to missing extlinux label: $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
  install_file_atomic "$tmp_conf" "$EXTLINUX_CONF"
  rm -f "$tmp_conf"
}

install_modules() {
  local row
  local source
  local target_subdir
  local target

  for row in "${MODULE_FILES[@]}"; do
    IFS='|' read -r source target_subdir <<< "$row"
    target="/lib/modules/$KERNEL_VER/$target_subdir/${source##*/}"
    backup_if_exists "$target"
    install_file_atomic "$source" "$target"
  done

  if ((${#MODULE_FILES[@]} > 0)); then
    depmod -a "$KERNEL_VER"
  fi
}

install_boot_files() {
  local row
  local source
  local target

  for row in "${BOOT_FILES[@]}"; do
    IFS='|' read -r source target <<< "$row"
    backup_if_exists "$target"
    install_file_atomic "$source" "$target"
  done
}

install_extlinux_entries() {
  local row
  local entry_module
  local label
  local fdt
  local overlay_keyword
  local overlay
  local phy_mode
  local installed=0

  require_file_nonempty "$EXTLINUX_CONF"
  backup_if_exists "$EXTLINUX_CONF"
  cleanup_tn_extlinux_labels

  for row in "${EXTLINUX_ENTRIES[@]}"; do
    IFS='|' read -r entry_module label fdt overlay_keyword overlay phy_mode <<< "$row"
    [ "$entry_module" = "$ARG_MODULE" ] || continue
    if [ -n "${phy_mode:-}" ] && [ "$phy_mode" != "$ARG_PHY_MODE" ]; then
      continue
    fi
    install_extlinux_entry "$label" "$fdt" "$overlay_keyword" "$overlay"
    installed=$((installed + 1))
  done

  [ "$installed" -gt 0 ] || die "no extlinux entries configured for module: $ARG_MODULE"
  set_default_label "$(default_label_for_module "$ARG_MODULE")"
}

show_summary() {
  local row
  local source
  local target_subdir
  local target
  local installed_files=()

  for row in "${MODULE_FILES[@]}"; do
    IFS='|' read -r source target_subdir <<< "$row"
    installed_files+=("/lib/modules/$KERNEL_VER/$target_subdir/${source##*/}")
  done

  for row in "${BOOT_FILES[@]}"; do
    IFS='|' read -r source target <<< "$row"
    installed_files+=("$target")
  done

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

main() {
  local script_dir

  parse_args "$@"
  prepare_install_choices

  for cmd in awk cat chmod cp date depmod find grep id install mkdir mktemp mv rm sed sha256sum sort sync uname; do
    require_cmd "$cmd"
  done
  require_root

  script_dir=$(CDPATH= cd -- "$(dirname -- "$0")" && pwd -P)
  cd "$script_dir"

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

  echo "Installing ${APP_DESCRIPTION}"
  echo "Selected module: $ARG_MODULE"
  if [ "$SUPPORTS_PHY" = "1" ]; then
    echo "Selected PHY mode: $ARG_PHY_MODE"
  fi
  
  remove_legacy_tncam_modules "tevs.ko"
  remove_legacy_tncam_modules "max96724.ko"
  remove_legacy_tncam_modules "max96717.ko"
  remove_legacy_tncam_service
  

  install_modules
  install_boot_files
  install_extlinux_entries
  show_summary

  echo "Finish Camera Driver Installation."
  if [ "$ARG_REBOOT" = "y" ]; then
    echo "Rebooting..."
    reboot
  else
    echo "Reboot the device when the user is ready."
  fi
}

main "$@"
