#!/bin/bash

# Thread-level CPU profiler for c-relay on remote server
# - Samples per-thread CPU every N seconds using /proc/<pid>/task/*/stat deltas
# - Optionally runs perf record in parallel for callgraph data
# - Pulls artifacts locally and generates a summary report

set -euo pipefail

REMOTE_HOST="${REMOTE_HOST:-ubuntu@laantungir.com}"
DURATION="${DURATION:-600}"          # seconds (default 10 minutes)
INTERVAL="${INTERVAL:-10}"           # seconds between samples
OUTDIR_BASE="${OUTDIR_BASE:-thread_profile_runs}"
RUN_TS="$(date -u +%Y%m%d_%H%M%S)"
RUN_DIR="${OUTDIR_BASE}/${RUN_TS}"
LOCAL_CSV="${RUN_DIR}/thread_samples.csv"
LOCAL_REPORT="${RUN_DIR}/thread_summary.txt"
REMOTE_DIR="/tmp/c_relay_thread_profile_${RUN_TS}"
ENABLE_PERF="${ENABLE_PERF:-1}"      # 1=run perf, 0=skip

mkdir -p "${RUN_DIR}"

echo "=========================================="
echo "C-Relay Thread CPU Profiler"
echo "=========================================="
echo "Remote host : ${REMOTE_HOST}"
echo "Duration    : ${DURATION}s"
echo "Interval    : ${INTERVAL}s"
echo "Run dir     : ${RUN_DIR}"
echo "Perf record : ${ENABLE_PERF}"
echo ""

ssh "${REMOTE_HOST}" "bash -s" <<EOF
set -euo pipefail

REMOTE_DIR="${REMOTE_DIR}"
DURATION="${DURATION}"
INTERVAL="${INTERVAL}"
ENABLE_PERF="${ENABLE_PERF}"

mkdir -p "\${REMOTE_DIR}"
CSV_FILE="\${REMOTE_DIR}/thread_samples.csv"
PERF_FILE="\${REMOTE_DIR}/perf.data"
LOG_FILE="\${REMOTE_DIR}/run.log"

PID=\$(pgrep -f '/usr/local/bin/c_relay/c_relay|c_relay' | head -1 || true)
if [ -z "\${PID}" ]; then
  echo "ERROR: c_relay process not found (checked by cmdline pattern)" | tee -a "\${LOG_FILE}"
  exit 1
fi

echo "Profiling PID: \${PID}" | tee -a "\${LOG_FILE}"
CLK_TCK=\$(getconf CLK_TCK)
echo "CLK_TCK=\${CLK_TCK}" | tee -a "\${LOG_FILE}"

echo "timestamp,tid,thread_name,cpu_pct,delta_ticks,total_ticks" > "\${CSV_FILE}"

# Start perf in background (if enabled and available)
PERF_PID=""
if [ "\${ENABLE_PERF}" = "1" ]; then
  if command -v perf >/dev/null 2>&1; then
    if sudo -n true >/dev/null 2>&1; then
      echo "Starting perf record in background..." | tee -a "\${LOG_FILE}"
      sudo perf record -g -p "\${PID}" -o "\${PERF_FILE}" -- sleep "\${DURATION}" >/dev/null 2>&1 &
      PERF_PID=\$!
      PERF_OWNER="\$(id -un)"
      echo "perf_pid=\${PERF_PID}" | tee -a "\${LOG_FILE}"
    else
      echo "Skipping perf: sudo -n not available" | tee -a "\${LOG_FILE}"
    fi
  else
    echo "Skipping perf: command not found" | tee -a "\${LOG_FILE}"
  fi
fi

# Capture previous totals by tid
declare -A PREV_TOTAL
capture_totals() {
  local pid="\$1"
  for statf in /proc/\${pid}/task/*/stat; do
    [ -f "\${statf}" ] || continue
    local tid
    tid=\${statf%/stat}
    tid=\${tid##*/}
    local total
    total=\$(awk '{print \$14+\$15}' "\${statf}" 2>/dev/null || echo 0)
    PREV_TOTAL[\${tid}]="\${total}"
  done
}

