#!/usr/bin/env bash
# Profile-Guided Optimization build for mise.
#
# Three-phase rustc PGO flow, adapted from jdx/aube's benchmarks/pgo.bash:
#
#   1. Build mise with -Cprofile-generate (instrumented binary).
#   2. Train via scripts/train-startup.bash against a hermetic, offline
#      workload covering the startup
#      hot paths: settings/config load, toolset resolution, hook-env
#      (both the full run and the per-prompt early-exit fast path),
#      env rendering, task listing, and exec. No network, no registry —
#      everything runs in a throwaway HOME with isolated MISE_*_DIR s.
#   3. Merge .profraw via the rustup toolchain's llvm-profdata,
#      recompile with -Cprofile-use.
#
# The startup path dominates every mise invocation (shell prompts via
# hook-env, shims, version checks), so the training workload is biased
# toward short-lived invocations rather than long installs.
#
# Env hooks (mirroring aube's CI contract):
#   MISE_PGO_PROFILE=<profile>  cargo profile for both phases
#                               (default: serious, matching release
#                               tarballs).
#   MISE_PGO_TARGET=<triple>    cross-compilation target (default:
#                               host). The instrumented binary must be
#                               runnable on THIS machine for training —
#                               same-arch targets only.
#   MISE_PGO_BUILD_TOOL=<tool>  `cargo` (default) or `cross`. cross is
#                               used in CI for Linux targets so the
#                               binary keeps cross's older glibc
#                               baseline; the cross-built instrumented
#                               binary still runs on the host for
#                               training (older glibc is
#                               forward-compatible).
#   MISE_PGO_SKIP_FINAL_BUILD=1 stop after merging .profraw, for CI
#                               setups that delegate the final build to
#                               a separate step picking up RUSTFLAGS.
#
# Additional arguments are passed through to both cargo invocations —
# scripts/build-tarball.sh uses this for --no-default-features
# --features "...".
set -euo pipefail

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

PGO_DATA_DIR="$REPO_ROOT/target/pgo-data"
PGO_PROFRAW_DIR="$PGO_DATA_DIR/profraw"
PGO_MERGED="$PGO_DATA_DIR/merged.profdata"

PGO_PROFILE="${MISE_PGO_PROFILE:-serious}"
PGO_TARGET="${MISE_PGO_TARGET:-}"
PGO_BUILD_TOOL="${MISE_PGO_BUILD_TOOL:-cargo}"

# target_arg stays unquoted at expansion sites: empty string disappears,
# "--target=foo" expands to one arg. Avoids bash 3.2 (macOS) array+set -u
# unbound-variable issues with "${arr[@]}".
target_arg=""
target_dir_part=""
if [ -n "$PGO_TARGET" ]; then
	target_arg="--target=$PGO_TARGET"
	target_dir_part="$PGO_TARGET/"
fi

RUSTC_HOST="$(rustc -vV | sed -n 's|^host: ||p')"
RUSTC_SYSROOT="$(rustc --print sysroot)"
LLVM_PROFDATA="$RUSTC_SYSROOT/lib/rustlib/$RUSTC_HOST/bin/llvm-profdata"
if [ ! -x "$LLVM_PROFDATA" ]; then
	echo "ERROR: llvm-profdata not found at $LLVM_PROFDATA" >&2
	echo "  Install with: rustup component add llvm-tools" >&2
	exit 1
fi

