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

EXTLINUX_CONF="/boot/extlinux/extlinux.conf"
BOOT_DTB_DIR="/boot/dtb"
DEFAULT_LABEL="afe-r750-camera"
DEFAULT_PHY_MODE="single"
SINGLE_PHY_DTBO="tegra234-p3767-camera-afe-r750-vls-gm2-overlay.dtbo"
SINGLE_PHY_FSYNC_DTBO="tegra234-p3767-camera-afe-r750-vls-gm2-fsync-overlay.dtbo"
DUAL_PHY_DTBO="tegra234-p3767-camera-afe-r750-dual-phy-vls-gm2-overlay.dtbo"
DUAL_PHY_FSYNC_DTBO="tegra234-p3767-camera-afe-r750-dual-phy-vls-gm2-fsync-overlay.dtbo"

print_usage() {
  cat <<'USAGE'
Usage: ./tn_install.sh [--yes] [--reboot yes|no] [--label LABEL_PREFIX] [--phy single|dual]

Installs AFE-R750 Orin NX camera modules and DTBOs from this package directory.
Without arguments, this script asks before installing and rebooting.

The package installs every DTBO in the package root and creates extlinux labels
for all AFE-R750 DTBOs. The selected PHY mode only chooses the DEFAULT label.
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
}

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
}

ARG_AUTO_YES=0
ARG_REBOOT=""
ARG_LABEL="$DEFAULT_LABEL"
ARG_PHY_MODE=""

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
      ;;
    --label)
      require_arg_value "$1" "${2:-}"
      ARG_LABEL="$2"
      shift 2
      ;;
    --phy)
      require_arg_value "$1" "${2:-}"
      ARG_PHY_MODE=$(normalize_phy_mode "$2")
      shift 2
      ;;
    --help|-h)
      print_usage
      exit 0
      ;;
    *)
      print_usage >&2
      die "unknown argument: $1"
      ;;
  esac
done

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_in() {
  local label=$1
  local conf=$2

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

label_exists() {
  local label=$1

  label_exists_in "$label" "$EXTLINUX_CONF"
}

get_default_label() {
  awk '
    /^[[:space:]]*(#|$)/ { next }
    toupper($1) == "DEFAULT" { print $2; found = 1; exit }
    END { exit found ? 0 : 1 }
  ' "$EXTLINUX_CONF" || die "cannot find DEFAULT label in $EXTLINUX_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
}

require_boot_path_nonempty() {
  local boot_path=$1
  local fs_path

  fs_path=$(resolve_boot_path "$boot_path")
  require_file_nonempty "$fs_path"
}

validate_extlinux_conf() {
  local conf=$1
  local label=${2:-}
  local overlay_path=${3:-}

  require_file_nonempty "$conf"
  grep -Eq '^[[:space:]]*DEFAULT[[:space:]]+[^[:space:]]+' "$conf" \
    || die "extlinux config has no DEFAULT entry: $conf"
  grep -Eq '^[[:space:]]*LABEL[[:space:]]+[^[:space:]]+' "$conf" \
    || die "extlinux config has no LABEL entry: $conf"

  if [ -n "$label" ]; then
    label_exists_in "$label" "$conf" \
      || die "generated extlinux config is missing label: $label"
  fi

  if [ -n "$overlay_path" ]; then
    awk -v overlay="$overlay_path" '
      toupper($1) == "OVERLAY" || toupper($1) == "OVERLAYS" || toupper($1) == "FDTOVERLAYS" {
        for (i = 2; i <= NF; i++) {
          if ($i == overlay) {
            found = 1
          }
        }
      }
      END { exit found ? 0 : 1 }
    ' "$conf" \
      || die "generated extlinux config is missing overlay: $overlay_path"
  fi
}

validate_extlinux_source() {
  validate_extlinux_conf "$EXTLINUX_CONF"
  find_base_label >/dev/null
}

install_extlinux_conf_file() {
  local source=$1
  local label=${2:-}
  local overlay_path=${3:-}
  local target_tmp="$EXTLINUX_CONF.tmp-$$"

  validate_extlinux_conf "$source" "$label" "$overlay_path"
  sudo_cmd rm -f "$target_tmp"
  sudo_cmd install -m 644 "$source" "$target_tmp"
  sudo_cmd mv "$target_tmp" "$EXTLINUX_CONF"
  validate_extlinux_conf "$EXTLINUX_CONF" "$label" "$overlay_path"
}

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

find_base_label() {
  if label_exists primary; then
    printf 'primary\n'
  else
    get_default_label
  fi
}

get_label_fdt() {
  local label=$1

  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 }
  ' "$EXTLINUX_CONF" || true
}

check_default_dtb() {
  local preferred_file=${1:-}
  local found_file

  [ -d "$BOOT_DTB_DIR" ] || die "missing boot DTB directory: $BOOT_DTB_DIR"

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

  found_file=$(find "$BOOT_DTB_DIR" -maxdepth 1 \( -type f -o -type l \) -printf '%f\n' |
    sort |
    sed -n '1p')

  [ -n "$found_file" ] || die "no DTB files found in $BOOT_DTB_DIR"
  require_file_nonempty "$BOOT_DTB_DIR/$found_file"
  printf '%s/%s\n' "$BOOT_DTB_DIR" "$found_file"
}

find_fdt_path() {
  local base_label=$1
  local fdt_path=${AFE_R750_FDT_PATH:-}
  local resolved_path
  local candidate

  if [ -n "$fdt_path" ]; then
    require_boot_path_nonempty "$fdt_path"
    printf '%s\n' "$fdt_path"
    return
  fi

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

  if candidate=$(check_default_dtb "kernel_tegra234-p3768-0000+p3767-0000-nv.dtb" 2>/dev/null); then
    printf '%s\n' "$candidate"
    return
  fi

  for candidate in \
    /boot/dtb/kernel_tegra234-p3768-0000+p3767-0000-nv.dtb \
    /boot/kernel_tegra234-p3768-0000+p3767-0000-nv.dtb \
    /boot/tegra234-p3768-0000+p3767-0000-nv.dtb \
    /boot/dtb/kernel_tegra234-p3768-0000+p3767-*-nv.dtb; do
    if [ -s "$candidate" ]; then
      printf '%s\n' "$candidate"
      return
    fi
  done

  die "cannot find a base FDT; set AFE_R750_FDT_PATH to the target DTB"
}

label_has_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" {
      found = 1
      exit
    }
    END { exit found ? 0 : 1 }
  ' "$conf"
}

