#!/bin/sh
set -eu

APP_NAME="brex"
BIN_NAME="brex"
BASE_URL="${BREX_RELEASE_BASE_URL:-https://brex.chaitanya.science}"
DEFAULT_MINISIGN_PUBKEY="RWSNb4j3C69NFxNVviY4Myshhsywk/gt4oH8LfHksIA2ep3O47sxE0FK"
MINISIGN_PUBKEY="${BREX_MINISIGN_PUBKEY:-${DEFAULT_MINISIGN_PUBKEY}}"

is_tty() {
  [ -t 1 ]
}

brand_paint() {
  if is_tty; then
    printf '\033[38;2;118;185;0m%s\033[0m\n' "$*"
  else
    printf '%s\n' "$*"
  fi
}

paint() {
  code="$1"
  shift
  if is_tty; then
    printf '\033[%sm%s\033[0m\n' "$code" "$*"
  else
    printf '%s\n' "$*"
  fi
}

info() {
  paint "1;38;2;118;185;0" "NVIDIA > $*"
}

note() {
  paint "0;37" "  · $*"
}

success() {
  paint "1;38;2;118;185;0" "  ✓ $*"
}

warn() {
  if is_tty; then
    printf '\033[1;33m  ! %s\033[0m\n' "$*" >&2
  else
    printf '! %s\n' "$*" >&2
  fi
}

die() {
  if is_tty; then
    printf '\033[1;31m  x %s\033[0m\n' "$*" >&2
  else
    printf 'x %s\n' "$*" >&2
  fi
  exit 1
}

have() {
  command -v "$1" >/dev/null 2>&1
}

json_escape() {
  printf '%s' "$1" | sed \
    -e 's/\\/\\\\/g' \
    -e 's/"/\\"/g' \
    -e 's/\r/\\r/g' \
    -e 's/\n/\\n/g' \
    -e 's/\t/\\t/g'
}

json_string() {
  printf '"%s"' "$(json_escape "$1")"
}

json_nullable_string() {
  if [ -n "$1" ]; then
    json_string "$1"
  else
    printf 'null'
  fi
}

json_number() {
  if [ -n "$1" ]; then
    printf '%s' "$1"
  else
    printf '0'
  fi
}

state_dir() {
  if [ -n "${BREX_INSTALL_STATE_DIR:-}" ]; then
    printf '%s' "$BREX_INSTALL_STATE_DIR"
    return
  fi
  printf '%s/.brex' "${HOME:-.}"
}

install_id_path() {
  printf '%s/install_id' "$(state_dir)"
}

generate_install_id() {
  if have uuidgen; then
    uuidgen | tr '[:upper:]' '[:lower:]'
    return
  fi
  if have openssl; then
    openssl rand -hex 16
    return
  fi
  if [ -r /proc/sys/kernel/random/uuid ]; then
    cat /proc/sys/kernel/random/uuid | tr '[:upper:]' '[:lower:]'
    return
  fi
  printf '%s-%s-%s' "$(date +%s 2>/dev/null || printf '0')" "$$" "${RANDOM:-0}"
}

load_install_id() {
  if [ -n "${BREX_INSTALL_ID:-}" ]; then
    INSTALL_ID="$BREX_INSTALL_ID"
    return
  fi

  path="$(install_id_path)"
  if [ -f "$path" ]; then
    INSTALL_ID="$(tr -d '\r\n' <"$path" || true)"
    if [ -n "$INSTALL_ID" ]; then
      return
    fi
  fi

  mkdir -p "$(dirname "$path")" 2>/dev/null || {
    INSTALL_ID="$(generate_install_id)"
    return
  }
  INSTALL_ID="$(generate_install_id)"
  printf '%s\n' "$INSTALL_ID" >"$path" 2>/dev/null || true
}

verify_minisign_file() {
  file="$1"
  sig="$2"
  if ! have minisign; then
    die "minisign is required to verify signed releases"
  fi
  minisign -Vm "$file" -x "$sig" -P "$MINISIGN_PUBKEY" >/dev/null 2>&1 \
    || die "Minisign signature verification failed"
}

fetch_text() {
  url="$1"
  if have curl; then
    curl -fsSL "$url"
    return
  fi
  if have wget; then
    wget -qO- "$url"
    return
  fi
  die "curl or wget is required to install ${APP_NAME}"
}

download_file() {
  url="$1"
  destination="$2"
  if have curl; then
    curl -fsSL "$url" -o "$destination"
    return
  fi
  if have wget; then
    wget -qO "$destination" "$url"
    return
  fi
  die "curl or wget is required to install ${APP_NAME}"
}

