#!/usr/bin/env bash

set -euo pipefail

usage() {
    echo "Usage: $0 <tag> <dpu-pci-bdf> <chunk-directory> <chunk-size> [run-label]"
    echo
    echo "Examples:"
    echo "  $0 VANILLA_1G 03:00.0 bins/1G 1M VANILLA_1G_run_1"
    echo "  $0 TRACE_5G   03:00.0 bins/5G 1M TRACE_5G_run_1"
}

if [[ $# -lt 4 || $# -gt 5 ]]; then
    usage
    exit 1
fi

TAG="$1"
DPU_PCI="$2"
CHUNK_DIR="$3"
CHUNK_SIZE="$4"
RUN_LABEL="${5:-${TAG}}"

DECOMPRESS_BIN="./decompress"

LOG_DIR="logs"
OUT_CSV="${LOG_DIR}/cpu_${TAG}.csv"

mkdir -p "$LOG_DIR"

TMP_DIR="$(mktemp -d)"
DECOMPRESS_LOG="${TMP_DIR}/decompress.log"
TIME_LOG="${TMP_DIR}/time.log"

cleanup() {
    rm -rf "$TMP_DIR"
}
trap cleanup EXIT INT TERM

die() {
    echo "ERROR: $*" >&2
    exit 1
}

read_proc_stat() {
    # Outputs: idle total softirq
    awk '
    /^cpu / {
        user_v=$2
        nice_v=$3
        system_v=$4
        idle_v=$5
        iowait_v=$6
        irq_v=$7
        softirq_v=$8
        steal_v=$9

        idle_all = idle_v + iowait_v
        non_idle = user_v + nice_v + system_v + irq_v + softirq_v + steal_v
        total = idle_all + non_idle

        printf "%.0f %.0f %.0f\n", idle_all, total, softirq_v
        exit
    }
    ' /proc/stat
}

calc_system_cpu_pct() {
    local before_idle="$1"
    local before_total="$2"
    local after_idle="$3"
    local after_total="$4"

    awk \
        -v bi="$before_idle" \
        -v bt="$before_total" \
        -v ai="$after_idle" \
        -v at="$after_total" \
        'BEGIN {
            total_delta = at - bt
            idle_delta = ai - bi

            if (total_delta <= 0) {
                printf ""
            } else {
                cpu = 100.0 * (total_delta - idle_delta) / total_delta
                printf "%.6f", cpu
            }
        }'
}

calc_softirq_cpu_pct() {
    local before_total="$1"
    local before_softirq="$2"
    local after_total="$3"
    local after_softirq="$4"

    awk \
        -v bt="$before_total" \
        -v bs="$before_softirq" \
        -v at="$after_total" \
        -v as="$after_softirq" \
        'BEGIN {
            total_delta = at - bt
            softirq_delta = as - bs

            if (total_delta <= 0) {
                printf ""
            } else {
                pct = 100.0 * softirq_delta / total_delta
                printf "%.6f", pct
            }
        }'
}

parse_time_cpu_pct() {
    local file="$1"

    awk -F: '
    /Percent of CPU this job got/ {
        val = $2
        gsub(/^[ \t]+|[ \t%]+$/, "", val)
        print val
        found = 1
    }
    END {
        if (!found)
            print ""
    }
    ' "$file" 2>/dev/null | tail -n1
}

write_header_if_needed() {
    if [[ ! -f "$OUT_CSV" ]]; then
        echo "system_cpu_pct,process_cpu_pct,softirq_cpu_pct" > "$OUT_CSV"
    fi
}

[[ -x "$DECOMPRESS_BIN" ]] || die "$DECOMPRESS_BIN not found or not executable"
[[ -d "$CHUNK_DIR" ]] || die "chunk directory not found: $CHUNK_DIR"

echo "[dpu] tag:        $TAG"
echo "[dpu] run label:  $RUN_LABEL"
echo "[dpu] chunk dir:  $CHUNK_DIR"
echo "[dpu] chunk size: $CHUNK_SIZE"
echo "[dpu] pci:        $DPU_PCI"
echo "[dpu] csv:        $OUT_CSV"

echo "[dpu] dropping filesystem caches..."
sudo -n sh -c 'sync; echo 3 > /proc/sys/vm/drop_caches'

read before_idle before_total before_softirq < <(read_proc_stat)

set +e
(
    /usr/bin/time -v sudo -n "$DECOMPRESS_BIN" \
        -p "$DPU_PCI" \
        -d "$CHUNK_DIR" \
        -S "$CHUNK_SIZE"
) > "$DECOMPRESS_LOG" 2> "$TIME_LOG"
DECOMPRESS_STATUS=$?
set -e

read after_idle after_total after_softirq < <(read_proc_stat)

system_cpu_pct="$(
    calc_system_cpu_pct "$before_idle" "$before_total" "$after_idle" "$after_total"
)"

softirq_cpu_pct="$(
    calc_softirq_cpu_pct "$before_total" "$before_softirq" "$after_total" "$after_softirq"
)"

process_cpu_pct="$(parse_time_cpu_pct "$TIME_LOG")"

write_header_if_needed

printf "%s,%s,%s\n" \
    "${system_cpu_pct:-}" \
    "${process_cpu_pct:-}" \
    "${softirq_cpu_pct:-}" \
    >> "$OUT_CSV"

# Keep these prints so the host script can parse chunks/compressed_bytes/bytes/time/BW.
cat "$DECOMPRESS_LOG"
cat "$TIME_LOG"

if [[ "$DECOMPRESS_STATUS" -ne 0 ]]; then
    echo "[dpu] decompressor failed with status $DECOMPRESS_STATUS"
    echo "[dpu] appended CSV: $OUT_CSV"
    exit "$DECOMPRESS_STATUS"
fi

echo "[dpu] decompressor completed successfully"
echo "[dpu] appended CSV: $OUT_CSV"