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

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ENV_FILE="$SCRIPT_DIR/.env"

HARDEN="${HARDEN:-true}"
while [[ $# -gt 0 ]]; do
    case "$1" in
        --no-harden) HARDEN="false"; shift ;;
        *) echo "ERROR: Unknown option: $1" >&2; exit 1 ;;
    esac
done

NVIDIA_VENDOR="0955"

declare -A PID_TO_MODEL=(
    ["7523"]="Orin Nano"
    ["7323"]="Orin NX 16GB"
    ["7423"]="Orin NX 8GB"
)

declare -A PID_TO_BOARD_CONFIG=(
    ["7523"]="jetson-orin-nano-devkit"
    ["7323"]="jetson-orin-nx-devkit"
    ["7423"]="jetson-orin-nx-devkit"
)

require_env_var() {
    local name="$1"
    local val="${!name:-}"
    if [[ -z "$val" ]]; then
        echo "ERROR: $name is not set in $ENV_FILE" >&2
        exit 1
    fi
}

api_post() {
    local url="$1"; shift
    local response
    response="$(curl -sS -w "\n%{http_code}" -X POST "$url" -H "Content-Type: application/json" "$@")"
    echo "$response"
}

provision_single_device() {
    local device_line="$1"
    local usb_id="$2"
    local board_config="$3"

    echo "[Step 1/5] Registering device with admin API..."
    local reg_response reg_http reg_body
    reg_response="$(api_post "${API_URL}/devices" -H "Authorization: Bearer $TOKEN")"
    reg_http="$(echo "$reg_response" | tail -1)"
    reg_body="$(echo "$reg_response" | sed '$d')"

    if [[ "$reg_http" != "200" && "$reg_http" != "201" ]]; then
        echo "ERROR: Device registration failed (HTTP $reg_http)" >&2
        echo "$reg_body" >&2
        RESULTS+=("$usb_id | FAILED | registration error HTTP $reg_http")
        return 1
    fi

    local serial dev_email dev_password
    serial="$(echo "$reg_body" | jq -r '.serial // .Serial // empty')"
    dev_email="$(echo "$reg_body" | jq -r '.email // .Email // empty')"
    dev_password="$(echo "$reg_body" | jq -r '.password // .Password // empty')"

    if [[ -z "$serial" || -z "$dev_email" || -z "$dev_password" ]]; then
        echo "ERROR: Incomplete response from POST /devices" >&2
        RESULTS+=("$usb_id | FAILED | incomplete API response")
        return 1
    fi

    echo "  Assigned serial: $serial"
    echo "  Email: $dev_email"

    echo "[Step 2/5] Writing device.conf to rootfs staging..."
    local conf_dir="${ROOTFS_DIR}/etc/azaion"
    sudo mkdir -p "$conf_dir"
    local conf_path="${conf_dir}/device.conf"
    printf 'AZAION_DEVICE_EMAIL=%s\nAZAION_DEVICE_PASSWORD=%s\n' "$dev_email" "$dev_password" \
        | sudo tee "$conf_path" > /dev/null
    sudo chmod 600 "$conf_path"
    echo "  Written: $conf_path"

    echo "[Step 3/5] Fusing device (odmfuse.sh)..."
    if ! sudo "$L4T_DIR/odmfuse.sh" --instance "$usb_id" 2>&1; then
        echo "ERROR: Fusing failed for $serial ($usb_id)" >&2
        RESULTS+=("$usb_id | $serial | FAILED | fuse error")
        return 1
    fi

    echo ""
    echo "[$serial] Fuse complete."
    read -rp "  Power-cycle the device and put it back in recovery mode. Press Enter when ready..."

    echo "[Step 4/5] Flashing device (flash.sh)..."
    if ! sudo "$L4T_DIR/flash.sh" "$board_config" "$FLASH_TARGET" --instance "$usb_id" 2>&1; then
        echo "ERROR: Flashing failed for $serial ($usb_id)" >&2
        RESULTS+=("$usb_id | $serial | FAILED | flash error")
        return 1
    fi

    echo ""
    echo "[$serial] Flash complete."
    echo "  >>> Apply sticker with serial: $serial <<<"
    read -rp "  Power-cycle for first boot. Press Enter when done..."

    echo "[Step 5/5] $serial provisioned successfully."
    RESULTS+=("$usb_id | $serial | OK")
}

