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

VERSION="6.2.0"
PACKAGE="isaac-sim-assets-complete-${VERSION}"
BASE_URL="${ISAAC_SIM_ASSETS_BASE_URL:-https://downloads.isaacsim.nvidia.com}"
OUT_DIR="${PWD}"
EXTRACT=true
DOWNLOAD_ONLY=false
FORCE_CONCAT=false
FORCE_EXTRACT=false

FILES=(
    "isaac-sim-assets-complete-6.2.0.001.zip"
    "isaac-sim-assets-complete-6.2.0.002.zip"
    "isaac-sim-assets-complete-6.2.0.003.zip"
    "isaac-sim-assets-complete-6.2.0.004.zip"
    "isaac-sim-assets-complete-6.2.0.005.zip"
)

MD5S=(
    "5206aa3b08a86ac8e0b7fed086db26f1"
    "a5f940171a402453ec146cdd2095119c"
    "397d0bdbc36d20a98eca983c871b8cb8"
    "be3a88b15d83ac8615c51520c2f2f95a"
    "5f47abe66975cc9929dae78b7184b284"
)

SIZES=(
    "16106127360"
    "16106127360"
    "16106127360"
    "16106127360"
    "710199986"
)

ARCHIVE_SIZE="65134709426"

usage() {
    cat <<EOF
Download, verify, combine, and extract the Isaac Sim ${VERSION} assets pack.

Usage:
  ${0##*/} [options]

Options:
  -o, --output DIR       Download and extract in DIR. Defaults to the current directory.
      --url-base URL     Override download base URL. Defaults to ${BASE_URL}.
      --download-only    Download and verify the split ZIP files, then stop.
      --no-extract       Download, verify, and combine the ZIP files, then stop.
      --force-concat     Rebuild the combined ZIP even if it already exists.
      --force-extract    Extract even if the extraction marker already exists.
  -h, --help             Show this help.

Examples:
  curl -fsSL https://downloads.isaacsim.nvidia.com/download-isaac-sim-assets-${VERSION}.sh | bash
  curl -fsSL https://downloads.isaacsim.nvidia.com/download-isaac-sim-assets-${VERSION}.sh | bash -s -- --output "\$HOME/Downloads/isaacsim-assets"
EOF
}

log() {
    printf '\n==> %s\n' "$*"
}

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

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

file_size() {
    if stat -c '%s' "$1" >/dev/null 2>&1; then
        stat -c '%s' "$1"
    else
        stat -f '%z' "$1"
    fi
}

file_md5() {
    if command -v md5sum >/dev/null 2>&1; then
        md5sum "$1" | awk '{print $1}'
    elif command -v md5 >/dev/null 2>&1; then
        md5 -q "$1"
    else
        die "Missing required command: md5sum or md5"
    fi
}

md5_ok() {
    local path="$1"
    local expected="$2"

    [[ -f "${path}" ]] && [[ "$(file_md5 "${path}")" == "${expected}" ]]
}

bytes_human() {
    local bytes="$1"
    if command -v numfmt >/dev/null 2>&1; then
        numfmt --to=iec --suffix=B "${bytes}"
    else
        printf '%s bytes' "${bytes}"
    fi
}

parse_args() {
    while (($#)); do
        case "$1" in
            -o|--output)
                (($# >= 2)) || die "Missing value for $1"
                OUT_DIR="$2"
                shift 2
                ;;
            --url-base)
                (($# >= 2)) || die "Missing value for $1"
                BASE_URL="$2"
                shift 2
                ;;
            --download-only)
                DOWNLOAD_ONLY=true
                EXTRACT=false
                shift
                ;;
            --no-extract)
                EXTRACT=false
                shift
                ;;
            --force-concat)
                FORCE_CONCAT=true
                shift
                ;;
            --force-extract)
                FORCE_EXTRACT=true
                shift
                ;;
            -h|--help)
                usage
                exit 0
                ;;
            *)
                die "Unknown option: $1"
                ;;
        esac
    done
}

