#!/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.

MODULE_FILE="tevs.ko"

REF_LEBEL_NAME="primary"
BOOT_DTB_DIR="/boot/dtb/"
REF_DTB_FILE="kernel_tegra234-p3768-0000+p3767-0005-nv.dtb"
EXTLINUX_CONF="/boot/extlinux/extlinux.conf"
LEGACY_SERVICE_FILE="/etc/systemd/system/tn_cam.service"
INSTALL_PLATFORM=""
INSTALL_MODULE=""
REBOOT_AFTER_INSTALL=""
NONINTERACTIVE=0

fail() {
    echo "ERROR: $*" >&2
    exit 27
}

sudo_or_fail() {
    sudo "$@" || fail "Command failed: sudo $*"
}

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" || {
    echo "ERROR: Cannot resolve installer directory." >&2
    exit 27
}
cd "$SCRIPT_DIR" || fail "Cannot enter installer directory: $SCRIPT_DIR"

usage() {
cat << EOF
Usage: ./tn_install.sh [--platform orin_nano|agx_orin] [--module tevs|vls-fpd3|vls-gm2] [--reboot yes|no]

Without arguments, the installer runs interactively.
With arguments, the installer runs non-interactively and --module is required.

Options:
  --platform VALUE   Target platform. Defaults to orin_nano in non-interactive mode.
  --module VALUE     Camera module to enable: tevs, vls-fpd3, or vls-gm2.
  --reboot VALUE     Reboot after installation: yes or no. Defaults to no in non-interactive mode.
  --yes              Compatibility flag for automation; prompts are already skipped when arguments are used.
  -h, --help         Show this help.

Aliases:
  platform: 1=orin_nano, 2=agx_orin
  module: 1=tevs, 2=vls-fpd3, 3=vls-gm2, vls3=vls-fpd3
EOF
}

normalize_platform() {
    local value
    value="$(printf '%s' "$1" | tr '[:upper:]' '[:lower:]')"
    value="${value//-/_}"
    value="${value// /_}"

    case "$value" in
        ""|1|orin_nano|orin_nano_evk|jetson_orin_nano)
            printf '%s\n' "orin_nano"
            ;;
        2|agx_orin|jetson_agx_orin)
            printf '%s\n' "agx_orin"
            ;;
        *)
            return 1
            ;;
    esac
}

normalize_module() {
    local value
    value="$(printf '%s' "$1" | tr '[:upper:]' '[:lower:]')"
    value="${value//-/_}"
    value="${value// /_}"

    case "$value" in
        1|tevs)
            printf '%s\n' "tevs"
            ;;
        2|vls3|vls_fpd3|vl316_vls)
            printf '%s\n' "vls-fpd3"
            ;;
        3|vls_gm2)
            printf '%s\n' "vls-gm2"
            ;;
        *)
            return 1
            ;;
    esac
}

normalize_reboot() {
    local value
    value="$(printf '%s' "$1" | tr '[:upper:]' '[:lower:]')"

    case "$value" in
        ""|no|n)
            printf '%s\n' "no"
            ;;
        yes|y)
            printf '%s\n' "yes"
            ;;
        *)
            return 1
            ;;
    esac
}

parse_args() {
    local reboot_choice

    while [ "$#" -gt 0 ]; do
        case "$1" in
            -h|--help)
                usage
                exit 0
                ;;
            --yes)
                NONINTERACTIVE=1
                shift
                ;;
            --platform)
                if [ "$#" -lt 2 ]; then
                    echo "Missing value for --platform" >&2
                    exit 2
                fi
                INSTALL_PLATFORM="$2"
                NONINTERACTIVE=1
                shift 2
                ;;
            --platform=*)
                INSTALL_PLATFORM="${1#*=}"
                NONINTERACTIVE=1
                shift
                ;;
            --module)
                if [ "$#" -lt 2 ]; then
                    echo "Missing value for --module" >&2
                    exit 2
                fi
                INSTALL_MODULE="$2"
                NONINTERACTIVE=1
                shift 2
                ;;
            --module=*)
                INSTALL_MODULE="${1#*=}"
                NONINTERACTIVE=1
                shift
                ;;
            --reboot)
                if [ "$#" -lt 2 ]; then
                    echo "Missing value for --reboot" >&2
                    exit 2
                fi
                REBOOT_AFTER_INSTALL="$2"
                NONINTERACTIVE=1
                shift 2
                ;;
            --reboot=*)
                REBOOT_AFTER_INSTALL="${1#*=}"
                NONINTERACTIVE=1
                shift
                ;;
            *)
                echo "Unknown option: $1" >&2
                usage >&2
                exit 2
                ;;
            esac
    done

    if [ -n "$REBOOT_AFTER_INSTALL" ]; then
        reboot_choice="$REBOOT_AFTER_INSTALL"
        REBOOT_AFTER_INSTALL="$(normalize_reboot "$reboot_choice")" || {
            echo "Cannot Support this --reboot value:$reboot_choice" >&2
            exit 2
        }
    fi
}