# --- main ---

if [[ ! -f "$ENV_FILE" ]]; then
    echo "ERROR: $ENV_FILE not found. Copy .env.example to .env and fill in values." >&2
    exit 1
fi

set -a
source "$ENV_FILE"
set +a

for var in ADMIN_EMAIL ADMIN_PASSWORD API_URL LOADER_IMAGE_TAR; do
    require_env_var "$var"
done
API_URL="${API_URL%/}"

RESOURCE_API_URL="${RESOURCE_API_URL:-$API_URL}"
LOADER_DEV_STAGE="${LOADER_DEV_STAGE:-main}"
LOADER_IMAGE="${LOADER_IMAGE:-localhost:5000/loader:arm}"
FLASH_TARGET="${FLASH_TARGET:-nvme0n1p1}"

L4T_VERSION="${L4T_VERSION:-r36.4.4}"
L4T_DIR="${L4T_DIR:-/opt/nvidia/Linux_for_Tegra}"
ROOTFS_DIR="${ROOTFS_DIR:-$L4T_DIR/rootfs}"

export L4T_VERSION L4T_DIR ROOTFS_DIR RESOURCE_API_URL LOADER_DEV_STAGE LOADER_IMAGE LOADER_IMAGE_TAR

echo "=== Installing host dependencies ==="
sudo apt-get update -qq
sudo apt-get install -y usbutils curl jq wget
[[ "$(uname -m)" != "aarch64" ]] && sudo apt-get install -y qemu-user-static binfmt-support
echo ""

echo "=== L4T BSP setup ==="
"$SCRIPT_DIR/ensure_l4t.sh"
echo ""

echo "=== Setting up rootfs (Docker + application) ==="
"$SCRIPT_DIR/setup_rootfs_docker.sh"
echo ""

if [[ "$HARDEN" == "true" ]]; then
    echo "=== Applying security hardening ==="
    "$SCRIPT_DIR/harden_rootfs.sh"
    echo ""
else
    echo "=== Security hardening SKIPPED (--no-harden) ==="
    echo ""
fi

echo "=== Authenticating with admin API ==="
LOGIN_JSON="$(printf '{"email":"%s","password":"%s"}' "$ADMIN_EMAIL" "$ADMIN_PASSWORD")"
LOGIN_RESPONSE="$(api_post "${API_URL}/login" -d "$LOGIN_JSON")"
LOGIN_HTTP="$(echo "$LOGIN_RESPONSE" | tail -1)"
LOGIN_BODY="$(echo "$LOGIN_RESPONSE" | sed '$d')"

if [[ "$LOGIN_HTTP" != "200" ]]; then
    echo "ERROR: Login failed (HTTP $LOGIN_HTTP)" >&2
    echo "$LOGIN_BODY" >&2
    exit 1
fi

TOKEN="$(echo "$LOGIN_BODY" | jq -r '.token // .Token // empty')"
if [[ -z "$TOKEN" ]]; then
    echo "ERROR: No token in login response" >&2
    echo "$LOGIN_BODY" >&2
    exit 1
fi
echo "Authenticated successfully."
echo ""

echo "=== Scanning for Jetson devices in recovery mode ==="
LSUSB_OUTPUT="$(lsusb -d "${NVIDIA_VENDOR}:" 2>/dev/null || true)"

if [[ -z "$LSUSB_OUTPUT" ]]; then
    echo "No Jetson devices found in recovery mode."
    echo "Ensure devices are connected via USB and in recovery mode (hold Force Recovery, press Power)."
    exit 0
fi

