#!/bin/bash
# Update upstream source script for CUDA packages
#
# This script:
# 1. Reads debian/changelog and extracts the upsteam source name and CUDA major and minor versions
# 2. Finds the latest redistrib JSON for that CUDA minor version
# 3. Downloads and renames architecture-specific tarballs (amd64, arm64)
# 4. Generates the dummy primary tarball
# 5. Updates debian/cuda-version
# 6. Runs dch

set -e

# Save initial working directory
ORIGINAL_PWD="$PWD"

# Trap to restore PWD on exit
trap 'cd "$ORIGINAL_PWD"' EXIT

REDIST_BASE_URL="https://developer.download.nvidia.com/compute/cuda/redist"

# Get package directory (script is in debian/scripts/)
PACKAGE_DIR="$(cd "$(dirname "$0")/../.." && pwd)"
TARBALL_DIR="$(dirname "$PACKAGE_DIR")"

log_info() {
    echo "INFO: $*" >&2
}

log_error() {
    echo "ERROR: $*" >&2
}

log_debug() {
    if [ "${VERBOSE:-0}" -eq 1 ]; then
        echo "DEBUG: $*" >&2
    fi
}

# Parse package name from changelog
parse_changelog() {
    PACKAGE_NAME=$(dpkg-parsechangelog -l "$PACKAGE_DIR/debian/changelog" -S Source)
    PREVIOUS_COMPONENT_VERSION=$(dpkg-parsechangelog -l "$PACKAGE_DIR/debian/changelog" -S Version | cut -d- -f1)

    log_info "Package name: $PACKAGE_NAME"
    log_info "Current upstream version: $PREVIOUS_COMPONENT_VERSION"
}

# Extract upstream name and CUDA version from package name
parse_package_name() {
    if [[ ! "$PACKAGE_NAME" =~ ^(.+)-([0-9]+)-([0-9]+)$ ]]; then
        log_error "Package name '$PACKAGE_NAME' doesn't match pattern 'name-MAJOR-MINOR'"
        exit 1
    fi

    UPSTREAM_NAME="${BASH_REMATCH[1]}"
    CUDA_MAJOR="${BASH_REMATCH[2]}"
    CUDA_MINOR="${BASH_REMATCH[3]}"

    # Convert hyphens to underscores for NVIDIA URLs
    UPSTREAM_NAME_UNDERSCORE="${UPSTREAM_NAME//-/_}"

    log_info "Upstream name: $UPSTREAM_NAME"
    log_info "CUDA version: $CUDA_MAJOR.$CUDA_MINOR"
}

# Find the latest redistrib JSON for the CUDA version
find_latest_redistrib_json() {
    log_info "Searching for latest CUDA $CUDA_MAJOR.$CUDA_MINOR.X redistrib manifest..."

    local LATEST_VERSION=""
    local LATEST_URL=""

    # Try patch versions from 0 upward
    for PATCH in {0..99}; do
        local VERSION="$CUDA_MAJOR.$CUDA_MINOR.$PATCH"
        local URL="$REDIST_BASE_URL/redistrib_$VERSION.json"

        log_debug "Checking $URL"

        # Use HEAD request to check existence
        if curl -f -s -I -A "Ubuntu-CUDA-Downloader/1.0" --connect-timeout 10 "$URL" > /dev/null 2>&1; then
            log_debug "Found: $URL"
            LATEST_VERSION="$VERSION"
            LATEST_URL="$URL"
        else
            log_debug "Not available: $URL"
            # If we've found at least one and now hitting errors, stop
            if [ -n "$LATEST_VERSION" ]; then
                break
            fi
        fi
    done

    if [ -z "$LATEST_VERSION" ]; then
        log_error "Could not find any redistrib JSON for CUDA $CUDA_MAJOR.$CUDA_MINOR.X"
        exit 1
    fi

    log_info "Found latest redistrib manifest: $LATEST_URL"

    # Download the manifest to disk with full version
    MANIFEST_FILE="$PACKAGE_DIR/debian/redistrib_${LATEST_VERSION}.json"
    curl -f -s -A "Ubuntu-CUDA-Downloader/1.0" -o "$MANIFEST_FILE" "$LATEST_URL"
}

