install

#!/usr/bin/env bash

set -euo pipefail

# Configuration
OWNER="NVIDIA"
PROJECT="aicr"
REPO="$OWNER/$PROJECT"
BIN_NAME="aicr"
CHECKSUMS_FILE="aicr_checksums.txt"
REQUIRED_TOOLS=(curl)
INSTALL_DIR="${1:-/usr/local/bin}"

# Colors
COLOR_RED='\033[0;31m'
COLOR_GREEN='\033[0;32m'
COLOR_YELLOW='\033[0;33m'
COLOR_CYAN='\033[0;36m'
COLOR_DIM='\033[2m'
COLOR_RESET='\033[0m'

# ==============================================================================
# Utility Functions
# ==============================================================================

usage() {
  cat <<EOF
Install utility for $BIN_NAME

Usage: $0 [-d install_dir]

Environment Variables:
  GITHUB_TOKEN           GitHub token to avoid API rate limits (no special scopes required)
  AICR_NO_COMPLETIONS    Set to 1 to skip shell completion installation

Examples:
  $0                    # Install to /usr/local/bin
  $0 -d ~/bin           # Install to custom directory
  GITHUB_TOKEN=xyz $0   # Use token to avoid GitHub API rate limits
EOF
  exit 1
}

err() {
  printf "${COLOR_RED}  ✗ %s${COLOR_RESET}\n" "$1" >&2
  exit 1
}

ok() {
  printf "${COLOR_GREEN}  ✓ %s${COLOR_RESET}\n" "$1" >&2
}

warn() {
  printf "${COLOR_YELLOW}  ! %s${COLOR_RESET}\n" "$1" >&2
}

step() {
  printf "${COLOR_CYAN}  → %s${COLOR_RESET}\n" "$1" >&2
}

info() {
  printf "${COLOR_DIM}    %s${COLOR_RESET}\n" "$1" >&2
}

msg() {
  printf "  %s\n" "$1" >&2
}

# Compare semver: returns 0 if $1 >= $2
version_gte() {
  local highest
  highest=$(printf '%s\n%s' "$1" "$2" | sort -V | tail -1)
  [[ "$highest" == "$1" ]]
}

has_tools() {
  local missing=()
  for tool in "$@"; do
    command -v "$tool" &>/dev/null || missing+=("$tool")
  done
  [[ ${#missing[@]} -eq 0 ]] || err "Required tools not installed: ${missing[*]}"
}

normalize_arch() {
  case "$1" in
    x86_64)  echo "amd64" ;;
    aarch64) echo "arm64" ;;
    *)       echo "$1" ;;
  esac
}

get_os() {
  case $(uname -s) in
    Darwin) echo "darwin" ;;
    Linux)  echo "linux" ;;
    *)      echo "windows" ;;
  esac
}

get_binary_name() {
  local ver="${1#v}"  # strip leading 'v'
  echo "${BIN_NAME}_${ver}_${2}_${3}"  # version, os, arch
}

get_archive_name() {
  local ver="${1#v}"  # strip leading 'v'
  echo "${BIN_NAME}_${ver}_${2}_${3}.tar.gz"  # version, os, arch
}

# ==============================================================================
# Shell Completion Setup
# ==============================================================================

write_completion() {
  local bin_path="$1" shell="$2" target="$3"
  local target_dir tmp_file
  target_dir=$(dirname "$target")
  mkdir -p "$target_dir" 2>/dev/null || sudo mkdir -p "$target_dir" 2>/dev/null || return 1
  tmp_file=$(mktemp) || return 1
  if ! "$bin_path" completion "$shell" > "$tmp_file" 2>/dev/null || [[ ! -s "$tmp_file" ]]; then
    rm -f "$tmp_file"
    return 1
  fi
  chmod 644 "$tmp_file"
  if [[ -w "$target_dir" ]]; then
    mv "$tmp_file" "$target" || { rm -f "$tmp_file"; return 1; }
  else
    sudo mv "$tmp_file" "$target" || { rm -f "$tmp_file"; return 1; }
  fi
}

