//
// An example collector/tool that prints out just the SYCL UR layer trace
// events
//
#include "xpti/xpti_trace_framework.h"

#include <chrono>
#include <cstdio>
#include <iostream>
#include <map>
#include <mutex>
#include <string>
#include <thread>
#include <unordered_map>

static uint8_t GStreamID = 0;
std::mutex GIOMutex;

// The lone callback function we are going to use to demonstrate how to attach
// the collector to the running executable
XPTI_CALLBACK_API void tpCallback(uint16_t trace_type,
                                  xpti::trace_event_data_t *parent,
                                  xpti::trace_event_data_t *event,
                                  uint64_t instance, const void *user_data);

// Based on the documentation, every subscriber MUST implement the
// xptiTraceInit() and xptiTraceFinish() APIs for their subscriber collector to
// be loaded successfully.
XPTI_CALLBACK_API void xptiTraceInit(unsigned int major_version,
                                     unsigned int minor_version,
                                     const char *version_str,
                                     const char *stream_name) {
  // The example basic collector under xptifw/samples/basic_collector takes in
  // streams from anyone as an example. In this collector, we will accept
  // streams from just the SYCL Unified Runtime (UR) layer.
  printf("Stream Name: %s\n", stream_name);
  if (std::string("ur.call") == stream_name) {
    // Register this stream to get the stream ID; This stream may already have
    // been registered by the framework and will return the previously
    // registered stream ID. In this sample, we subscribe to only the events
    // generated by the SYCL UR layer. The protocol described by the UR layer
    // is that `user_data` field in the tpCallback() function will contain the
    // SYCL UR function name.
    GStreamID = xptiRegisterStream(stream_name);
    xptiRegisterCallback(GStreamID,
                         (uint16_t)xpti::trace_point_type_t::function_begin,
                         tpCallback);
    xptiRegisterCallback(GStreamID,
                         (uint16_t)xpti::trace_point_type_t::function_end,
                         tpCallback);

    printf("Registered all callbacks\n");
  } else {
    // handle the case when a stream name provided is not supported
    std::cerr << "Stream[" << stream_name
              << "] not supported by this collector - no callbacks registered "
                 "for the stream!\n";
  }
}

XPTI_CALLBACK_API void xptiTraceFinish(const char *stream_name) {
  // We do nothing here
}

XPTI_CALLBACK_API void tpCallback(uint16_t TraceType,
                                  xpti::trace_event_data_t *Parent,
                                  xpti::trace_event_data_t *Event,
                                  uint64_t Instance, const void *UserData) {
  // For `function_begin` trace point type, the Parent and Event parameters can
  // be null. However, the `UserData` field must be present and contain the
  // function name that these trace points are defined to trace.
  if (UserData) {
    const char *Name = (const char *)UserData;
    const char *be = "BEGIN";

    // Lock while we print information
    std::lock_guard<std::mutex> Lock(GIOMutex);
    // Print the record information
    if (TraceType & 0x1)
      be = "END";
    else
      be = "BEGIN";
    printf("SYCL_UR: %-35s %s\n", Name, be);
  }
}