// Copyright 2026 The Parca Authors
// SPDX-License-Identifier: Apache-2.0

#include "Driver/GPU/CudaApi.h"
#include "Driver/GPU/CuptiApi.h"
#include "pc_sampling.h"
#include "probes.h"
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <memory>
#include <stdexcept>
#include <tuple>

namespace parcagpu {

// CUDA driver version for 12.8.1 (minimum for PC sampling)
// Version format: major * 1000 + minor * 10 + patch
#define CUDA_VERSION_12_8_1 12081

// CUPTI version that added correlationId to CUpti_PCSamplingPCData,
// breaking ABI compatibility.
#define CUPTI_CUDA12_4_VERSION 22

// noinline wrappers so each USDT probe has exactly one call site in the
// binary.  Multiple call sites produce multiple .note.stapsdt entries with
// different argument encodings, which complicates BPF attachment.
__attribute__((noinline)) void fireCubinLoaded(uint64_t crc, const char *cubin,
                                               uint64_t size) {
  PARCAGPU_CUBIN_LOADED(crc, cubin, size);
}

__attribute__((noinline)) void fireCubinUnloaded(uint64_t crc) {
  PARCAGPU_CUBIN_UNLOADED(crc);
}

__attribute__((noinline)) void fireGpuConfig(uint32_t dev, uint32_t factor,
                                             uint32_t clockKHz,
                                             uint32_t smCount) {
  PARCAGPU_GPU_CONFIG(dev, factor, clockKHz, smCount);
}

static uint64_t emitMetadataNowNs();

// Max records per pc_sample_batch probe invocation.
static constexpr uint32_t PCSampleBatchSize = 128;

__attribute__((noinline)) void firePCSampleBatch(
    const void **ptrs, uint32_t count) {
  PARCAGPU_PC_SAMPLE_BATCH(ptrs, count);
}

namespace {

// CUPTI helper functions (adapted from Proton's CuptiPCSamplingUtils.h)
// These wrap Proton's cupti API calls with PARCAGPU-specific setup

uint64_t getCubinCrc(const char *cubin, size_t size) {
  CUpti_GetCubinCrcParams cubinCrcParams = {
      /*size=*/CUpti_GetCubinCrcParamsSize,
      /*cubinSize=*/size,
      /*cubin=*/cubin,
      /*cubinCrc=*/0,
  };
  proton::cupti::getCubinCrc<true>(&cubinCrcParams);
  return cubinCrcParams.cubinCrc;
}

void enablePCSampling(CUcontext context) {
  CUpti_PCSamplingEnableParams params = {
      /*size=*/CUpti_PCSamplingEnableParamsSize,
      /*pPriv=*/NULL,
      /*ctx=*/context,
  };
  proton::cupti::pcSamplingEnable<true>(&params);
}

bool startPCSampling(CUcontext context) {
  // CUPTI requires the GPU to be idle before starting PC sampling.
  proton::cuda::ctxSynchronize<false>();
  CUpti_PCSamplingStartParams params = {
      /*size=*/CUpti_PCSamplingStartParamsSize,
      /*pPriv=*/NULL,
      /*ctx=*/context,
  };
  auto ret = proton::cupti::pcSamplingStart<false>(&params);
  if (ret != CUPTI_SUCCESS) {
    DEBUG_PRINTF("cuptiPCSamplingStart failed: %d\n", ret);
    return false;
  }
  return true;
}

bool stopPCSampling(CUcontext context) {
  CUpti_PCSamplingStopParams params = {
      /*size=*/CUpti_PCSamplingStopParamsSize,
      /*pPriv=*/NULL,
      /*ctx=*/context,
  };
  auto ret = proton::cupti::pcSamplingStop<false>(&params);
  if (ret != CUPTI_SUCCESS) {
    DEBUG_PRINTF("cuptiPCSamplingStop failed: %d\n", ret);
    return false;
  }
  return true;
}

void disablePCSampling(CUcontext context) {
  CUpti_PCSamplingDisableParams params = {
      /*size=*/CUpti_PCSamplingDisableParamsSize,
      /*pPriv=*/NULL,
      /*ctx=*/context,
  };
  // Use <false>: invoked from a libcupti callback whose caller wasn't built
  // with C++ EH, so a throw here would propagate up and abort the process.
  auto ret = proton::cupti::pcSamplingDisable<false>(&params);
  if (ret != CUPTI_SUCCESS) {
    DEBUG_PRINTF("cuptiPCSamplingDisable failed: %d (ctx=%p)\n", ret, context);
  }
}

// Returns CUPTI_SUCCESS on success, or the raw CUptiResult on failure so
// callers can distinguish CUPTI_ERROR_OUT_OF_MEMORY (error 8) — a wedged
// internal-buffer condition we recover from — from other errors.
CUptiResult getPCSamplingData(CUcontext context,
                              CUpti_PCSamplingData *pcSamplingData) {
  CUpti_PCSamplingGetDataParams params = {
      /*size=*/CUpti_PCSamplingGetDataParamsSize,
      /*pPriv=*/NULL,
      /*ctx=*/context,
      /*pcSamplingData=*/pcSamplingData,
  };
  auto result = proton::cupti::pcSamplingGetData<false>(&params);
  if (result != CUPTI_SUCCESS) {
    DEBUG_PRINTF("cuptiPCSamplingGetData failed: error %d (ctx=%p)\n", result,
                 context);
  }
  return result;
}

void setConfigurationAttribute(
    CUcontext context,
    std::vector<CUpti_PCSamplingConfigurationInfo> &configurationInfos) {
  CUpti_PCSamplingConfigurationInfoParams infoParams = {
      /*size=*/CUpti_PCSamplingConfigurationInfoParamsSize,
      /*pPriv=*/NULL,
      /*ctx=*/context,
      /*numAttributes=*/configurationInfos.size(),
      /*pPCSamplingConfigurationInfo=*/configurationInfos.data(),
  };
  proton::cupti::pcSamplingSetConfigurationAttribute<true>(&infoParams);
}

std::tuple<uint32_t, std::string, std::string>
getSassToSourceCorrelation(const char *functionName, uint64_t pcOffset,
                           const char *cubin, size_t cubinSize) {
  CUpti_GetSassToSourceCorrelationParams sassToSourceParams = {
      /*size=*/CUpti_GetSassToSourceCorrelationParamsSize,
      /*cubin=*/cubin,
      /*functionName=*/functionName,
      /*cubinSize=*/cubinSize,
      /*lineNumber=*/0,
      /*pcOffset=*/pcOffset,
      /*fileName=*/NULL,
      /*dirName=*/NULL,
  };
  // Get source can fail if the line mapping is not available
  proton::cupti::getSassToSourceCorrelation<false>(&sassToSourceParams);
  auto fileNameStr = sassToSourceParams.fileName
                         ? std::string(sassToSourceParams.fileName)
                         : "";
  auto dirNameStr =
      sassToSourceParams.dirName ? std::string(sassToSourceParams.dirName) : "";
  // Free the memory
  if (sassToSourceParams.fileName)
    std::free(sassToSourceParams.fileName);
  if (sassToSourceParams.dirName)
    std::free(sassToSourceParams.dirName);
  return std::make_tuple(sassToSourceParams.lineNumber, fileNameStr,
                         dirNameStr);
}

// Double-checked locking helper
template <typename CheckFn, typename ActionFn>
void doubleCheckedLock(CheckFn check, std::mutex &mutex, ActionFn action) {
  if (check()) {
    std::lock_guard<std::mutex> lock(mutex);
    if (check()) {
      action();
    }
  }
}

// Helper to get PARCAGPU's custom sampling frequency from environment
uint32_t getGPUSamplingFrequency() {
  // Default frequency for PARCAGPU is 20 (Proton uses 10)
  constexpr uint32_t PARCAGPU_DEFAULT_FREQUENCY = 20;

  uint32_t samplingPeriod = PARCAGPU_DEFAULT_FREQUENCY;
  const char *sampling_factor_env = getenv("PARCAGPU_SAMPLING_FACTOR");
  if (sampling_factor_env) {
    int factor = atoi(sampling_factor_env);
    if (factor >= 5 && factor <= 31) {
      samplingPeriod = factor;
      DEBUG_PRINTF("Using PARCAGPU_SAMPLING_FACTOR=%u\n", samplingPeriod);
    } else if (factor != 0) {
      fprintf(stderr,
              "[PARCAGPU] Warning: PARCAGPU_SAMPLING_FACTOR=%d out of range "
              "[5,31], using default %u\n",
              factor, PARCAGPU_DEFAULT_FREQUENCY);
    }
  }
  return samplingPeriod;
}

// Get number of stall reasons
size_t getNumStallReasons(CUcontext context) {
  size_t numStallReasons = 0;
  CUpti_PCSamplingGetNumStallReasonsParams numStallReasonsParams = {
      /*size=*/CUpti_PCSamplingGetNumStallReasonsParamsSize,
      /*pPriv=*/NULL,
      /*ctx=*/context,
      /*numStallReasons=*/&numStallReasons};
  proton::cupti::pcSamplingGetNumStallReasons<true>(&numStallReasonsParams);
  return numStallReasons;
}

// Get stall reason names and indices
std::pair<char **, uint32_t *>
getStallReasonNamesAndIndices(CUcontext context, size_t numStallReasons) {
  char **stallReasonNames =
      static_cast<char **>(std::calloc(numStallReasons, sizeof(char *)));
  for (size_t i = 0; i < numStallReasons; i++) {
    stallReasonNames[i] = static_cast<char *>(
        std::calloc(CUPTI_STALL_REASON_STRING_SIZE, sizeof(char)));
  }
  uint32_t *stallReasonIndices =
      static_cast<uint32_t *>(std::calloc(numStallReasons, sizeof(uint32_t)));
  CUpti_PCSamplingGetStallReasonsParams stallReasonsParams = {
      /*size=*/CUpti_PCSamplingGetStallReasonsParamsSize,
      /*pPriv=*/NULL,
      /*ctx=*/context,
      /*numStallReasons=*/numStallReasons,
      /*stallReasonIndex=*/stallReasonIndices,
      /*stallReasons=*/stallReasonNames,
  };
  proton::cupti::pcSamplingGetStallReasons<true>(&stallReasonsParams);
  return std::make_pair(stallReasonNames, stallReasonIndices);
}

// Match stall reasons to indices (PARCAGPU emits all stall reasons)
size_t matchStallReasonsToIndices(
    size_t numStallReasons, char **stallReasonNames,
    uint32_t *stallReasonIndices,
    std::map<size_t, size_t> &stallReasonIndexToMetricIndex,
    std::set<size_t> &notIssuedStallReasonIndices) {
  // PARCAGPU emits all stall reasons
  size_t numValidStalls = 0;
  for (size_t i = 0; i < numStallReasons; i++) {
    std::string cuptiStallName = std::string(stallReasonNames[i]);
    bool notIssued = cuptiStallName.find("not_issued") != std::string::npos ||
                     cuptiStallName.find("Not Issued") != std::string::npos;

    if (notIssued)
      notIssuedStallReasonIndices.insert(stallReasonIndices[i]);
    stallReasonIndexToMetricIndex[stallReasonIndices[i]] = i;
    numValidStalls++;
  }
  return numValidStalls;
}

// Allocate PC sampling data buffer
CUpti_PCSamplingData allocPCSamplingData(size_t collectNumPCs,
                                         size_t numValidStallReasons) {
  CUpti_PCSamplingData pcSamplingData{
      /*size=*/sizeof(CUpti_PCSamplingData),
      /*collectNumPcs=*/collectNumPCs,
      /*totalSamples=*/0,
      /*droppedSamples=*/0,
      /*totalNumPcs=*/0,
      /*remainingNumPcs=*/0,
      /*rangeId=*/0,
      /*pPcData=*/
      static_cast<CUpti_PCSamplingPCData *>(
          std::calloc(collectNumPCs, sizeof(CUpti_PCSamplingPCData)))};
  for (size_t i = 0; i < collectNumPCs; ++i) {
    pcSamplingData.pPcData[i].size = sizeof(CUpti_PCSamplingPCData);
    pcSamplingData.pPcData[i].stallReason =
        static_cast<CUpti_PCSamplingStallReason *>(std::calloc(
            numValidStallReasons, sizeof(CUpti_PCSamplingStallReason)));
  }
  return pcSamplingData;
}

} // namespace

// ConfigureData implementation

CUpti_PCSamplingConfigurationInfo ConfigureData::configureStallReasons() {
  numStallReasons = getNumStallReasons(context);
  std::tie(this->stallReasonNames, this->stallReasonIndices) =
      getStallReasonNamesAndIndices(context, numStallReasons);
  numValidStallReasons = matchStallReasonsToIndices(
      numStallReasons, stallReasonNames, stallReasonIndices,
      stallReasonIndexToMetricIndex, notIssuedStallReasonIndices);

  CUpti_PCSamplingConfigurationInfo stallReasonInfo{};
  stallReasonInfo.attributeType =
      CUPTI_PC_SAMPLING_CONFIGURATION_ATTR_TYPE_STALL_REASON;
  stallReasonInfo.attributeData.stallReasonData.stallReasonCount =
      numValidStallReasons;
  stallReasonInfo.attributeData.stallReasonData.pStallReasonIndex =
      stallReasonIndices;
  return stallReasonInfo;
}

CUpti_PCSamplingConfigurationInfo ConfigureData::configureSamplingPeriod() {
  CUpti_PCSamplingConfigurationInfo samplingPeriodInfo{};
  samplingPeriodInfo.attributeType =
      CUPTI_PC_SAMPLING_CONFIGURATION_ATTR_TYPE_SAMPLING_PERIOD;

  // Use PARCAGPU's custom sampling frequency
  uint32_t frequency = getGPUSamplingFrequency();

  samplingPeriodInfo.attributeData.samplingPeriodData.samplingPeriod =
      frequency;
  return samplingPeriodInfo;
}

CUpti_PCSamplingConfigurationInfo ConfigureData::configureSamplingBuffer() {
  CUpti_PCSamplingConfigurationInfo samplingBufferInfo{};
  samplingBufferInfo.attributeType =
      CUPTI_PC_SAMPLING_CONFIGURATION_ATTR_TYPE_SAMPLING_DATA_BUFFER;
  this->pcSamplingData =
      allocPCSamplingData(DataBufferPCCount, numValidStallReasons);
  samplingBufferInfo.attributeData.samplingDataBufferData.samplingDataBuffer =
      &this->pcSamplingData;
  return samplingBufferInfo;
}

CUpti_PCSamplingConfigurationInfo ConfigureData::configureScratchBuffer() {
  CUpti_PCSamplingConfigurationInfo scratchBufferInfo{};
  scratchBufferInfo.attributeType =
      CUPTI_PC_SAMPLING_CONFIGURATION_ATTR_TYPE_SCRATCH_BUFFER_SIZE;
  scratchBufferInfo.attributeData.scratchBufferSizeData.scratchBufferSize =
      ScratchBufferSize;
  return scratchBufferInfo;
}

CUpti_PCSamplingConfigurationInfo ConfigureData::configureHardwareBufferSize() {
  CUpti_PCSamplingConfigurationInfo hardwareBufferInfo{};
  hardwareBufferInfo.attributeType =
      CUPTI_PC_SAMPLING_CONFIGURATION_ATTR_TYPE_HARDWARE_BUFFER_SIZE;
  hardwareBufferInfo.attributeData.hardwareBufferSizeData.hardwareBufferSize =
      HardwareBufferSize;
  return hardwareBufferInfo;
}

CUpti_PCSamplingConfigurationInfo ConfigureData::configureCollectionMode() {
  CUpti_PCSamplingConfigurationInfo collectionModeInfo{};
  collectionModeInfo.attributeType =
      CUPTI_PC_SAMPLING_CONFIGURATION_ATTR_TYPE_COLLECTION_MODE;
  collectionModeInfo.attributeData.collectionModeData.collectionMode =
      CUPTI_PC_SAMPLING_COLLECTION_MODE_KERNEL_SERIALIZED;
  return collectionModeInfo;
}

CUpti_PCSamplingConfigurationInfo ConfigureData::configureStartStopControl() {
  CUpti_PCSamplingConfigurationInfo startStopControlInfo{};
  startStopControlInfo.attributeType =
      CUPTI_PC_SAMPLING_CONFIGURATION_ATTR_TYPE_ENABLE_START_STOP_CONTROL;
  startStopControlInfo.attributeData.enableStartStopControlData
      .enableStartStopControl = true;
  return startStopControlInfo;
}

void ConfigureData::initialize(CUcontext context) {
  this->context = context;
  proton::cupti::getContextId<true>(context, &contextId);

  DEBUG_PRINTF("Initializing PC sampling for context %p (id %u)\n", context,
               contextId);

  configurationInfos.emplace_back(configureStallReasons());
  configurationInfos.emplace_back(configureCollectionMode());
  configurationInfos.emplace_back(configureStartStopControl());
  configurationInfos.emplace_back(configureSamplingBuffer());
  // Bigger scratch + hardware buffers so a busy workload (PyTorch-class)
  // doesn't overflow CUPTI's defaults within minutes and start returning
  // CUPTI_ERROR_OUT_OF_MEMORY from cuptiPCSamplingGetData.
  configurationInfos.emplace_back(configureScratchBuffer());
  configurationInfos.emplace_back(configureHardwareBufferSize());
  // Don't set sampling period — let CUPTI use its default.
  // Explicit period values silently break sampling on some GPUs (e.g.
  // Blackwell).

  setConfigurationAttribute(context, configurationInfos);

  // Allocate a separate output buffer for getPCSamplingData calls.
  // The configured pcSamplingData buffer is owned by CUPTI internally;
  // we must pass a different buffer to getPCSamplingData.
  this->outputData =
      allocPCSamplingData(DataBufferPCCount, numValidStallReasons);

  DEBUG_PRINTF("PC sampling configured with %u stall reasons (%u valid)\n",
               numStallReasons, numValidStallReasons);
}

// GPUPCSampling implementation

bool PCSampling::isSupported() {
  // PC sampling is off by default. Setting PARCAGPU_PC_SAMPLING_RATE to a
  // non-negative number opts in; the user does not need to also set
  // PARCAGPU_SAMPLING_FACTOR (it has a default).
  const char *env = getenv("PARCAGPU_PC_SAMPLING_RATE");
  if (!env) {
    DEBUG_PRINTF("PC sampling disabled (set PARCAGPU_PC_SAMPLING_RATE to enable)\n");
    return false;
  }

  // Check CUDA driver version >= 12.8.1
  int driverVersion = 0;
  proton::cuda::driverGetVersion<true>(&driverVersion);

  if (driverVersion < CUDA_VERSION_12_8_1) {
    int major = driverVersion / 1000;
    int minor = (driverVersion % 1000) / 10;
    int patch = driverVersion % 10;
    DEBUG_PRINTF("PC sampling not supported: CUDA driver version %d.%d.%d < "
                 "required 12.8.1\n",
                 major, minor, patch);
    fireError(driverVersion,
              "CUDA driver version too low for PC sampling (need >= 12.8.1)",
              "pc_sampling");
    return false;
  }

  // Check CUPTI API/driver version compatibility.
  // CUPTI 12.4 (v22) added correlationId to CUpti_PCSamplingPCData, breaking
  // ABI. Mixing compile-time and runtime versions across this boundary crashes.
  uint32_t cuptiVersion = 0;
  proton::cupti::getVersion<true>(&cuptiVersion);

  if ((cuptiVersion < CUPTI_CUDA12_4_VERSION &&
       CUPTI_API_VERSION >= CUPTI_CUDA12_4_VERSION) ||
      (cuptiVersion >= CUPTI_CUDA12_4_VERSION &&
       CUPTI_API_VERSION < CUPTI_CUDA12_4_VERSION)) {
    DEBUG_PRINTF(
        "PC sampling disabled: CUPTI API version %d and driver version %d "
        "are incompatible across the 12.4 (v22) ABI boundary\n",
        CUPTI_API_VERSION, cuptiVersion);
    fireError((int32_t)cuptiVersion,
              "CUPTI API/driver version mismatch (12.4 ABI boundary)",
              "pc_sampling");
    return false;
  }

  // Attempt a lightweight permission probe. CUPTI PC sampling requires
  // either root, CAP_SYS_ADMIN, or the NVIDIA module parameter
  // NVreg_RestrictProfilingToAdminUsers=0.
  // We cannot easily pre-check permissions without attempting CUPTI calls,
  // so we defer the real check to initialize() where enablePCSampling()
  // will fail with a CUPTI error if permissions are insufficient.
  // TODO: Add explicit permission pre-check.
  // Reference: https://developer.nvidia.com/nvidia-development-tools-solutions-err_nvgpuctrperm-permission-issue-performance-counters

  int major = driverVersion / 1000;
  int minor = (driverVersion % 1000) / 10;
  int patch = driverVersion % 10;
  DEBUG_PRINTF("PC sampling supported: CUDA %d.%d.%d, CUPTI v%u (API v%d)\n",
               major, minor, patch, cuptiVersion, CUPTI_API_VERSION);
  return true;
}

ConfigureData *PCSampling::getConfigureData(uint32_t contextId) {
  return &contextIdToConfigureData[contextId];
}

CubinData *PCSampling::getCubinData(uint64_t cubinCrc) {
  return &(cubinCrcToCubinData[cubinCrc].first);
}

void PCSampling::initialize(CUcontext context) {
  uint32_t contextId = 0;
  proton::cupti::getContextId<true>(context, &contextId);

  doubleCheckedLock(
      [&]() {
        return !contextInitialized.contain(contextId) &&
               !contextFailed.contain(contextId);
      },
      contextMutex,
      [&]() {
        // enablePCSampling can fail due to insufficient permissions
        // (ERR_NVGPUCTRPERM). Catch and degrade gracefully.
        CUpti_PCSamplingEnableParams enableParams = {
            /*size=*/CUpti_PCSamplingEnableParamsSize,
            /*pPriv=*/NULL,
            /*ctx=*/context,
        };
        auto result =
            proton::cupti::pcSamplingEnable<false>(&enableParams);
        if (result != CUPTI_SUCCESS) {
          DEBUG_PRINTF(
              "Failed to enable PC sampling for context %u: CUPTI error %d\n"
              "This may be a permission issue. See:\n"
              "https://developer.nvidia.com/nvidia-development-tools-solutions-"
              "err_nvgpuctrperm-permission-issue-performance-counters\n",
              contextId, result);
          fireError((int32_t)result,
                    "Failed to enable PC sampling (possible permission issue)",
                    "pc_sampling");
          contextFailed.insert(contextId);
          return;
        }

        auto *configData = getConfigureData(contextId);
        configData->initialize(context);

        // Build contiguous stall reason map for USDT probe emission.
        stallReasonMap.build(configData->numStallReasons,
                             configData->stallReasonIndices,
                             configData->stallReasonNames);

        contextInitialized.insert(contextId);
        initializedContextIds.push_back(contextId);
        DEBUG_PRINTF(
            "PC sampling initialized (serialized mode) for context %u\n",
            contextId);

        // Capture per-device config so the agent can convert PC sample counts
        // to nanoseconds. cuCtxGetDevice reads the current context; the
        // CUPTI ENTER callback that drove us here ran with `context` current.
        // contextMutex is already held by doubleCheckedLock above, so push
        // into loadedConfigs without re-locking (std::mutex is not recursive).
        CUdevice dev = 0;
        int clockKHz = 0;
        int smCount = 0;
        proton::cuda::ctxGetDevice<false>(&dev);
        proton::cuda::deviceGetAttribute<false>(
            &clockKHz, CU_DEVICE_ATTRIBUTE_CLOCK_RATE, dev);
        proton::cuda::deviceGetAttribute<false>(
            &smCount, CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT, dev);
        const uint32_t factor = getGPUSamplingFrequency();
        const uint64_t emittedNs = emitMetadataNowNs();
        fireGpuConfig((uint32_t)dev, factor, (uint32_t)clockKHz,
                      (uint32_t)smCount);
        loadedConfigs.push_back({(uint32_t)dev, factor, (uint32_t)clockKHz,
                                 (uint32_t)smCount, emittedNs});
      });
}

void PCSampling::start(CUcontext context) {
  std::lock_guard<std::mutex> lock(pcSamplingMutex);
  if (samplingActive)
    return;
  if (startPCSampling(context)) {
    samplingActive = true;
    samplingContext = context;
    DEBUG_PRINTF("PC sampling started (kernels serialized)\n");
  }
}

void PCSampling::stop(CUcontext context) {
  std::lock_guard<std::mutex> lock(pcSamplingMutex);
  if (!samplingActive)
    return;
  stopPCSampling(context);
  samplingActive = false;
  DEBUG_PRINTF("PC sampling stopped (kernels concurrent)\n");
  // Drain data collected during this window.
  collectData(context);
}

__attribute__((noinline)) void fireError(int32_t code, const char *message,
                                         const char *component) {
  PARCAGPU_ERROR(code, message, component);
}

void PCSampling::processPCSamplingData(ConfigureData *configureData) {
  auto *pcSamplingData = &configureData->outputData;

  if (pcSamplingData->totalNumPcs == 0) {
    return;
  }

  DEBUG_PRINTF("Processing %zu PCs (remaining: %zu)\n",
               pcSamplingData->totalNumPcs, pcSamplingData->remainingNumPcs);

  // Count distinct (PC, stallReason) pairs with non-zero samples — that is
  // the unit the agent sees as a "sample" record on the receive side, and the
  // right rate to steer with the controller. The raw CUPTI per-cell hardware
  // counts are weights, not records.
  uint64_t batchSamples = 0;
  for (size_t i = 0; i < pcSamplingData->totalNumPcs; ++i) {
    auto *pcData = pcSamplingData->pPcData + i;

    uint64_t totalSamples = 0;
    uint64_t stalledSamples = 0;
    uint64_t distinctPairs = 0;
    for (size_t j = 0; j < pcData->stallReasonCount; ++j) {
      auto *stallReason = &pcData->stallReason[j];
      if (stallReason->samples == 0)
        continue;
      ++distinctPairs;
      totalSamples += stallReason->samples;
      bool isNotIssued = configureData->notIssuedStallReasonIndices.count(
                             stallReason->pcSamplingStallReasonIndex) > 0;
      if (!isNotIssued)
        stalledSamples += stallReason->samples;
    }
    batchSamples += distinctPairs;

    if (debug_enabled) {
      auto *cubinData = getCubinData(pcData->cubinCrc);
      auto key =
          CubinData::LineInfoKey{pcData->functionIndex, pcData->pcOffset};
      if (cubinData->lineInfo.find(key) == cubinData->lineInfo.end()) {
        auto [lineNumber, fileName, dirName] =
            getSassToSourceCorrelation(pcData->functionName, pcData->pcOffset,
                                       cubinData->cubin, cubinData->cubinSize);
        cubinData->lineInfo.try_emplace(key, lineNumber,
                                        std::string(pcData->functionName),
                                        dirName, fileName);
      }
      auto &lineInfo = cubinData->lineInfo[key];
      std::string fullPath = lineInfo.fileName.size()
                                 ? lineInfo.dirName + "/" + lineInfo.fileName
                                 : "";
      DEBUG_PRINTF("  [%zu] func=%s pc=0x%lx total=%lu stalled=%lu %s:%u\n", i,
                   lineInfo.functionName.c_str(), pcData->pcOffset,
                   totalSamples, stalledSamples, fullPath.c_str(),
                   lineInfo.lineNumber);
    }
  }
  recordPCSamples(batchSamples);

  // Emit batched PC sample probes as a bag of pointers (like activity_batch).
  // Using pointers avoids depending on the CUPTI struct stride, which can
  // change across CUDA versions.
  const void *batchPtrs[PCSampleBatchSize];
  uint32_t batchCount = 0;

  for (size_t i = 0; i < pcSamplingData->totalNumPcs; ++i) {
    batchPtrs[batchCount++] = &pcSamplingData->pPcData[i];
    if (batchCount == PCSampleBatchSize) {
      firePCSampleBatch(batchPtrs, batchCount);
      batchCount = 0;
    }
  }
  if (batchCount > 0) {
    firePCSampleBatch(batchPtrs, batchCount);
  }
}

// Period between forced metadata refreshes regardless of probe-arm
// transitions. Bounds how long an undetected ABA semaphore cycle (consumer
// detaches and re-attaches between two of our reads) can leave a tracer
// missing cubins or the stall-reason map.
static constexpr uint64_t kEmitRefreshPeriodNs = 30ULL * 1000 * 1000 * 1000;

static uint64_t emitMetadataNowNs() {
  struct timespec ts;
  clock_gettime(CLOCK_MONOTONIC, &ts);
  return uint64_t(ts.tv_sec) * 1000000000ULL + uint64_t(ts.tv_nsec);
}

// Read the raw USDT semaphore (a refcount of attached uprobe consumers).
// __atomic_load_n with relaxed ordering — the kernel updates this from
// another address space via RefCtrOffset; this prevents the compiler from
// caching the value across calls.
static inline uint16_t readSem(unsigned short &sem) {
  return __atomic_load_n(reinterpret_cast<uint16_t *>(&sem), __ATOMIC_RELAXED);
}

void PCSampling::emitMetadata() {
  // Re-emit stall_reason_map and cubin_loaded for late-attaching tracers.
  //
  // The USDT semaphore is a refcount of attached uprobe consumers — every
  // attach increments, every detach decrements. Re-emit whenever the count
  // increases (so 1→2 fires the metadata to a second consumer joining
  // while the first is still attached), plus on a periodic refresh that
  // bounds staleness from ABA cycles which net the same count between our
  // reads.

  const uint64_t nowNs = emitMetadataNowNs();
  const uint64_t lastRefresh = lastRefreshNs.load(std::memory_order_relaxed);
  const bool refresh = (nowNs - lastRefresh) > kEmitRefreshPeriodNs;
  if (refresh) {
    lastRefreshNs.store(nowNs, std::memory_order_relaxed);
  }

  // Stall reason map (small, ~4KB) — no per-entry state needed.
  const uint16_t stallSem = readSem(parcagpu_stall_reason_map_semaphore);
  const uint16_t prevStall =
      prevStallSem.exchange(stallSem, std::memory_order_acq_rel);
  const bool stallJoin = stallSem > prevStall;
  if (stallReasonMap.data() && stallSem > 0 && (stallJoin || refresh)) {
    PARCAGPU_STALL_REASON_MAP(stallReasonMap.data(),
                              stallReasonMap.numEntries());
  }

  // Cubin loaded — re-fire all cubins on a consumer-join (count increase);
  // on a periodic refresh, re-fire only entries that have gone stale.
  const uint16_t cubinSem = readSem(parcagpu_cubin_loaded_semaphore);
  const uint16_t prevCubin =
      prevCubinSem.exchange(cubinSem, std::memory_order_acq_rel);
  const bool cubinJoin = cubinSem > prevCubin;
  if (cubinSem > 0 && (cubinJoin || refresh)) {
    std::lock_guard<std::mutex> lock(contextMutex);
    for (auto &ref : loadedCubins) {
      if (cubinJoin ||
          (nowNs - ref.lastEmittedNs) > kEmitRefreshPeriodNs) {
        fireCubinLoaded(ref.crc, ref.data, ref.size);
        ref.lastEmittedNs = nowNs;
      }
    }
  }

  // GPU config — same edge-on-consumer-join + stale-refresh pattern as cubins.
  const uint16_t cfgSem = readSem(parcagpu_gpu_config_semaphore);
  const uint16_t prevCfg =
      prevConfigSem.exchange(cfgSem, std::memory_order_acq_rel);
  const bool cfgJoin = cfgSem > prevCfg;
  if (cfgSem > 0 && (cfgJoin || refresh)) {
    std::lock_guard<std::mutex> lock(contextMutex);
    for (auto &ref : loadedConfigs) {
      if (cfgJoin ||
          (nowNs - ref.lastEmittedNs) > kEmitRefreshPeriodNs) {
        fireGpuConfig(ref.dev, ref.samplingFactor, ref.clockKHz, ref.smCount);
        ref.lastEmittedNs = nowNs;
      }
    }
  }
}

void PCSampling::collectData(CUcontext context) {
  uint32_t contextId = 0;
  proton::cupti::getContextId<true>(context, &contextId);

  if (!contextInitialized.contain(contextId)) {
    DEBUG_PRINTF("Context %u not initialized, skipping data collection\n",
                 contextId);
    return;
  }

  auto *configureData = getConfigureData(contextId);
  DEBUG_PRINTF("Collecting PC sampling data for context %u (cfg total=%zu "
               "remaining=%zu)\n",
               contextId, configureData->pcSamplingData.totalNumPcs,
               configureData->pcSamplingData.remainingNumPcs);

  // Drain all available PCs in a loop. Each getData call returns at most
  // DataBufferPCCount (1024) PCs; a single sampling window can produce
  // tens of thousands. Failing to drain leaves data in CUPTI's internal
  // buffers, which eventually causes CUPTI_ERROR_OUT_OF_MEMORY (error 8).
  do {
    CUptiResult res = getPCSamplingData(context, &configureData->outputData);
    DEBUG_PRINTF("getData: res=%d output total=%zu remaining=%zu "
                 "cfg total=%zu remaining=%zu\n",
                 res, configureData->outputData.totalNumPcs,
                 configureData->outputData.remainingNumPcs,
                 configureData->pcSamplingData.totalNumPcs,
                 configureData->pcSamplingData.remainingNumPcs);
    if (res == CUPTI_SUCCESS) {
      processPCSamplingData(configureData);
      continue;
    }
    if (res == CUPTI_ERROR_OUT_OF_MEMORY) {
      // CUPTI's PC-sampling state for this context is wedged — stop+start
      // alone does NOT unwedge it (confirmed empirically: 7800 starts
      // produced 0 samples after the first OOM). The fix is a full reset:
      // disable + erase tracking + re-enable + re-configure, which is
      // exactly what finalize() + initialize() do back-to-back.
      //
      // try_lock-probe on pcSamplingMutex first: collectData is sometimes
      // called from inside PCSampling::stop() which already holds the
      // mutex. Calling finalize() in that re-entry would self-deadlock
      // on the non-recursive mutex (finalize re-takes it). If the probe
      // fails (someone holds it), bail; the next callback not nested
      // under stop()/start() will retry.
      {
        std::unique_lock<std::mutex> probe(pcSamplingMutex, std::try_to_lock);
        if (!probe.owns_lock()) {
          break;
        }
      }
      // probe released; finalize() + initialize() each take their own
      // locks (contextMutex, and pcSamplingMutex internally).
      DEBUG_PRINTF("Recovering from CUPTI_ERROR_OUT_OF_MEMORY: full reinit on "
                   "ctx=%p (disable+enable to clear wedged state)\n", context);
      finalize(context);
      initialize(context);
    }
    break;
  } while (configureData->outputData.remainingNumPcs > 0);
}

void PCSampling::collectAllData() {
  std::lock_guard<std::mutex> lock(contextMutex);
  for (auto contextId : initializedContextIds) {
    auto result = contextIdToConfigureData.find(contextId);
    if (!result) {
      DEBUG_PRINTF("Context %u in initializedContextIds but not in map, "
                   "skipping\n",
                   contextId);
      continue;
    }
    auto *configureData = &result->get();
    DEBUG_PRINTF("Draining PC sampling data for context %u\n", contextId);
    // Fetch and drain all pending data from CUPTI.
    do {
      if (getPCSamplingData(configureData->context,
                            &configureData->outputData) != CUPTI_SUCCESS) {
        break;
      }
      processPCSamplingData(configureData);
    } while (configureData->outputData.remainingNumPcs > 0);
  }
}

void PCSampling::finalize(CUcontext context) {
  uint32_t contextId = 0;
  proton::cupti::getContextId<true>(context, &contextId);

  if (!contextInitialized.contain(contextId)) {
    // Clean up failed context tracking if applicable.
    contextFailed.erase(contextId);
    return;
  }

  // Hold contextMutex for the entire finalize to prevent collectAllData
  // from racing with us (it iterates initializedContextIds under this lock).
  std::lock_guard<std::mutex> lock(contextMutex);

  DEBUG_PRINTF("Finalizing PC sampling for context %p\n", context);

  // Remove from iteration list first so collectAllData won't touch this context
  initializedContextIds.erase(std::remove(initializedContextIds.begin(),
                                          initializedContextIds.end(),
                                          contextId),
                              initializedContextIds.end());

  // Stop sampling if it was started on this context.
  {
    std::lock_guard<std::mutex> lock2(pcSamplingMutex);
    if (samplingActive) {
      stopPCSampling(context);
      samplingActive = false;
    }
  }

  // Drain all remaining PC data before disabling.
  auto *configureData = getConfigureData(contextId);
  do {
    if (getPCSamplingData(context, &configureData->outputData) !=
        CUPTI_SUCCESS) {
      break;
    }
    processPCSamplingData(configureData);
  } while (configureData->outputData.remainingNumPcs > 0);

  disablePCSampling(context);

  contextIdToConfigureData.erase(contextId);
  contextInitialized.erase(contextId);
}

void PCSampling::loadModule(const char *cubin, size_t cubinSize) {
  auto cubinCrc = getCubinCrc(cubin, cubinSize);

  if (cubinCrcToCubinData.contain(cubinCrc)) {
    // Increment reference count
    cubinCrcToCubinData[cubinCrc].second++;
    DEBUG_PRINTF("Module 0x%lx loaded (refcount=%zu)\n", cubinCrc,
                 cubinCrcToCubinData[cubinCrc].second);
  } else {
    // New module — getCubinData after the contain() check so operator[]
    // doesn't auto-insert before we test.
    auto *cubinData = getCubinData(cubinCrc);
    cubinData->cubinCrc = cubinCrc;
    cubinData->cubinSize = cubinSize;
    cubinData->cubin = cubin;
    cubinCrcToCubinData[cubinCrc].second = 1;
    DEBUG_PRINTF("Module 0x%lx loaded (new)\n", cubinCrc);
    const uint64_t emittedNs = emitMetadataNowNs();
    fireCubinLoaded(cubinCrc, cubin, cubinSize);
    {
      std::lock_guard<std::mutex> lock(contextMutex);
      loadedCubins.push_back({cubinCrc, cubin, cubinSize, emittedNs});
    }
  }
}

void PCSampling::unloadModule(const char *cubin, size_t cubinSize) {
  auto cubinCrc = getCubinCrc(cubin, cubinSize);

  if (!cubinCrcToCubinData.contain(cubinCrc))
    return;

  auto count = cubinCrcToCubinData[cubinCrc].second;
  if (count > 1) {
    cubinCrcToCubinData[cubinCrc].second = count - 1;
    DEBUG_PRINTF("Module 0x%lx unloaded (refcount=%zu)\n", cubinCrc, count - 1);
  } else {
    cubinCrcToCubinData.erase(cubinCrc);
    DEBUG_PRINTF("Module 0x%lx unloaded (removed)\n", cubinCrc);
    fireCubinUnloaded(cubinCrc);
    {
      std::lock_guard<std::mutex> lock(contextMutex);
      loadedCubins.erase(std::remove_if(loadedCubins.begin(),
                                        loadedCubins.end(),
                                        [cubinCrc](const CubinRef &r) {
                                          return r.crc == cubinCrc;
                                        }),
                         loadedCubins.end());
    }
  }
}

} // namespace parcagpu