setup_bash_completions() {
  local bin_path="$1" os="$2" xdg_data="$3" user_only="$4"
  local sys_dir
  if [[ "$os" == "darwin" ]]; then
    local brew_prefix="/usr/local"
    [[ -d /opt/homebrew ]] && brew_prefix="/opt/homebrew"
    sys_dir="${brew_prefix}/etc/bash_completion.d"
  else
    sys_dir="/usr/share/bash-completion/completions"
  fi

  if [[ "$user_only" != "true" ]] && write_completion "$bin_path" bash "${sys_dir}/aicr"; then
    ok "Bash completions installed"
  elif write_completion "$bin_path" bash "${xdg_data}/bash-completion/completions/aicr"; then
    ok "Bash completions installed (user-local)"
  else
    warn "Could not install bash completions"
  fi
}

setup_zsh_completions() {
  local bin_path="$1" os="$2" xdg_data="$3" user_only="$4"
  local sys_dir
  if [[ "$os" == "darwin" ]]; then
    local brew_prefix="/usr/local"
    [[ -d /opt/homebrew ]] && brew_prefix="/opt/homebrew"
    sys_dir="${brew_prefix}/share/zsh/site-functions"
  else
    sys_dir="/usr/local/share/zsh/site-functions"
  fi

  if [[ "$user_only" != "true" ]] && write_completion "$bin_path" zsh "${sys_dir}/_aicr"; then
    ok "Zsh completions installed"
  elif write_completion "$bin_path" zsh "${xdg_data}/zsh/site-functions/_aicr"; then
    ok "Zsh completions installed (user-local)"
  else
    warn "Could not install zsh completions"
  fi
}

setup_fish_completions() {
  local bin_path="$1" user_only="$2"
  local sys_dir="/usr/share/fish/vendor_completions.d"
  local user_dir="${XDG_CONFIG_HOME:-$HOME/.config}/fish/completions"

  if [[ "$user_only" != "true" ]] && [[ -d "$sys_dir" ]] && write_completion "$bin_path" fish "${sys_dir}/aicr.fish"; then
    ok "Fish completions installed"
  elif [[ -d "${user_dir%/*}" ]] && write_completion "$bin_path" fish "${user_dir}/aicr.fish"; then
    ok "Fish completions installed (user-local)"
  fi
}

setup_completions() {
  if [[ "${AICR_NO_COMPLETIONS:-0}" == "1" ]]; then
    info "Skipping shell completion setup (AICR_NO_COMPLETIONS=1)"
    return 0
  fi

  local bin_path="${INSTALL_DIR}/${BIN_NAME}"
  local os xdg_data user_only="false"
  os=$(get_os)
  xdg_data="${XDG_DATA_HOME:-$HOME/.local/share}"

  # If the binary install did not need sudo, skip system completion dirs
  # to avoid unexpected sudo prompts (e.g., ./install -d ~/bin).
  [[ "${USED_SUDO:-false}" == "false" ]] && user_only="true"

  step "Setting up shell completions..."

  setup_bash_completions "$bin_path" "$os" "$xdg_data" "$user_only"
  setup_zsh_completions "$bin_path" "$os" "$xdg_data" "$user_only"
  setup_fish_completions "$bin_path" "$user_only"
}

# ==============================================================================
# GitHub API Functions
# ==============================================================================

github_api_curl() {
  local url="$1"
  local curl_opts=(-s)

  [[ -n "${GITHUB_TOKEN:-}" ]] && curl_opts+=(-H "Authorization: token $GITHUB_TOKEN")

  curl "${curl_opts[@]}" "$url"
}