download_parts() {
    need_cmd aria2c
    need_cmd awk

    for i in "${!FILES[@]}"; do
        local file="${FILES[$i]}"
        local md5="${MD5S[$i]}"
        local expected_size="${SIZES[$i]}"
        local path="${OUT_DIR}/${file}"
        local url="${BASE_URL%/}/${file}"

        if md5_ok "${path}" "${md5}"; then
            log "Using existing verified ${file}"
            rm -f "${path}.aria2"
            continue
        fi

        if [[ -f "${path}" ]]; then
            local current_size
            current_size="$(file_size "${path}")"
            if (( current_size >= expected_size )); then
                log "Removing invalid complete ${file}; expected MD5 ${md5}"
                rm -f "${path}" "${path}.aria2"
            else
                log "Resuming ${file} at $(bytes_human "${current_size}") of $(bytes_human "${expected_size}")"
            fi
        else
            log "Downloading ${file}"
        fi

        aria2c \
            --continue=true \
            --auto-file-renaming=false \
            --allow-overwrite=true \
            --max-connection-per-server=8 \
            --split=8 \
            --min-split-size=10M \
            --max-tries=20 \
            --retry-wait=10 \
            --connect-timeout=30 \
            --timeout=60 \
            --summary-interval=60 \
            --dir="${OUT_DIR}" \
            --out="${file}" \
            --checksum="md5=${md5}" \
            "${url}"

        md5_ok "${path}" "${md5}" || die "Checksum failed after download: ${file}"
        rm -f "${path}.aria2"
        log "Verified ${file}"
    done
}

combine_parts() {
    local archive="${OUT_DIR}/${PACKAGE}.zip"
    local tmp_archive="${archive}.part"

    if [[ "${FORCE_CONCAT}" == "false" && -f "${archive}" && "$(file_size "${archive}")" == "${ARCHIVE_SIZE}" ]]; then
        log "Using existing combined archive ${archive}"
        return
    fi

    log "Combining split ZIP files into ${archive}"
    rm -f "${tmp_archive}"
    local parts=()
    for file in "${FILES[@]}"; do
        parts+=("${OUT_DIR}/${file}")
    done
    cat "${parts[@]}" > "${tmp_archive}"

    [[ "$(file_size "${tmp_archive}")" == "${ARCHIVE_SIZE}" ]] || die "Combined archive has the wrong size"
    mv -f "${tmp_archive}" "${archive}"
    log "Created ${archive}"
}

extract_archive() {
    local archive="${OUT_DIR}/${PACKAGE}.zip"
    local marker="${OUT_DIR}/.${PACKAGE}.extracted"

    if [[ "${FORCE_EXTRACT}" == "false" && -f "${marker}" ]]; then
        log "Extraction marker exists; skipping extraction. Use --force-extract to extract again."
        return
    fi

    [[ -f "${archive}" ]] || die "Missing combined archive: ${archive}"

    log "Extracting ${archive} into ${OUT_DIR}"
    if command -v unzip >/dev/null 2>&1; then
        unzip -q -o "${archive}" -d "${OUT_DIR}"
    elif command -v 7z >/dev/null 2>&1; then
        7z x -y "-o${OUT_DIR}" "${archive}"
    else
        die "Missing extractor: install unzip or 7z"
    fi

    {
        printf 'package=%s\n' "${PACKAGE}"
        printf 'archive=%s\n' "${archive}"
        printf 'extracted_at=%s\n' "$(date -Iseconds)"
    } > "${marker}"
    log "Extraction complete"
}

main() {
    parse_args "$@"

    mkdir -p "${OUT_DIR}"
    OUT_DIR="$(cd "${OUT_DIR}" && pwd -P)"
    BASE_URL="${BASE_URL%/}"

    log "Output directory: ${OUT_DIR}"
    log "Download base URL: ${BASE_URL}"

    download_parts

    if [[ "${DOWNLOAD_ONLY}" == "true" ]]; then
        log "Download verification complete"
        exit 0
    fi

    combine_parts

    if [[ "${EXTRACT}" == "true" ]]; then
        extract_archive
    else
        log "Done. Extraction skipped."
    fi
}

main "$@"