is_agx_supported_l4t() {
    case "${L4T_VER:-}" in
        36.4.3|36.4.4|36.5|36.5.*)
            return 0
            ;;
        *)
            return 1
            ;;
    esac
}

is_supported_l4t() {
    case "${L4T_VER:-}" in
        36.4.0|36.4.3|36.4.4|36.5|36.5.*)
            return 0
            ;;
        *)
            return 1
            ;;
    esac
}

validate_l4t_version() {
    if [ -z "${L4T_VER:-}" ]; then
        fail "Cannot detect NVIDIA L4T version from nvidia-l4t-core."
    fi
    if ! is_supported_l4t; then
        fail "Cannot Support this L4T version:$L4T_VER. Supported L4T versions: 36.4.0, 36.4.3, 36.4.4, 36.5.x."
    fi
}

prompt_continue() {
    local choice
    cat << EOF
****** TechNexion Camera Driver Installation ******
This installation is easy to install TechnNexion Camera Drivers for Nvidia
Jetson Development Kits. Before start to install camera driver,
You should BACKUP your image to avoid any file you lost while installing process.
EOF
    read -p "Do you want to continue?[Y/n]" choice
    if [ -z "$choice" ] || [ "$choice" = "y" ] || [ "$choice" = "Y" ]; then
        return 0
    fi
    return 1
}

prompt_module() {
cat << EOF
Select modules:
    [1]: TEVS: TEVS Series MIPI Cameras with TEV-RPI22 Adaptor
    [2]: VLS-FPD3: VLS-FPD3 Series Cameras with VL-FPD3-8CAM-RPI22 Adaptor
    [3]: VLS-GM2: VLS-GM2 Series Cameras with VL-GM2-8CAM-RPI22 Adaptor
EOF
    read -p "Which modules do you select?[default:1]" INSTALL_MODULE
    INSTALL_MODULE=${INSTALL_MODULE:-"1"}
}

resolve_platform() {
    local platform_choice

    if [ "$NONINTERACTIVE" = "1" ]; then
        platform_choice=${INSTALL_PLATFORM:-"orin_nano"}
        INSTALL_PLATFORM="$(normalize_platform "$platform_choice")" || {
            echo "Cannot Support this Platform:$platform_choice" >&2
            exit 27
        }
    elif is_agx_supported_l4t; then
cat << EOF
Select platform:
    [1]: Jetson Orin Nano
    [2]: Jetson AGX Orin (only support JP6.2.x)
EOF
        read -p "Which platform do you select?[default:1]" platform_choice
        platform_choice=${platform_choice:-"1"}
        INSTALL_PLATFORM="$(normalize_platform "$platform_choice")" || {
            echo "Cannot Support this Platform:$platform_choice" >&2
            exit 27
        }
    else
        INSTALL_PLATFORM="orin_nano"
    fi

    if [ "$INSTALL_PLATFORM" = "agx_orin" ] && ! is_agx_supported_l4t; then
        echo "Cannot Support this Platform:$INSTALL_PLATFORM for L4T $L4T_VER" >&2
        exit 27
    fi
}

resolve_module() {
    if [ "$NONINTERACTIVE" = "1" ] && [ -z "$INSTALL_MODULE" ]; then
        echo "Missing required --module for non-interactive installation." >&2
        exit 2
    fi

    if [ "$INSTALL_PLATFORM" = "agx_orin" ]; then
        INSTALL_MODULE=${INSTALL_MODULE:-"vls-gm2"}
    elif [ "$NONINTERACTIVE" != "1" ]; then
        prompt_module
    fi

    INSTALL_MODULE="$(normalize_module "$INSTALL_MODULE")" || {
        echo "Cannot Support this Module:$INSTALL_MODULE" >&2
        exit 27
    }
}

validate_selection() {
    if [ "$INSTALL_PLATFORM" = "agx_orin" ] && [ "$INSTALL_MODULE" != "vls-gm2" ]; then
        echo "Cannot Support this Module:$INSTALL_MODULE on AGX Orin" >&2
        exit 27
    fi
}

require_package_file() {
    if [ ! -f "$1" ]; then
        fail "Required package file is missing: $1"
    fi
}

