#!/bin/bash
# This script sets up the repository for INEX software installation

cleanup() {
    if [ -f "/etc/apt/sources.list.d/inex.list" ]; then
        $SUDO rm -f "/etc/apt/sources.list.d/inex.list"
    fi
    if [ -f "/etc/apt/sources.list.d/ubuntu-mirrors.list" ]; then
        $SUDO rm -f "/etc/apt/sources.list.d/ubuntu-mirrors.list"
    fi
    if [ -f "/etc/apt/sources.list.d/nvidia-mirrors.list" ]; then
        $SUDO rm -f "/etc/apt/sources.list.d/nvidia-mirrors.list"
    fi
}
trap cleanup ERR


# Exit immediately if a command exits with a non-zero status
# Treat unset variables as an error when substituting
set -eaux

# Set the production repository address
PROD_REPO_ADDRESS="https://repo.iz-cloud.com"

# Check if REPO_ADDRESS is set, otherwise use the default
REPO_ADDRESS="${REPO_ADDRESS:-$PROD_REPO_ADDRESS}"

# Validate REPO_ADDRESS as a URL format for apt sources
case "$REPO_ADDRESS" in
    http://*|https://*) ;;
    *)
        echo "Error: REPO_ADDRESS ($REPO_ADDRESS) is not a valid URL. Use the format: http://<address> or https://<address>"
        exit 1
        ;;
esac

# Display the repository address
echo "Using repository address: $REPO_ADDRESS"

# Set paths for the GPG key and INEX repository
GPG_KEY_URL="$REPO_ADDRESS/zamir.pub"
GPG_KEY_PATH="/etc/apt/keyrings/inex.asc"
INEX_REPO_PATH="$REPO_ADDRESS/inex"

# Check if running in Docker container
IN_CONTAINER=0
if [ -f /.dockerenv ] || [ -f /run/.containerenv ] || grep -q docker /proc/1/cgroup; then
    IN_CONTAINER=1
    echo "Warning: Script is running inside Docker container"
fi

# Set DIST, PLATFORM, and CODENAME based on Ubuntu version
UBUNTU_VERSION=$(grep 'VERSION_ID=' /etc/os-release | cut -d '"' -f 2)
CODENAME=$(grep -E '(VERSION_CODENAME|UBUNTU_CODENAME)=' /etc/os-release | head -n1 | cut -d'=' -f2)

# Check and set Jetson versions
if [ -z "${JETSON_JETPACK:-}" ]; then
    if JETPACK_VER=$(dpkg-query -W -f='${Version}' nvidia-jetpack 2>/dev/null); then
        JETSON_JETPACK=$(echo "$JETPACK_VER" | grep -oE '^[0-9]+\.[0-9]+\.[0-9]+' || echo "")
    else
        JETPACK_VER=$(apt-cache policy nvidia-jetpack 2>/dev/null | grep -oE '[0-9]+\.[0-9]+\.[0-9]+' | head -n1 || echo "")
        JETSON_JETPACK=$(echo "$JETPACK_VER" | grep -oE '^[0-9]+\.[0-9]+\.[0-9]+' || echo "")  # Add version processing
    fi
fi

if [ -z "${JETSON_L4T:-}" ]; then
    # Try to get L4T version from package, fallback to policy check
    if L4T_VER=$(dpkg-query -W -f='${Version}' nvidia-l4t-core 2>/dev/null); then
        JETSON_L4T=$(echo "$L4T_VER" | grep -oE '^[0-9]+\.[0-9]+\.[0-9]+' || echo "")
    else
        L4T_VER=$(apt-cache policy nvidia-l4t-core 2>/dev/null | grep -oE '[0-9]+\.[0-9]+\.[0-9]+' | head -n1 || echo "")
        JETSON_L4T=${L4T_VER}
    fi
fi

# Debug output
echo "Detected JetPack version: ${JETSON_JETPACK:-not found}"
echo "Detected L4T version: ${JETSON_L4T:-not found}"

# Exit if jetpack version is not detected
if [ -z "${JETSON_JETPACK}" ] && [ -z "${JETSON_L4T}" ]; then
    echo "Error: Failed to detect required Jetson versions."
    echo "Please ensure either nvidia-jetpack or nvidia-l4t-core packages are installed or available."
    exit 1
fi

# Set DIST based on JETSON_L4T version
DIST="r$(echo "$JETSON_L4T" | cut -d'.' -f1,2)"

# Debug output for DIST
echo "Using L4T distribution: $DIST"