sha256_file() {
  file="$1"
  if have sha256sum; then
    sha256sum "$file" | awk '{print $1}'
    return
  fi
  if have shasum; then
    shasum -a 256 "$file" | awk '{print $1}'
    return
  fi
  if have openssl; then
    openssl dgst -sha256 "$file" | awk '{print $NF}'
    return
  fi
  die "No SHA-256 tool found (need sha256sum, shasum, or openssl)"
}

send_install_beacon() {
  event_type="$1"
  error_summary="${2:-}"
  beacon_url="${BASE_URL%/}/analytics/beacon"
  timestamp="$(date -u +%Y-%m-%dT%H:%M:%SZ 2>/dev/null || date)"
  payload="$(printf '{"events":[{"event_type":%s,"timestamp":%s,"install_id":%s,"app_version":%s,"os":%s,"arch":%s,"route_source":%s,"auth_method":null,"provider":null,"model":null,"session_id":null,"input_tokens":0,"output_tokens":0,"cache_read_tokens":0,"cache_creation_tokens":0,"cost_usd":0,"duration_ms":0,"error_summary":%s,"target_version":%s}]}' \
    "$(json_string "$event_type")" \
    "$(json_string "$timestamp")" \
    "$(json_string "$INSTALL_ID")" \
    "$(json_string "$VERSION")" \
    "$(json_string "$PLATFORM_OS")" \
    "$(json_string "$PLATFORM_ARCH")" \
    "$(json_string "installer")" \
    "$(json_nullable_string "$error_summary")" \
    "$(json_string "$VERSION")")"

  if have curl; then
    curl -fsSL -X POST -H "content-type: application/json" --data "$payload" "$beacon_url" >/dev/null 2>&1 || true
    return
  fi
  if have wget; then
    wget -qO- --method=POST --header="content-type: application/json" --body-data="$payload" "$beacon_url" >/dev/null 2>&1 || true
    return
  fi
}

detect_platform() {
  os="$(uname -s 2>/dev/null | tr '[:upper:]' '[:lower:]')"
  arch="$(uname -m 2>/dev/null)"

  case "$os" in
    darwin) PLATFORM_OS="darwin" ;;
    *)
      die "brex currently ships installer builds only for Apple Silicon Macs"
      ;;
  esac

  case "$arch" in
    arm64|aarch64) PLATFORM_ARCH="aarch64" ;;
    x86_64|amd64)
      die "Intel Macs are not supported by the brex installer. Use an Apple Silicon Mac."
      ;;
    *)
      die "Unsupported CPU architecture: ${arch}. brex currently supports Apple Silicon only."
      ;;
  esac
}

resolve_install_dir() {
  if [ -n "${BREX_INSTALL_DIR:-}" ]; then
    INSTALL_DIR="$BREX_INSTALL_DIR"
    return
  fi

  if have "$BIN_NAME"; then
    existing_bin="$(command -v "$BIN_NAME" || true)"
    existing_dir="$(dirname "$existing_bin")"
    if [ -n "$existing_dir" ] && [ -w "$existing_dir" ]; then
      INSTALL_DIR="$existing_dir"
      return
    fi
  fi

  if printf '%s' ":${PATH}:" | grep -q ":${HOME}/.local/bin:"; then
    INSTALL_DIR="${HOME}/.local/bin"
    return
  fi

  if [ -d "/usr/local/bin" ] && [ -w "/usr/local/bin" ]; then
    INSTALL_DIR="/usr/local/bin"
    return
  fi

  INSTALL_DIR="${HOME}/.local/bin"
}

print_path_hint() {
  case ":${PATH}:" in
    *":${INSTALL_DIR}:"*) return ;;
  esac

  warn "${INSTALL_DIR} is not on your PATH yet, so ${BIN_NAME} will not resolve by name yet."
  brand_paint '  Add Brex CLI to PATH:'
  shell_name="$(basename "${SHELL:-sh}")"
  case "$shell_name" in
    zsh)
      printf '   echo '\''export PATH="%s:$PATH"'\'' >> ~/.zshrc\n' "$INSTALL_DIR"
      ;;
    bash)
      printf '   echo '\''export PATH="%s:$PATH"'\'' >> ~/.bashrc\n' "$INSTALL_DIR"
      ;;
    fish)
      printf "   fish_add_path %s\n" "$INSTALL_DIR"
      ;;
    *)
      printf '   export PATH="%s:$PATH"\n' "$INSTALL_DIR"
      ;;
  esac
}

