#!/bin/bash

set -e
set -o pipefail
set -o nounset

function pVerbose () {
    if [ "${VERBOSE}" == "1" ]; then
        echo $@
    fi
}

function pErr () {
    local msg="$@"

    if [ "${GUI}" == "1" ]; then
        notify-send $(basename $0) "$msg" --icon=dialog-error
    fi

    echo -e "$msg"
}

function die () {
    if [ $# -ne 0 ]; then
        pErr $@
    fi
    exit 1
}

function pushTrap () {
    local newHandler=$1; shift;
    local signals=$@;

    for signal in $signals; do
        local currHandler=$(trap -p $signal)
        currHandler=${currHandler#*\'}
        currHandler=${currHandler%\' $signal}

        trap "$newHandler; $currHandler" $signal
    done
}

function waitForProcess () {
    local pid="$1"; shift;
    local pidFile="/proc/$pid"

    if [ ! -e ${pidFile} ]; then
        pErr "Could not find process: $pid. Did it already exit?"
        return
    fi

    pVerbose "Waiting for process: $pid"

    # Can't do event based wait for PID unless it is a child PID
    # We are most likely interested in an external PID
    while [ -e /proc/$pid ]; do
        sleep ${PID_POLL_PRECISION};
    done

    pVerbose "Finished waiting for process: $pid"
}

function runAsRoot() {
    if [ "${GUI}" == "0" ]; then
        sudo "$@"
        return
    fi

    # Prefer pkexec as it automatically detects CLI fallback
    if [ -e "$(which pkexec)" ]; then
        pkexec "$@"
        return
    fi

    if [ -e "$(which gksudo)" ]; then
        gksudo "$@"
        return
    fi

    die "No suitable GUI sudo helper found, consider installing pkexec or gksu packages"
}

function setSuid () {
    local path="$1"; shift;
    if [ ! -u "$path" ]; then
        pVerbose "Setting $path as suid"
        runAsRoot chmod 4755 "$path"
        runAsRoot chown root:root "$path"
    fi
}

# Avoid cli defaults EDITOR and VISUAL since we might be running
# in a background process
function getDefaultReportEditor () {
    if [ -e "$(which gvim)" ]; then
        echo gvim
        return
    fi

    if [ -e "$(which gedit)" ]; then
        echo gedit
        return
    fi

    if [ -e "$(which kate)" ]; then
        echo gedit
        return
    fi

    pErr "Warning: No valid default report editor found"
}

function getDefaultReportVisualizer () {
    if [ -e "$(which gpuvis)" ]; then
        echo gpuvis
        return
    fi

    pErr "Warning: No valid default report visualizer found"
}

# helper variables
UNINIT="UNINITMAGIC"
TMPDIR="$(mktemp -d)"
pushTrap "rm -rf $TMPDIR"

# parameter defaults
VERBOSE=0
EXIT_SIG=2 #SIGINT
SNAPSHOT_SIG=5 #SIGTRAP
PID="${UNINIT}"
PID_POLL_PRECISION="0.1"
TRACE_OUTPUT="amdgpu-trace.dat"
DAT_OUTPUT="trace.dat"
REPORT=0
VISUALIZE=1
GUI=0
TRACE_CMD="$(which trace-cmd)"
REPORT_EDITOR="$(getDefaultReportEditor)"
REPORT_GPUVIS="$(getDefaultReportVisualizer)"
TRACEFS_ROOT=""

function printUsage() {
    pErr "Usage: $(basename $0) [options]"
    pErr
    pErr "Options:"
    pErr "    -p, --wait-for-pid <pid>  Trace for the lifetime of <pid>"
    pErr "    -s, --snapshot-signal <n> Snapshot trace when signal <n> is received (default: ${SNAPSHOT_SIG}"
    pErr "    -e, --exit-signal <n>     Trace until signal <n> is received (default: ${EXIT_SIG}"
    pErr "    -r, --report              Open a report in a text editor (default: ${REPORT_EDITOR})"
    pErr "        --editor <path>       Select a custom editor for a report"
    pErr "    -i, --vis                 Open a trace visualizer (default: ${REPORT_GPUVIS})"
    pErr "        --visualizer <path>   Select a custom trace visualizer"
    pErr "    -o, --output <path>       Specify output file location"
    pErr "    -d, --trace-dat <path>    Specify output location for dat file"
    pErr "    -g, --gui                 Run in a GUI friendly environment"
    pErr "    -v, --verbose             Verbose Output"
    pErr "    -h, --help                Print this help menu and exit"
}

while [[ $# -gt 0 ]]; do
    case "$1" in
        -v|--verbose) VERBOSE=1; shift ;;
        -p|--wait-for-pid) PID=$2; shift 2 ;;
        -s|--snapshot-signal) SNAPSHOT_SIG=$2; shift 2 ;;
        -e|--exit-signal) EXIT_SIG=$2; shift 2 ;;
        -o|--output) TRACE_OUTPUT="$2"; shift 2 ;;
        -g|--gui) GUI=1; shift ;;
        -r|--report) REPORT=1; shift ;;
        -i|--vis) VISUALIZE=1; shift ;;
        -d|--trace-dat) DAT_OUTPUT=$2; shift 2;;
        --editor) REPORT_EDITOR=$2; shift 2 ;;
        --visualizer) REPORT_GPUVIS=$2; shift 2 ;;
        -h|--help) printUsage; exit ;;
        *) echo "Unknown option $1"; printUsage; die;;
    esac