validate_common_module_files() {
    require_package_file "technexion/tevs/tevs.ko"
    require_package_file "technexion/vls3/vls3.ko"
    require_package_file "technexion/maxim-gmsl/max96724.ko"
    require_package_file "technexion/maxim-gmsl/max96716a.ko"
    require_package_file "technexion/maxim-gmsl/max96717.ko"
}

validate_orin_nano_dtbo_files() {
    require_package_file "dtbs/tegra234-p3767-camera-p3768-tevs-dual.dtbo"
    require_package_file "dtbs/tegra234-p3767-camera-p3768-vls3.dtbo"
    require_package_file "dtbs/tegra234-p3767-camera-p3768-vls-gm2.dtbo"
    require_package_file "dtbs/tegra234-p3767-camera-p3768-vls-gm2-fsync.dtbo"
    require_package_file "dtbs/tegra234-p3767-camera-p3768-vls-gm2-tunnel.dtbo"
    require_package_file "dtbs/tegra234-p3767-camera-p3768-vls-gm2-tunnel-fsync.dtbo"
    require_package_file "dtbs/tegra234-p3767-camera-p3768-vls-gm2-fsync-external.dtbo"
}

validate_agx_orin_dtbo_files() {
    require_package_file "dtbs/tegra234-p3737-camera-vls-gm2-overlay.dtbo"
    require_package_file "dtbs/tegra234-p3737-camera-vls-gm2-fsync-overlay.dtbo"
    require_package_file "dtbs/tegra234-p3737-camera-vls-gm2-tunnel-overlay.dtbo"
    require_package_file "dtbs/tegra234-p3737-camera-vls-gm2-tunnel-fsync-overlay.dtbo"
    require_package_file "dtbs/tegra234-p3737-camera-vls-gm2-fsync-external-overlay.dtbo"
}

validate_required_files() {
    validate_common_module_files

    if [ "$INSTALL_PLATFORM" = "agx_orin" ]; then
        validate_agx_orin_dtbo_files
    else
        validate_orin_nano_dtbo_files
    fi
}

install_common_modules() {
    install_tncam_modules "technexion/tevs" "tevs.ko"
    install_tncam_modules "technexion/vls3" "vls3.ko"
    install_tncam_modules "technexion/maxim-gmsl" "max96724.ko"
    install_tncam_modules "technexion/maxim-gmsl" "max96716a.ko"
    install_tncam_modules "technexion/maxim-gmsl" "max96717.ko"
    refresh_tncam_modules
    remove_legacy_tncam_service
}

install_orin_nano_dtbo_files() {
    install_tncam_dtbo "tegra234-p3767-camera-p3768-tevs-dual.dtbo"
    install_tncam_dtbo "tegra234-p3767-camera-p3768-vls3.dtbo"
    install_tncam_dtbo "tegra234-p3767-camera-p3768-vls-gm2.dtbo"
    install_tncam_dtbo "tegra234-p3767-camera-p3768-vls-gm2-fsync.dtbo"
    install_tncam_dtbo "tegra234-p3767-camera-p3768-vls-gm2-tunnel.dtbo"
    install_tncam_dtbo "tegra234-p3767-camera-p3768-vls-gm2-tunnel-fsync.dtbo"
    install_tncam_dtbo "tegra234-p3767-camera-p3768-vls-gm2-fsync-external.dtbo"
}

install_agx_orin_dtbo_files() {
    install_tncam_dtbo "tegra234-p3737-camera-vls-gm2-overlay.dtbo"
    install_tncam_dtbo "tegra234-p3737-camera-vls-gm2-fsync-overlay.dtbo"
    install_tncam_dtbo "tegra234-p3737-camera-vls-gm2-tunnel-overlay.dtbo"
    install_tncam_dtbo "tegra234-p3737-camera-vls-gm2-tunnel-fsync-overlay.dtbo"
    install_tncam_dtbo "tegra234-p3737-camera-vls-gm2-fsync-external-overlay.dtbo"
}