capture_totals "\${PID}"

START_TS=\$(date +%s)
END_TS=\$((START_TS + DURATION))

while [ \$(date +%s) -lt \${END_TS} ]; do
  NOW=\$(date +%s)

  for statf in /proc/\${PID}/task/*/stat; do
    [ -f "\${statf}" ] || continue

    tid=\${statf%/stat}
    tid=\${tid##*/}
    total=\$(awk '{print \$14+\$15}' "\${statf}" 2>/dev/null || echo 0)
    prev=\${PREV_TOTAL[\${tid}]:-\${total}}
    delta=\$((total - prev))

    if [ \${delta} -lt 0 ]; then
      delta=0
    fi

    PREV_TOTAL[\${tid}]="\${total}"

    name_file="/proc/\${PID}/task/\${tid}/comm"
    if [ -f "\${name_file}" ]; then
      tname=\$(tr -d '\n' < "\${name_file}")
    else
      tname="unknown"
    fi

    cpu_pct=\$(awk -v d="\${delta}" -v hz="\${CLK_TCK}" -v iv="\${INTERVAL}" 'BEGIN { printf "%.2f", (d / hz) * 100.0 / iv }')
    echo "\${NOW},\${tid},\${tname},\${cpu_pct},\${delta},\${total}" >> "\${CSV_FILE}"
  done

  sleep "\${INTERVAL}"
done

if [ -n "\${PERF_PID}" ]; then
  wait "\${PERF_PID}" || true
  if [ -f "\${PERF_FILE}" ]; then
    sudo perf report --stdio -i "\${PERF_FILE}" --sort=comm,symbol --no-children -n > "\${REMOTE_DIR}/perf_report.txt" 2>/dev/null || true
    sudo chown "\${PERF_OWNER}:\${PERF_OWNER}" "\${PERF_FILE}" "\${REMOTE_DIR}/perf_report.txt" 2>/dev/null || true
  fi
fi

echo "Done. Artifacts in \${REMOTE_DIR}" | tee -a "\${LOG_FILE}"
EOF

echo "Fetching artifacts..."
scp -q "${REMOTE_HOST}:${REMOTE_DIR}/thread_samples.csv" "${LOCAL_CSV}"
scp -q "${REMOTE_HOST}:${REMOTE_DIR}/run.log" "${RUN_DIR}/run.log" || true
scp -q "${REMOTE_HOST}:${REMOTE_DIR}/perf.data" "${RUN_DIR}/perf.data" || true
scp -q "${REMOTE_HOST}:${REMOTE_DIR}/perf_report.txt" "${RUN_DIR}/perf_report.txt" || true

echo "Generating summary..."
{
  echo "=========================================="
  echo "Thread CPU Summary (avg + max over run)"
  echo "=========================================="
  echo "Run timestamp (UTC): ${RUN_TS}"
  echo "Duration: ${DURATION}s, Interval: ${INTERVAL}s"
  echo ""
  awk -F, '
    NR==1 {next}
    {
      key=$2"|"$3;
      cpu=$4+0;
      sum[key]+=cpu;
      cnt[key]++;
      if (cpu > max[key]) max[key]=cpu;
    }
    END {
      printf("%-10s %-18s %12s %12s\n", "TID", "THREAD", "AVG_CPU%", "MAX_CPU%");
      for (k in sum) {
        split(k, a, "|");
        avg=sum[k]/cnt[k];
        printf("%-10s %-18s %12.2f %12.2f\n", a[1], a[2], avg, max[k]);
      }
    }
  ' "${LOCAL_CSV}" | sort -k3,3nr
} | tee "${LOCAL_REPORT}"

echo ""
echo "Completed. Files:"
echo "  ${LOCAL_CSV}"
echo "  ${LOCAL_REPORT}"
if [ -f "${RUN_DIR}/perf_report.txt" ]; then
  echo "  ${RUN_DIR}/perf_report.txt"
fi

echo ""
echo "Top 20 summary:"
head -n 22 "${LOCAL_REPORT}"