done

startTrace_TraceStarted="0"
function startTrace () {

    local traceEvents=""

    echo "Starting trace, please wait..."

    # https://github.com/mikesart/gpuvis/wiki/TechDocs-Linux-Scheduler
    traceEvents+=" -e sched:sched_switch"
    traceEvents+=" -e sched:sched_process_fork"
    traceEvents+=" -e sched:sched_process_exec"
    traceEvents+=" -e sched:sched_process_exit"

    traceEvents+=" -e drm:drm_vblank_event"
    traceEvents+=" -e drm:drm_vblank_event_queued"
    traceEvents+=" -e drm:drm_vblank_event_delivered"

    # https://github.com/mikesart/gpuvis/wiki/TechDocs-AMDGpu
    traceEvents+=" -e amdgpu:amdgpu_vm_flush"
    traceEvents+=" -e amdgpu:amdgpu_cs_ioctl"
    traceEvents+=" -e amdgpu:amdgpu_sched_run_job"
    traceEvents+=" -e amdgpu:amdgpu_ttm_bo_move"
    traceEvents+=" -e *fence:*fence_signaled"

    # https://github.com/mikesart/gpuvis/wiki/TechDocs-Intel
    #
    # NOTE: the i915_gem_request_submit, i915_gem_request_in, i915_gem_request_out
    # tracepoints require the CONFIG_DRM_I915_LOW_LEVEL_TRACEPOINTS Kconfig option to
    # be enabled.
    traceEvents+=" -e i915:i915_flip_request"
    traceEvents+=" -e i915:i915_flip_complete"
    traceEvents+=" -e i915:intel_gpu_freq_change"
    traceEvents+=" -e i915:i915_gem_request_add"
    traceEvents+=" -e i915:i915_gem_request_submit"
    traceEvents+=" -e i915:i915_gem_request_in"
    traceEvents+=" -e i915:i915_gem_request_out"
    traceEvents+=" -e i915:intel_engine_notify"
    traceEvents+=" -e i915:i915_gem_request_wait_begin"
    traceEvents+=" -e i915:i915_gem_request_wait_end"

    #${TRACE_CMD} reset
    ${TRACE_CMD} start -b 8000 -D -i ${traceEvents}

    startTrace_TraceStarted="1"
    echo "Capture started"
}

function stopTrace () {
    local tracingOn="$(cat ${TRACEFS_ROOT}/tracing_on)"

    if [ "${tracingOn}" != "1" ]; then
        return
    fi

    pVerbose "Stopping trace"

    ${TRACE_CMD} reset
    ${TRACE_CMD} snapshot -f
}

function captureTrace () {
    pVerbose "Capturing trace"

    ${TRACE_CMD} stop

    if [ "${startTrace_TraceStarted}" != "1" ]; then
        die "FATAL: trace interrupted before capture started"
    fi

    ${TRACE_CMD} extract -k -o ${TRACE_OUTPUT}
    ${TRACE_CMD} restart
}

function reportTrace () {
    pVerbose "Reporting trace"
    ${GPUVIS_CMD} ${TRACE_OUTPUT} &
}

function processReport () {
    stopTrace

    if [ "${VISUALIZE}" == "1" ]; then
        if [ -x "${REPORT_GPUVIS} " ]; then
            die "Invalid report editor: '${REPORT_EDITOR}'"
        fi
        ${REPORT_GPUVIS} ${TRACE_OUTPUT} &
    fi

    if [ "${REPORT}" == "1" ]; then
        if [ -x "${REPORT_EDITOR} " ]; then
            die "Invalid report editor: '${REPORT_EDITOR}'"
        fi
        ${REPORT_EDITOR} ${DAT_OUTPUT} &
    fi
}

function setupTraceCmd () {
    local path="${TRACE_CMD}"

    if [ "${path}" == "" ]; then
        die "Could not find required application: trace-cmd"
    fi

    setSuid "$path"

    # Invoking trace-cmd will mount tracefs if not mounted
    ${TRACE_CMD} stat > /dev/null
    TRACEFS_ROOT=$(mount | grep tracefs | head -1 | awk '{ print $3 }')
}

setupTraceCmd
pushTrap "captureTrace; processReport; exit" $EXIT_SIG
pushTrap "captureTrace; processReport" $SNAPSHOT_SIG
startTrace

if [ "${PID}" != "$UNINIT" ]; then
    waitForProcess $PID
    captureTrace
    processReport
else
    sleep inf
fi