# Download architecture-specific tarball
download_architecture_tarball() {
    local ARCH="$1"

    local PLATFORM
    case "$ARCH" in
        amd64)
            PLATFORM="linux-x86_64"
            ;;
        arm64)
            PLATFORM="linux-sbsa"
            ;;
        *)
            log_error "Unknown architecture: $ARCH"
            exit 1
            ;;
    esac

    local FILENAME="${UPSTREAM_NAME_UNDERSCORE}-${PLATFORM}-${NEW_COMPONENT_VERSION}-archive.tar.xz"
    local URL="$REDIST_BASE_URL/$UPSTREAM_NAME_UNDERSCORE/$PLATFORM/$FILENAME"

    local OUTPUT_FILENAME="${PACKAGE_NAME}_${NEW_COMPONENT_VERSION}.orig-${ARCH}.tar.xz"
    local OUTPUT_PATH="$TARBALL_DIR/$OUTPUT_FILENAME"

    if [ -f "$OUTPUT_PATH" ]; then
        log_info "Tarball already exists: $OUTPUT_FILENAME"
        return
    fi

    log_info "Downloading $ARCH tarball from $URL"

    if ! curl -f -# -L -A "Ubuntu-CUDA-Downloader/1.0" -o "$OUTPUT_PATH" "$URL"; then
        rm -f "$OUTPUT_PATH"
        log_error "Failed to download $ARCH tarball"
        exit 1
    fi

    log_info "Downloaded: $OUTPUT_FILENAME"
}

create_dummy_tarball() {
    local OUTPUT_FILENAME="${PACKAGE_NAME}_${PREVIOUS_COMPONENT_VERSION}.orig.tar.xz"
    local OUTPUT_PATH="$TARBALL_DIR/$OUTPUT_FILENAME"

    if [ -f "$OUTPUT_PATH" ]; then
        log_info "Dummy tarball already exists: $OUTPUT_FILENAME"
        return
    fi

    log_info "Creating dummy tarball with manifest: $OUTPUT_FILENAME"

    # Create tarball with manifest from debian/ directory  
    # Find the redistrib_*.json file
    local MANIFEST_NAME=$(basename "$PACKAGE_DIR"/debian/redistrib_*.json)
    
    # Create reproducible tarball with deterministic options
    tar cJf "$OUTPUT_PATH" \
        --sort=name \
        --mtime='1970-01-01 00:00:00' \
        --owner=0 \
        --group=0 \
        --numeric-owner \
        -C "$PACKAGE_DIR/debian" \
        "$MANIFEST_NAME"

    log_info "Created: $OUTPUT_FILENAME (with manifest)"
}

update_cuda_version_file() {
    local CUDA_VERSION
    CUDA_VERSION=$(jq -r '.release_label' "$MANIFEST_FILE")
    echo "$CUDA_VERSION" > "$PACKAGE_DIR/debian/cuda-version"
    log_info "Updated cuda-version: $CUDA_VERSION"
}

dch_if_needed() {
    if [ "$NEW_COMPONENT_VERSION" != "$PREVIOUS_COMPONENT_VERSION" ]; then
        log_info "Upstream version changed: $PREVIOUS_COMPONENT_VERSION -> $NEW_COMPONENT_VERSION"
        log_info "Running dch to update changelog..."

        cd "$PACKAGE_DIR"
        dch --newversion="${NEW_COMPONENT_VERSION}-0ubuntu1" --distribution=UNRELEASED

        log_info "Changelog updated successfully"
    else
        log_info "Upstream version unchanged; no changelog update needed."
    fi
}

# Main execution
main() {
    parse_changelog
    parse_package_name
    find_latest_redistrib_json

    create_dummy_tarball

    log_info "Package refresh complete!"
}

# Handle --verbose flag
if [ "$1" = "--verbose" ] || [ "$1" = "-v" ]; then
    VERBOSE=1
fi

main "$@"