configure_orin_nano_module() {
    local module
    module="$INSTALL_MODULE"

    if [ "$module" = "tevs" ]; then
        echo "Installing module: tevs"
        DTB_FILE="tegra234-p3767-camera-p3768-tevs-dual.dtbo"
        remove_tncam_confs
        add_tncam_conf "tevs"
        set_default_conf "tevs"
    elif [ "$module" = "vls-fpd3" ]; then
        echo "Installing module: vls-fpd3"
        DTB_FILE="tegra234-p3767-camera-p3768-vls3.dtbo"
        remove_tncam_confs
        add_tncam_conf "vl316-vls"
        set_default_conf "vl316-vls"
    elif [ "$module" = "vls-gm2" ]; then
        echo "Installing module: vls-gm2"
        DTB_FILE="tegra234-p3767-camera-p3768-vls-gm2.dtbo"
        remove_tncam_confs
        add_tncam_conf "vls-gm2"
        module="vls-gm2-fsync"
        echo "Installing module: $module"
        DTB_FILE="tegra234-p3767-camera-p3768-$module.dtbo"
        add_tncam_conf "$module"
        module="vls-gm2-tunnel"
        echo "Installing module: $module"
        DTB_FILE="tegra234-p3767-camera-p3768-$module.dtbo"
        add_tncam_conf "$module"
        module="vls-gm2-tunnel-fsync"
        echo "Installing module: $module"
        DTB_FILE="tegra234-p3767-camera-p3768-$module.dtbo"
        add_tncam_conf "$module"
        module="vls-gm2-fsync-external"
        echo "Installing module: $module"
        DTB_FILE="tegra234-p3767-camera-p3768-$module.dtbo"
        add_tncam_conf "$module"
        set_default_conf "vls-gm2"
    else
        echo "Cannot Support this Module:$INSTALL_MODULE" >&2
        exit 27
    fi
}

configure_agx_orin_module() {
    local module
    module="$INSTALL_MODULE"
    if [ "$module" != "vls-gm2" ]; then
        echo "Cannot Support this Module:$module on AGX Orin" >&2
        exit 27
    fi

    echo "Installing module: vls-gm2"
    DTB_FILE="tegra234-p3737-camera-vls-gm2-overlay.dtbo"
    remove_tncam_confs
    add_tncam_conf "vls-gm2"
    module="vls-gm2-fsync"
    echo "Installing module: $module"
    DTB_FILE="tegra234-p3737-camera-$module-overlay.dtbo"
    add_tncam_conf "$module"
    module="vls-gm2-tunnel"
    echo "Installing module: $module"
    DTB_FILE="tegra234-p3737-camera-$module-overlay.dtbo"
    add_tncam_conf "$module"
    module="vls-gm2-tunnel-fsync"
    echo "Installing module: $module"
    DTB_FILE="tegra234-p3737-camera-$module-overlay.dtbo"
    add_tncam_conf "$module"
    module="vls-gm2-fsync-external"
    echo "Installing module: $module"
    DTB_FILE="tegra234-p3737-camera-$module-overlay.dtbo"
    add_tncam_conf "$module"
    set_default_conf "vls-gm2"
}

maybe_reboot() {
    local choice
    local reboot_choice

    echo "Finish Camera Driver Installation. Return Code:[0]"
    echo "You should Reboot Device to enable $INSTALL_MODULE Cameras."

    if [ "$NONINTERACTIVE" = "1" ]; then
        reboot_choice="$(printf '%s' "${REBOOT_AFTER_INSTALL:-no}" | tr '[:upper:]' '[:lower:]')"
        if [ "$reboot_choice" = "yes" ] || [ "$reboot_choice" = "y" ]; then
            echo "Rebooting...."
            sudo_or_fail reboot
        else
            echo "Reboot manually later."
        fi
        return
    fi

    read -p "Do you want to reboot now?[y/N]" choice
    if [ "$choice" = "y" ] || [ "$choice" = "Y" ]; then
        echo "Rebooting...."
        sudo_or_fail reboot
    else
        echo "Reboot manually later."
    fi
}

run_installation() {
    if [ "$NONINTERACTIVE" != "1" ]; then
        if ! prompt_continue; then
            echo "Installation aborted."
            exit 0
        fi
    fi

    echo "Continuing with the installation..."
    get_kernel_version
    get_l4t_version
    validate_l4t_version
    resolve_platform
    resolve_module
    validate_selection
    validate_required_files
    check_default_dtb
    sudo_or_fail -v
    install_common_modules

    if [ "$INSTALL_PLATFORM" = "agx_orin" ]; then
        install_agx_orin_dtbo_files
        configure_agx_orin_module
    else
        install_orin_nano_dtbo_files
        configure_orin_nano_module
    fi

    maybe_reboot
}

