#!/bin/bash
# SPDX-License-Identifier: BSD-3-Clause
#
# ghidra_batch_decompile.sh - decompile the whole code range of an analyzed
# CUDA program in parallel. Chunks already present are skipped, so the command
# is safe to re-run and to interrupt.
#
# Ghidra takes an exclusive project lock even in -readOnly mode, so a single
# project cannot be opened by several processes. To parallelize, this script
# makes one copy of the analyzed project per worker under <outdir>/wk/ and
# round-robins the chunks across the copies. Copies are kept (they make
# resuming cheap) and can be removed with: rm -rf <outdir>/wk
#
# Usage:
#   ghidra_batch_decompile.sh <project_dir> <project_name> <program> <outdir> \
#                             <start_hex> <end_hex> [chunk_size_hex] [jobs] [timeout_s]
#
# Defaults: chunk 0x100000, jobs 5, timeout 120.
#   env: GHIDRA_HOME, GHIDRA_HEADLESS_MAXMEM (default 4G for batch jobs), MAX_CPU
#
# After it completes, merge per-chunk indexes with:
#   ghidra_index_merge.sh <outdir>
set -euo pipefail

HERE="$(cd "$(dirname "$0")" && pwd)"
GHIDRA_HOME="${GHIDRA_HOME:-/snap/ghidra/current/ghidra}"
ANALYZE="$GHIDRA_HOME/support/analyzeHeadless"
export GHIDRA_HEADLESS_MAXMEM="${GHIDRA_HEADLESS_MAXMEM:-4G}"

PROJECT_DIR="${1:?usage: $0 <project_dir> <project_name> <program> <outdir> <start_hex> <end_hex> [chunk_hex] [jobs] [timeout_s]}"
PROJECT_NAME="${2:?}"
PROGRAM="${3:?}"
OUTDIR="${4:?}"
START="${5:?}"
END="${6:?}"
CHUNK="${7:-0x100000}"
JOBS="${8:-5}"
TIMEOUT="${9:-120}"

mkdir -p "$OUTDIR/buckets" "$OUTDIR/index" "$OUTDIR/logs" "$OUTDIR/wk"

START_D=$((START)); END_D=$((END)); CHUNK_D=$((CHUNK))

# Build the chunk list first: every chunk gets its OWN project copy, so two
# concurrent jobs can never race on the same Ghidra project lock.
ranges=""
n=0
for ((a=START_D; a<END_D; a+=CHUNK_D)); do
    b=$((a+CHUNK_D)); [ "$b" -gt "$END_D" ] && b=$END_D
    ranges+=$(printf "0x%x 0x%x %d\n" "$a" "$b" "$n")$'\n'
    n=$((n+1))
done

# Create project copies for the chunks that still need work.
while read -r _a _b w; do
    [ -z "$w" ] && continue
    out="$OUTDIR/buckets/bucket_${_a#0x}_${_b#0x}.c"
    idx="$OUTDIR/index/bucket_${_a#0x}_${_b#0x}.tsv"
    if [ -f "$out" ] && [ -f "$idx" ]; then continue; fi
    wdir="$OUTDIR/wk"
    if [ ! -d "$wdir/$PROJECT_NAME.w$w.rep" ]; then
        echo "[*] creating project copy $w"
        cp -a "$PROJECT_DIR/$PROJECT_NAME.gpr" "$wdir/$PROJECT_NAME.w$w.gpr"
        cp -a "$PROJECT_DIR/$PROJECT_NAME.rep" "$wdir/$PROJECT_NAME.w$w.rep"
    fi
done <<< "$ranges"

echo "[*] decompiling $PROGRAM in $n chunks across $JOBS workers"
echo "$ranges" | xargs -P "$JOBS" -n 3 bash -c '
    set -e
    a="$1"; b="$2"; w="$3"
    tag="${a#0x}_${b#0x}"
    log="'$OUTDIR'/logs/decompile_${tag}.log"
    GHIDRA_HEADLESS_MAXMEM="'$GHIDRA_HEADLESS_MAXMEM'" "'$ANALYZE'" "'$OUTDIR'/wk" "'$PROJECT_NAME'.w$w" \
        -process "'$PROGRAM'" -readOnly -noanalysis \
        -scriptPath "'$HERE'/ghidra" \
        -postScript BatchDecompile.java "$a" "$b" "'$OUTDIR'" "'$TIMEOUT'" \
        > "$log" 2>&1 || { echo "[FAIL] $a-$b (see $log)"; exit 1; }
    tail -1 "$log"
' _

echo "[ok] batch decompilation done -> $OUTDIR/buckets"
echo "     merge indexes: $HERE/ghidra_index_merge.sh $OUTDIR"