ensure_label_has_fdt() {
  local label=$1
  local fdt_path=$2
  local tmp_conf

  label_exists "$label" || die "extlinux label does not exist: $label"
  require_boot_path_nonempty "$fdt_path"

  if label_has_fdt "$label"; then
    return
  fi

  tmp_conf=$(mktemp)
  awk -v label="$label" -v fdt="$fdt_path" '
    function line_key(line,    trimmed, fields) {
      trimmed = line
      sub(/^[ \t]+/, "", trimmed)
      split(trimmed, fields, /[ \t]+/)
      return toupper(fields[1])
    }

    function insert_fdt() {
      if (capture && !inserted) {
        print "      FDT " fdt
        inserted = 1
      }
    }

    {
      key = line_key($0)

      if (key == "LABEL") {
        if (capture) {
          insert_fdt()
        }

        capture = ($2 == label)
        inserted = 0

        if (capture) {
          found = 1
        }
      }

      if (capture) {
        if (key == "FDT") {
          inserted = 1
        } else if (!inserted &&
            (key == "APPEND" || key == "OVERLAY" || key == "OVERLAYS" || key == "FDTOVERLAYS")) {
          insert_fdt()
        }
      }

      print
    }

    END {
      if (capture) {
        insert_fdt()
      }
      exit found ? 0 : 33
    }
  ' "$EXTLINUX_CONF" > "$tmp_conf" || {
    rm -f "$tmp_conf"
    die "cannot add FDT to extlinux label: $label"
  }

  label_has_fdt "$label" "$tmp_conf" || die "cannot add FDT to extlinux label: $label"
  install_extlinux_conf_file "$tmp_conf" "$label"
  rm -f "$tmp_conf"

  label_has_fdt "$label" || die "extlinux label is still missing FDT: $label"
}

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

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

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

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

install_extlinux_entry() {
  local label=$1
  local dtbo=$2
  local overlay_path="/boot/$dtbo"
  local base_label
  local fdt_path
  local tmp_block
  local tmp_conf
  local tmp_new

  require_file_nonempty "$overlay_path"
  validate_extlinux_source

  base_label=$(find_base_label)
  fdt_path=$(find_fdt_path "$base_label")
  ensure_label_has_fdt "$base_label" "$fdt_path"
  tmp_block=$(write_extlinux_entry "$base_label" "$label" "$overlay_path" "$fdt_path")
  require_file_nonempty "$tmp_block"
  tmp_conf=$(mktemp)
  tmp_new=$(mktemp)
  {
    remove_extlinux_label "$label" "$tmp_conf"
    awk '
      { lines[NR] = $0 }
      END {
        n = NR
        while (n > 0 && lines[n] ~ /^[[:space:]]*$/) {
          n--
        }
        for (i = 1; i <= n; i++) {
          print lines[i]
        }
      }
    ' "$tmp_conf"
    printf '\n\n'
    cat "$tmp_block"
  } > "$tmp_new"
  install_extlinux_conf_file "$tmp_new" "$label" "$overlay_path"
  rm -f "$tmp_block" "$tmp_conf" "$tmp_new"
}