check_default_dtb() {
    if [ -n "$1" ] && [ -f "$BOOT_DTB_DIR/$1" ]; then
        echo "File $1 exists in $BOOT_DTB_DIR"
    else
        if [ -n "$1" ]; then
            echo "File $1 does not exist in $BOOT_DTB_DIR"
        else
            echo "No FDT found in reference label. Detecting DTB from $BOOT_DTB_DIR"
        fi

        # Find the first file in the directory
        EXISTING_FILE=$(ls -1 "$BOOT_DTB_DIR" 2>/dev/null | head -n 1)

        if [ -n "$EXISTING_FILE" ]; then
            REF_DTB_FILE="$EXISTING_FILE"
            echo "Using detected DTB file: $REF_DTB_FILE"
        else
            echo "No DTB files found in $BOOT_DTB_DIR"
        fi
    fi
}

add_tncam_conf() {
    echo "Add TN-CAM Configuration for modules: $1"
    # sudo sed -i "s/DEFAULT .*/DEFAULT tn-$1/" $EXTLINUX_CONF
    CHECKED_CONF=$(grep -n "^[[:space:]]*LABEL[[:space:]]\+tn-$1[[:space:]]*$" "$EXTLINUX_CONF" || true)
    if [ -n "$CHECKED_CONF" ]; then
        echo "TN-CAM Configuration for modules $1 already exists, skip append."
        return 0
    fi
    TEVS_CONF=$(awk "/LABEL $REF_LEBEL_NAME/,/APPEND /" "$EXTLINUX_CONF" | \
    sed -e "s|LINUX /boot/Image|&\n      FDT /boot/dtb/$REF_DTB_FILE\n      OVERLAYS /boot/$DTB_FILE|" | \
    sed "s/LABEL primary/LABEL tn-$1/")
    if [ -z "$TEVS_CONF" ]; then
        echo "Cannot find reference LABEL $REF_LEBEL_NAME in $EXTLINUX_CONF"
        exit 27
    fi
    printf '\n\n%s\n' "$TEVS_CONF" | sudo tee -a "$EXTLINUX_CONF" > /dev/null || fail "Failed to append TN-CAM configuration to $EXTLINUX_CONF"
}

remove_tncam_confs() {
    echo "Remove old TN-CAM boot configurations..."
    TMP_EXTLINUX_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_EXTLINUX_CONF"
    if [ ! -s "$TMP_EXTLINUX_CONF" ]; then
        echo "Generated empty $EXTLINUX_CONF, abort TN-CAM configuration cleanup."
        rm -f "$TMP_EXTLINUX_CONF"
        exit 27
    fi
    sudo tee "$EXTLINUX_CONF" < "$TMP_EXTLINUX_CONF" > /dev/null || {
        rm -f "$TMP_EXTLINUX_CONF"
        fail "Failed to rewrite $EXTLINUX_CONF"
    }
    rm -f "$TMP_EXTLINUX_CONF"
}

set_default_conf() {
    echo "Set Default Configuration for modules: $1"
    sudo_or_fail sed -i "s/^[[:space:]]*DEFAULT[[:space:]].*/DEFAULT tn-$1/" "$EXTLINUX_CONF"
}

install_tncam_dtbo() {
    echo "Install TN-CAM DTBO file: $1"
    if [ ! -f "dtbs/$1" ]; then
        echo "ERROR: TN-CAM DTBO file is missing: dtbs/$1" >&2
        exit 27
    fi
    TGT_DTB_FILE="/boot/"$1
    sudo_or_fail cp "dtbs/$1" "$TGT_DTB_FILE"
    echo "Installed TN-CAM DTB file Done."
}

get_kernel_version() {
    TEGRA_KERNEL_VER=$(uname -r)
}

get_l4t_version() {
	L4T_VER=$(dpkg-query -W -f='${Version}' nvidia-l4t-core 2>/dev/null \
			| cut -d- -f1)
}

install_tncam_modules() {
    echo "Install TN-CAM modules: $2"
    if [ ! -f "$1/$2" ]; then
        fail "TN-CAM module file is missing: $1/$2"
    fi
    TGT_MODULE_FILE="/lib/modules/$TEGRA_KERNEL_VER/updates/drivers/media/i2c/$1"
    sudo_or_fail mkdir -p "$TGT_MODULE_FILE"
    sudo_or_fail cp "$1/$2" "$TGT_MODULE_FILE/"
}

refresh_tncam_modules() {
    echo "Update module dependency for kernel: $TEGRA_KERNEL_VER"
    sudo_or_fail depmod -a "$TEGRA_KERNEL_VER"
}

remove_legacy_tncam_service() {
    if [ -f "$LEGACY_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_or_fail rm -f "$LEGACY_SERVICE_FILE"
        if command -v systemctl >/dev/null 2>&1; then
            sudo_or_fail systemctl daemon-reload
        fi
    fi
}

parse_args "$@"
run_installation