if [ "$IN_CONTAINER" -eq 0 ] && [ -f "/proc/device-tree/nvidia,dtsfilename" ]; then
    DTS_FILENAME=$(tr -d '\0' < "/proc/device-tree/nvidia,dtsfilename")
    TEGRA_NUM=$(basename "$DTS_FILENAME" | grep -oE 'tegra[0-9]+' | grep -oE '[0-9]+')
    if [ -n "$TEGRA_NUM" ]; then
        PLATFORM="t${TEGRA_NUM}"
        echo "Detected platform from device tree: $PLATFORM"
    else
        echo "Warning: Could not detect platform from device tree"
    fi
fi

# Check for root privileges and set SUDO command
CAN_ROOT=""
SUDO=""
if [ "$(id -u)" = "0" ]; then
    CAN_ROOT=1
    SUDO=""
elif command -v sudo >/dev/null 2>&1; then
    CAN_ROOT=1
    SUDO="sudo"
elif command -v doas >/dev/null 2>&1; then
    CAN_ROOT=1
    SUDO="doas"
fi

# Exit if root privileges are not available
if [ "$CAN_ROOT" != "1" ]; then
    echo "This installer needs to run commands as root."
    echo "We tried looking for 'sudo' and 'doas', but couldn't find them."
    echo "Either re-run this script as root, or set up sudo/doas."
    exit 1
fi

# Check for wget or curl and define download command
DOWNLOAD_CMD=""
if command -v wget >/dev/null 2>&1; then
    DOWNLOAD_CMD="wget -q --tries=3 --timeout=10 -O"
elif command -v curl >/dev/null 2>&1; then
    DOWNLOAD_CMD="curl -s --retry 3 --max-time 10 -o"
else
    echo "Error: Neither wget nor curl is installed. Please install one of these utilities."
    exit 1
fi

# Create keyrings directory and download INEX public key
$SUDO mkdir -p /etc/apt/keyrings
if ! $SUDO $DOWNLOAD_CMD "$GPG_KEY_PATH" "$GPG_KEY_URL"; then
    echo "Error: Failed to download INEX public key from $GPG_KEY_URL."
    exit 1
fi
if ! [ -s "$GPG_KEY_PATH" ]; then
    echo "Error: Downloaded GPG key file is empty or missing"
    exit 1
fi
$SUDO chmod a+r "$GPG_KEY_PATH"

# Add INEX repository to sources list
echo "deb [arch=arm64 signed-by=$GPG_KEY_PATH] $INEX_REPO_PATH $DIST main" | $SUDO tee /etc/apt/sources.list.d/inex.list > /dev/null

PIN_FILE="/etc/apt/preferences.d/inex.pref"

echo "Creating apt pinning file: $PIN_FILE"

sudo tee "$PIN_FILE" > /dev/null <<EOF
Package: *
Pin: release o=inex $DIST
Pin-Priority: 700
EOF

echo "APT pinning configuration added at $PIN_FILE:"

# Add Ubuntu mirrors to sources list
for dist in "$CODENAME" "${CODENAME}-updates" "${CODENAME}-security" "${CODENAME}-backports"; do
    echo "deb [arch=arm64 signed-by=$GPG_KEY_PATH] $REPO_ADDRESS/mirrors $dist main"
done | $SUDO tee /etc/apt/sources.list.d/ubuntu-mirrors.list > /dev/null

# Add NVIDIA mirrors to sources list
echo "deb [arch=arm64 signed-by=$GPG_KEY_PATH] $REPO_ADDRESS/mirrors/jetson/common $DIST main" | $SUDO tee /etc/apt/sources.list.d/nvidia-mirrors.list > /dev/null

if [ -n "${PLATFORM:-}" ]; then
    echo "deb [arch=arm64 signed-by=$GPG_KEY_PATH] $REPO_ADDRESS/mirrors/jetson/$PLATFORM $DIST main" | $SUDO tee -a /etc/apt/sources.list.d/nvidia-mirrors.list > /dev/null
fi

# Add FFmpeg mirror for L4T 36.x.x
if [[ "${JETSON_L4T}" =~ ^36\. ]]; then
    echo "deb [arch=arm64 signed-by=$GPG_KEY_PATH] $REPO_ADDRESS/mirrors/jetson/ffmpeg $DIST main" | $SUDO tee -a /etc/apt/sources.list.d/nvidia-mirrors.list > /dev/null
fi

# Update package lists with error handling
if ! $SUDO apt-get update; then
    echo "Error: Failed to update package lists. Check your network connection or repository configuration."
    exit 1
fi

echo "Repository setup complete."