download_release_asset() {
  local url="$1"
  local output="$2"
  local desc="${3:-file}"
  local curl_opts=(-# -L -o "$output" -w "%{http_code}")
  local http_code

  [[ -n "${GITHUB_TOKEN:-}" ]] && curl_opts+=(-H "Authorization: token $GITHUB_TOKEN")

  step "Downloading $desc..."
  http_code=$(curl "${curl_opts[@]}" "$url" | tail -c 3)

  [[ "$http_code" == "200" ]] || err "Failed to download $desc (HTTP $http_code): $url"
}

fetch_latest_version() {
  local redirect_url version

  step "Fetching latest release from GitHub..."

  # Use the redirect from /releases/latest to extract the tag — no API call needed,
  # so this works without GITHUB_TOKEN and is not subject to API rate limits.
  redirect_url=$(curl -sI "https://github.com/${REPO}/releases/latest" | grep -i '^location:' | tr -d '\r' | awk '{print $2}')

  if [[ -z "$redirect_url" ]]; then
    # Fallback to GitHub API (requires GITHUB_TOKEN if rate-limited)
    local api_url="https://api.github.com/repos/${REPO}/releases/latest"
    local response
    response=$(github_api_curl "$api_url")

    [[ -z "$response" ]] && err "Failed to fetch release information from GitHub"

    if grep -q "API rate limit exceeded" <<< "$response"; then
      err "GitHub API rate limit exceeded. Set GITHUB_TOKEN to avoid rate limits, or try again later."
    fi

    if grep -q "Not Found" <<< "$response"; then
      err "Repository not found or no releases available."
    fi

    version=$(grep -m 1 '"tag_name"' <<< "$response" | sed -E 's/.*"tag_name"[[:space:]]*:[[:space:]]*"([^"]+)".*/\1/')
  else
    # Extract tag from redirect URL: https://github.com/NVIDIA/aicr/releases/tag/v1.2.3
    version="${redirect_url##*/}"
  fi

  [[ -z "$version" || "$version" == "null" ]] && err "Failed to determine latest version"

  echo "$version"
}

extract_version() {
  local json="$1"
  local version

  version=$(grep -m 1 '"tag_name"' <<< "$json" | sed -E 's/.*"tag_name"[[:space:]]*:[[:space:]]*"([^"]+)".*/\1/')

  [[ -z "$version" || "$version" == "null" ]] && err "Failed to extract version from GitHub API response"

  echo "$version"
}

# Extract asset download URL from release JSON by asset name.
extract_asset_url() {
  local json="$1"
  local asset_name="$2"
  local url

  # Find the line number of the matching asset name
  local name_line
  name_line=$(echo "$json" | grep -n "\"name\".*\"${asset_name}\"" | head -1 | cut -d: -f1)

  [[ -z "$name_line" ]] && err "Asset '$asset_name' not found in release"

  # Extract browser_download_url (appears after "name" in the JSON)
  url=$(echo "$json" | sed -n "${name_line},$((name_line+30))p" | grep -oE '"browser_download_url"[[:space:]]*:[[:space:]]*"[^"]+"' | head -1 | sed -E 's/.*"browser_download_url"[[:space:]]*:[[:space:]]*"([^"]+)".*/\1/')

  [[ -z "$url" ]] && err "Failed to extract download URL for '$asset_name'"

  echo "$url"
}


# ==============================================================================
# Main Installation Logic
# ==============================================================================

main() {
  # Parse arguments
  while getopts "d:h" opt; do
    case $opt in
      d) INSTALL_DIR="$OPTARG" ;;
      h) usage ;;
      *) usage ;;
    esac
  done

  # Detect platform
  local os arch version binary_name temp_dir
  os=$(get_os)
  arch=$(normalize_arch "$(uname -m)")
  
  # Validate platform
  [[ $os == "linux" || $os == "darwin" ]] || err "Unsupported OS: $os"
  [[ $arch == "amd64" || $arch == "arm64" ]] || err "Unsupported architecture: $arch"
  has_tools "${REQUIRED_TOOLS[@]}"
  
  # Fetch latest version (uses redirect, no API rate limit)
  local archive_name
  version=$(fetch_latest_version)
  archive_name=$(get_archive_name "$version" "$os" "$arch")

  msg ""
  msg "Platform: $os/$arch"
  msg "Version:  $version"
  msg ""

  # Download archive and checksums using predictable GitHub release URLs
  temp_dir=$(mktemp -d)
  trap "rm -rf $temp_dir" EXIT

  local base_url="https://github.com/${REPO}/releases/download/${version}"
  local archive_url="${base_url}/${archive_name}"
  local checksum_url="${base_url}/${CHECKSUMS_FILE}"

  download_release_asset "$archive_url" "${temp_dir}/${archive_name}" "$archive_name"
  download_release_asset "$checksum_url" "${temp_dir}/checksums.txt" "checksums"

  # Verify archive checksum
  step "Verifying checksum (SHA-256)..."
  info "Ensures archive integrity — detects tampering or corruption in transit"
  (cd "$temp_dir" && grep "$archive_name" checksums.txt | shasum -a 256 --check --strict) \
    || err "Checksum verification failed — archive may have been tampered with"
  ok "Checksum verified — archive integrity confirmed"

  # Extract archive
  step "Extracting archive..."
  tar -xzf "${temp_dir}/${archive_name}" -C "$temp_dir"

  # Optional: verify attestation if cosign is available
  COSIGN_MIN_VERSION="v3.0.4"
  if command -v cosign &>/dev/null && [[ -f "${temp_dir}/${BIN_NAME}-attestation.sigstore.json" ]]; then
    cosign_version=$(cosign version 2>/dev/null | grep -oE 'v[0-9]+\.[0-9]+\.[0-9]+' | head -1) || true

    if [[ -n "$cosign_version" ]] && ! version_gte "$cosign_version" "$COSIGN_MIN_VERSION"; then
      warn "Installed cosign version ($cosign_version) is older than the minimum required ($COSIGN_MIN_VERSION)"
      info "Update cosign to verify provenance: https://docs.sigstore.dev/cosign/system_config/installation/"
    else
      step "Verifying provenance attestation (Sigstore cosign)..."
      info "Confirms this binary was built by NVIDIA/aicr CI — not a third-party rebuild"
      cosign_stderr=$(mktemp)
      if cosign verify-blob-attestation \
        --bundle "${temp_dir}/${BIN_NAME}-attestation.sigstore.json" \
        --type https://slsa.dev/provenance/v1 \
        --certificate-oidc-issuer https://token.actions.githubusercontent.com \
        --certificate-identity-regexp 'https://github.com/NVIDIA/aicr/.github/workflows/on-tag\.yaml@refs/tags/.*' \
        "${temp_dir}/${BIN_NAME}" 2>"${cosign_stderr}"; then
        ok "Provenance verified — binary built by github.com/NVIDIA/aicr CI pipeline"
      else
        warn "Attestation verification failed — cannot confirm binary provenance"
        info "The binary may still be valid, but its build origin could not be cryptographically verified"
        if [[ -s "${cosign_stderr}" ]]; then
          while IFS= read -r line; do
            info "cosign: $line"
          done < "${cosign_stderr}"
        fi
      fi
      rm -f "${cosign_stderr}"
    fi
  elif [[ -f "${temp_dir}/${BIN_NAME}-attestation.sigstore.json" ]]; then
    warn "cosign not found — skipping provenance verification"
    info "Install cosign to verify binary provenance: https://docs.sigstore.dev/cosign/system_config/installation/"
  fi

  # Install binary and attestation
  chmod +x "${temp_dir}/${BIN_NAME}"
  step "Installing $BIN_NAME to $INSTALL_DIR"
  mkdir -p "$INSTALL_DIR"
  USED_SUDO=false
  if [[ -w "$INSTALL_DIR" ]]; then
    mv "${temp_dir}/${BIN_NAME}" "${INSTALL_DIR}/${BIN_NAME}"
    [[ -f "${temp_dir}/${BIN_NAME}-attestation.sigstore.json" ]] && \
      mv "${temp_dir}/${BIN_NAME}-attestation.sigstore.json" "${INSTALL_DIR}/${BIN_NAME}-attestation.sigstore.json"
  else
    USED_SUDO=true
    sudo mv "${temp_dir}/${BIN_NAME}" "${INSTALL_DIR}/${BIN_NAME}"
    [[ -f "${temp_dir}/${BIN_NAME}-attestation.sigstore.json" ]] && \
      sudo mv "${temp_dir}/${BIN_NAME}-attestation.sigstore.json" "${INSTALL_DIR}/${BIN_NAME}-attestation.sigstore.json"
  fi

  # Verify installation
  msg "$BIN_NAME installed successfully!"
  "${BIN_NAME}" --version
  [[ -f "${INSTALL_DIR}/${BIN_NAME}-attestation.sigstore.json" ]] && \
    msg "Attestation: ${INSTALL_DIR}/${BIN_NAME}-attestation.sigstore.json"

  # Install shell completions
  setup_completions || warn "Shell completion setup failed (non-fatal)"

  # Fetch latest Sigstore trusted root for offline verification.
  # The trusted root enables 'aicr verify' to check attestation signatures
  # without contacting Sigstore infrastructure. If this fails (e.g., no network),
  # verification still works using the embedded TUF root compiled into the binary.
  msg "Updating Sigstore trusted root..."
  if ! "${BIN_NAME}" trust update 2>/dev/null; then
    msg "Warning: could not fetch latest trusted root (network may be unavailable)"
    msg "  Verification will use the embedded root. To update later, run:"
    msg "    ${BIN_NAME} trust update"
  fi
}

# Run main function
main "$@"