set_default_label() {
  local label=$1
  local tmp_conf

  label_exists "$label" || die "refusing to set DEFAULT to missing 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_extlinux_conf_file "$tmp_conf" "$label"
  rm -f "$tmp_conf"
}

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
}

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

dtbo_for_phy_mode() {
  case "$1" in
    single)
      printf '%s\n' "$SINGLE_PHY_DTBO"
      ;;
    dual)
      printf '%s\n' "$DUAL_PHY_DTBO"
      ;;
    *)
      die "cannot map PHY mode to DTBO: $1"
      ;;
  esac
}

label_for_phy_mode() {
  case "$1" in
    single)
      printf '%s\n' "$ARG_LABEL"
      ;;
    dual)
      printf '%s-dual-phy\n' "$ARG_LABEL"
      ;;
    *)
      die "cannot map PHY mode to extlinux label: $1"
      ;;
  esac
}

validate_required_extlinux_dtbos() {
  local dtbo_file

  for dtbo_file in "${EXTLINUX_DTBOS[@]}"; do
    require_file_nonempty "$dtbo_file"
  done
}

install_extlinux_entries() {
  local index

  for index in "${!EXTLINUX_LABELS[@]}"; do
    install_extlinux_entry "${EXTLINUX_LABELS[$index]}" "${EXTLINUX_DTBOS[$index]}"
  done
}

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

cat <<'BANNER'
****** TechNexion Camera Driver Installation ******
This installs AFE-R750 Orin NX camera modules and 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
read_phy_choice
SELECTED_PHY_MODE="$choice"
SELECTED_DTBO_FILE=$(dtbo_for_phy_mode "$SELECTED_PHY_MODE")
SELECTED_LABEL=$(label_for_phy_mode "$SELECTED_PHY_MODE")
EXTLINUX_LABELS=(
  "$ARG_LABEL"
  "$ARG_LABEL-fsync"
  "$ARG_LABEL-dual-phy"
  "$ARG_LABEL-dual-phy-fsync"
)
EXTLINUX_DTBOS=(
  "$SINGLE_PHY_DTBO"
  "$SINGLE_PHY_FSYNC_DTBO"
  "$DUAL_PHY_DTBO"
  "$DUAL_PHY_FSYNC_DTBO"
)
echo "Selected camera PHY mode: $SELECTED_PHY_MODE"
echo "Selected default extlinux label: $SELECTED_LABEL"
echo "Selected default DTBO: $SELECTED_DTBO_FILE"

sudo_cmd true >/dev/null

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

while IFS= read -r file; do
  MODULE_FILES+=("$file")
done < <(find technexion -type f -name '*.ko' | sort)

while IFS= read -r file; do
  DTBO_FILES+=("$file")
done < <(find . -maxdepth 1 -type f -name '*.dtbo' -printf '%P\n' | sort)

((${#MODULE_FILES[@]} > 0)) || die "no modules found under technexion/"
((${#DTBO_FILES[@]} > 0)) || die "no DTBO files found in package root"
validate_required_extlinux_dtbos
validate_extlinux_source

for module_file in "${MODULE_FILES[@]}"; do
  module_name=${module_file##*/}
  target_path="/lib/modules/$KERNEL_VER/updates/drivers/media/i2c/$module_file"
  stale_extra_path="/lib/modules/$KERNEL_VER/extra/$module_name"

  backup_if_exists "$target_path" "$TS"
  install_file_atomic "$module_file" "$target_path"

  if [ -e "$stale_extra_path" ]; then
    backup_if_exists "$stale_extra_path" "$TS"
    sudo_cmd rm -f "$stale_extra_path"
  fi
done

for dtbo_file in "${DTBO_FILES[@]}"; do
  backup_if_exists "/boot/$dtbo_file" "$TS"
  install_file_atomic "$dtbo_file" "/boot/$dtbo_file"
done

sudo_cmd depmod -a "$KERNEL_VER"

backup_if_exists "$EXTLINUX_CONF" "$TS"
install_extlinux_entries
set_default_label "$SELECTED_LABEL"

echo "Installed artifacts:"
installed_files=()
for module_file in "${MODULE_FILES[@]}"; do
  installed_files+=("/lib/modules/$KERNEL_VER/updates/drivers/media/i2c/$module_file")
done
for dtbo_file in "${DTBO_FILES[@]}"; do
  installed_files+=("/boot/$dtbo_file")
done
sha256sum "${installed_files[@]}"

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

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