DEVICES=()
DEVICE_PIDS=()
DEVICE_MODELS=()
while IFS= read -r line; do
    pid="$(echo "$line" | grep -oP "${NVIDIA_VENDOR}:\K[0-9a-fA-F]+")"
    if [[ -n "${PID_TO_BOARD_CONFIG[$pid]:-}" ]]; then
        DEVICES+=("$line")
        DEVICE_PIDS+=("$pid")
        DEVICE_MODELS+=("${PID_TO_MODEL[$pid]:-Unknown (PID $pid)}")
    fi
done <<< "$LSUSB_OUTPUT"

DEVICE_COUNT="${#DEVICES[@]}"

if [[ "$DEVICE_COUNT" -eq 0 ]]; then
    echo "No supported Jetson devices found."
    echo "Detected NVIDIA USB devices but none matched known Jetson Orin product IDs."
    exit 0
fi

echo ""
echo "Select device(s) to provision."
echo "  one device, e.g. 1"
echo "  some devices, e.g. 1 3 4"
echo "  or all devices: 0"
echo ""
echo "--------------------------------------------"
echo "Connected Jetson devices (recovery mode):"
echo "--------------------------------------------"
for i in "${!DEVICES[@]}"; do
    num=$((i + 1))
    printf "%-3s %-16s %s\n" "$num" "[${DEVICE_MODELS[$i]}]" "${DEVICES[$i]}"
done
echo "--------------------------------------------"
echo "0 - provision all devices"
echo ""

read -rp "Your selection: " SELECTION

SELECTED_INDICES=()
if [[ "$SELECTION" == "0" ]]; then
    for i in "${!DEVICES[@]}"; do
        SELECTED_INDICES+=("$i")
    done
else
    for num in $SELECTION; do
        if [[ "$num" =~ ^[0-9]+$ ]] && (( num >= 1 && num <= DEVICE_COUNT )); then
            SELECTED_INDICES+=("$((num - 1))")
        else
            echo "ERROR: Invalid selection: $num (must be 1-$DEVICE_COUNT or 0 for all)" >&2
            exit 1
        fi
    done
fi

if [[ ${#SELECTED_INDICES[@]} -eq 0 ]]; then
    echo "No devices selected."
    exit 0
fi

echo ""
echo "=== Provisioning ${#SELECTED_INDICES[@]} device(s) ==="
echo ""

RESULTS=()

for idx in "${SELECTED_INDICES[@]}"; do
    DEVICE_LINE="${DEVICES[$idx]}"
    USB_ID="$(echo "$DEVICE_LINE" | grep -oP 'Bus \K[0-9]+')-$(echo "$DEVICE_LINE" | grep -oP 'Device \K[0-9]+')"
    BOARD_CONFIG="${PID_TO_BOARD_CONFIG[${DEVICE_PIDS[$idx]}]:-}"
    if [[ -z "$BOARD_CONFIG" ]]; then
        echo "ERROR: Unknown Jetson product ID: $NVIDIA_VENDOR:${DEVICE_PIDS[$idx]}" >&2
        RESULTS+=("$USB_ID | FAILED | unknown product ID")
        continue
    fi
    echo "--------------------------------------------"
    echo "Device: ${DEVICE_MODELS[$idx]} — $DEVICE_LINE"
    echo "USB instance: $USB_ID"
    echo "Board config: $BOARD_CONFIG"
    echo "--------------------------------------------"

    provision_single_device "$DEVICE_LINE" "$USB_ID" "$BOARD_CONFIG" || true
    echo ""
done

CONF_CLEANUP="${ROOTFS_DIR}/etc/azaion/device.conf"
if [[ -f "$CONF_CLEANUP" ]]; then
    sudo rm -f "$CONF_CLEANUP"
fi

echo ""
echo "========================================"
echo "         Provisioning Summary"
echo "========================================"
printf "%-12s | %-10s | %s\n" "USB ID" "Serial" "Status"
echo "----------------------------------------"
for r in "${RESULTS[@]}"; do
    echo "$r"
done
echo "========================================"