print_banner() {
  printf '\n'
  brand_paint 'NVIDIA BREX'
  brand_paint '========================================'
  paint "1;37" 'Brex CLI installer'
  paint "0;90" 'Green path. Verified install. Apple Silicon macOS.'
}

print_banner

have uname || die "uname is required"
have tar || die "tar is required"
have mktemp || die "mktemp is required"
[ "$MINISIGN_PUBKEY" != "__BREX_""MINISIGN_PUBKEY__" ] || die "Installer is missing embedded minisign public key"

detect_platform
resolve_install_dir

VERSION="${BREX_VERSION:-}"
if [ -z "$VERSION" ]; then
  info "Resolving latest release"
  VERSION="$(fetch_text "${BASE_URL}/latest" | tr -d '\r\n')"
fi

[ -n "$VERSION" ] || die "Could not resolve a brex release version"
load_install_id

ARCHIVE_URL="${BASE_URL}/download/${PLATFORM_OS}/${PLATFORM_ARCH}?version=${VERSION}"
CHECKSUM_URL="${BASE_URL}/checksum/${PLATFORM_OS}/${PLATFORM_ARCH}?version=${VERSION}"
SIGNATURE_URL="${BASE_URL}/signature/${PLATFORM_OS}/${PLATFORM_ARCH}?version=${VERSION}"

info "Preparing NVIDIA Brex ${VERSION}"
note "platform: ${PLATFORM_OS}/${PLATFORM_ARCH}"
note "target:   ${INSTALL_DIR}/${BIN_NAME}"
send_install_beacon "install_started"

TMPDIR_ROOT="${TMPDIR:-/tmp}"
WORKDIR="$(mktemp -d "${TMPDIR_ROOT%/}/brex-install.XXXXXX")"
ARCHIVE_PATH="${WORKDIR}/brex.tar.gz"
SIGNATURE_PATH="${WORKDIR}/brex.tar.gz.minisig"
trap 'rm -rf "$WORKDIR"' EXIT INT TERM HUP

EXPECTED_SHA="$(fetch_text "$CHECKSUM_URL" | tr -d '\r\n')"
[ -n "$EXPECTED_SHA" ] || die "Checksum endpoint returned no data"

info "Downloading signed release payload"
download_file "$ARCHIVE_URL" "$ARCHIVE_PATH"
download_file "$SIGNATURE_URL" "$SIGNATURE_PATH"

ACTUAL_SHA="$(sha256_file "$ARCHIVE_PATH")"
[ "$ACTUAL_SHA" = "$EXPECTED_SHA" ] || die "Checksum mismatch for downloaded archive"
success "Checksum verified"

verify_minisign_file "$ARCHIVE_PATH" "$SIGNATURE_PATH"
success "Signature verified"

info "Extracting Brex CLI"
tar -xzf "$ARCHIVE_PATH" -C "$WORKDIR"

EXTRACTED_BIN="${WORKDIR}/${BIN_NAME}"
if [ ! -x "$EXTRACTED_BIN" ]; then
  EXTRACTED_BIN="$(find "$WORKDIR" -type f -name "$BIN_NAME" | head -n 1 || true)"
fi
[ -n "$EXTRACTED_BIN" ] && [ -f "$EXTRACTED_BIN" ] || die "Archive did not contain ${BIN_NAME}"

mkdir -p "$INSTALL_DIR"
TARGET_BIN="${INSTALL_DIR}/${BIN_NAME}"

info "Placing binary"
if have install; then
  install -m 0755 "$EXTRACTED_BIN" "$TARGET_BIN"
else
  cp "$EXTRACTED_BIN" "$TARGET_BIN"
  chmod 0755 "$TARGET_BIN"
fi

INSTALLED_VERSION="$("$TARGET_BIN" --version 2>/dev/null || true)"
send_install_beacon "install_completed"
success "NVIDIA Brex is installed"
printf '\n'
brand_paint '  Ready:'
printf '   %s\n' "${INSTALLED_VERSION:-${TARGET_BIN}}"
printf '   %s\n' "Command: ${BIN_NAME}"
case ":${PATH}:" in
  *":${INSTALL_DIR}:"*)
    printf '   %s\n' "Try: ${BIN_NAME} --help"
    ;;
  *)
    printf '   %s\n' "Try once: ${TARGET_BIN} --help"
    ;;
esac
printf '\n'
print_path_hint