mkdir -p "$PGO_PROFRAW_DIR"
rm -f "$PGO_PROFRAW_DIR"/*.profraw "$PGO_MERGED"

# With MISE_PGO_BUILD_TOOL=cross, rustc runs inside a container that
# mounts the project at a container path, so when phase 3 reads
# -Cprofile-use=<host-path> from RUSTFLAGS the file is invisible.
# Bind-mount PGO_DATA_DIR at the same host path inside the container so
# the RUSTFLAGS value resolves. Harmless for phase 1.
if [ "$PGO_BUILD_TOOL" = "cross" ]; then
	export CROSS_CONTAINER_OPTS="${CROSS_CONTAINER_OPTS:-} -v $PGO_DATA_DIR:$PGO_DATA_DIR:rw"
fi

# Keep PGO Rust-only on macOS. The cc crate mirrors -Cprofile-generate /
# -Cprofile-use from RUSTFLAGS into the C compiler command line
# (-fprofile-generate / -fprofile-use), so vendored C code (openssl, lua)
# gets instrumented by Apple clang — whose profile-runtime format follows
# Xcode's LLVM, not rustc's. When those disagree (Xcode 16 = raw version 8,
# rustc 1.97 = 10), the instrumented binary aborts every profile write with
#   LLVM Profile Error: Runtime and instrumentation version mismatch
# and training produces zero .profraw files. C-side profiles couldn't be
# consumed by a mismatched clang in phase 3b anyway. cc appends environment
# flags after rustc-inherited ones, so these -fno- overrides win.
if [ "$(uname -s)" = "Darwin" ]; then
	export CFLAGS="${CFLAGS:-} -fno-profile-generate -fno-profile-use"
	export CXXFLAGS="${CXXFLAGS:-} -fno-profile-generate -fno-profile-use"
fi

# ---------- Phase 1: instrumented build ----------
echo ">>> [1/3] Building instrumented binary ($PGO_BUILD_TOOL, profile=$PGO_PROFILE${PGO_TARGET:+, target=$PGO_TARGET})"
# shellcheck disable=SC2086 # intentional word-splitting on $target_arg
RUSTFLAGS="${RUSTFLAGS:-} -Cprofile-generate=$PGO_PROFRAW_DIR" \
	"$PGO_BUILD_TOOL" build --profile="$PGO_PROFILE" $target_arg --bin mise "$@"

INSTRUMENTED_BIN="$REPO_ROOT/target/${target_dir_part}${PGO_PROFILE}/mise"
if [ ! -x "$INSTRUMENTED_BIN" ]; then
	echo "ERROR: instrumented binary missing at $INSTRUMENTED_BIN" >&2
	exit 1
fi

# ---------- Phase 2: training ----------
echo ">>> [2/3] Training against hermetic offline workload"
# Force the instrumented binary to write profraw to a path we control,
# regardless of what -Cprofile-generate=<dir> baked in at compile time
# (with cross, the compile-time path is a container path that doesn't
# resolve on the host). %m disambiguates per module; %p per process.
export LLVM_PROFILE_FILE="$PGO_PROFRAW_DIR/mise-%m-%p.profraw"
bash "$SCRIPT_DIR/train-startup.bash" "$INSTRUMENTED_BIN"
unset LLVM_PROFILE_FILE

# Sanity check: confirm training actually wrote profraw. On cross-built
# targets a bad LLVM_PROFILE_FILE path silently produces zero files and
# llvm-profdata merges nothing. Fail loudly here instead.
profraw_count=$(find "$PGO_PROFRAW_DIR" -maxdepth 1 -name '*.profraw' -type f | wc -l | tr -d ' ')
if [ "$profraw_count" -eq 0 ]; then
	echo "ERROR: no .profraw files written to $PGO_PROFRAW_DIR after training" >&2
	echo "  Training ran but the instrumented binary did not record profile data." >&2
	exit 1
fi
echo ">>> $profraw_count .profraw files collected"

# ---------- Phase 3a: merge ----------
echo ">>> [3/3] Merging profile data"
"$LLVM_PROFDATA" merge -o "$PGO_MERGED" "$PGO_PROFRAW_DIR"

# Defense in depth: a version mismatch between the rustc that
# instrumented and the llvm-profdata that merges can be a 0-exit silent
# no-op; phase 3b would then fail with an opaque missing-file error.
if [ ! -f "$PGO_MERGED" ]; then
	echo "ERROR: $PGO_MERGED was not produced by llvm-profdata merge" >&2
	exit 1
fi
echo ">>> merged profile written: $(stat -c %s "$PGO_MERGED" 2>/dev/null || stat -f %z "$PGO_MERGED") bytes"

if [ -n "${MISE_PGO_SKIP_FINAL_BUILD:-}" ]; then
	echo ">>> Skipping final optimized build (MISE_PGO_SKIP_FINAL_BUILD=1)"
	echo ">>> Profile ready at: $PGO_MERGED"
	exit 0
fi

# ---------- Phase 3b: optimize ----------
echo ">>> Rebuilding with -Cprofile-use"
# -Cllvm-args=-pgo-warn-missing-function=false: coverage gaps are
# expected (training can't exercise every path) and LLVM otherwise logs
# a note per uncovered symbol. Uncovered functions compile normally,
# just without PGO data.
bolt_rustflags=""
if [ -n "${MISE_BOLT:-}" ]; then
	# BOLT needs text relocations to safely reorder whole functions. Only the
	# final binary needs them; keeping them out of the instrumented link saves
	# memory and disk during phase 1.
	bolt_rustflags="-Clink-arg=-Wl,--emit-relocs"
fi
# shellcheck disable=SC2086 # intentional word-splitting on $target_arg
RUSTFLAGS="${RUSTFLAGS:-} -Cprofile-use=$PGO_MERGED -Cllvm-args=-pgo-warn-missing-function=false $bolt_rustflags" \
	"$PGO_BUILD_TOOL" build --profile="$PGO_PROFILE" $target_arg --bin mise "$@"

# Phase 3b wrote to the same path as phase 1, so the file at
# $INSTRUMENTED_BIN is now the PGO-optimized build, not the instrumented
# one. Alias for clarity.
PGO_FINAL_BIN="$INSTRUMENTED_BIN"
echo ">>> PGO build complete: $PGO_FINAL_BIN"
ls -lh "$PGO_FINAL_BIN"
