#ifndef THIRD_PARTY_XPROF_CONVERT_TRACE_OPTIONS_H_
#define THIRD_PARTY_XPROF_CONVERT_TRACE_OPTIONS_H_

#include <cstdint>
#include <memory>
#include <optional>

#include "absl/container/flat_hash_set.h"
#include "xprof/convert/tool_options.h"
#include "xprof/convert/trace_viewer/trace_events_filter_interface.h"
#include "xprof/convert/trace_viewer/trace_events_to_json.h"
#include "xprof/convert/trace_viewer/trace_utils.h"
#include "plugin/xprof/protobuf/trace_filter_config.pb.h"

namespace tensorflow {
namespace profiler {

enum class TraceDeviceType {
  kUnknownDevice,
  kTpu,
  kGpu,
};

inline constexpr const char* kFullDma = "full_dma";
inline constexpr const char* kOverlay = "overlay";
inline constexpr const char* kDerivedTimeline = "derived_timeline";
inline constexpr const char* kRawGpuStreams = "raw_gpu_streams";
inline constexpr const char* kMaxHostThreads = "max_threads";
inline constexpr const char* kShowHloCostModel = "show_hlo_cost_model";
inline constexpr const char* kMpmdPipelineView = "mpmd_pipeline_view";
inline constexpr const char* kEnableLegacyDcn = "enable_legacy_dcn";

// Options used to select TraceEvents (e.g., for visualization or further
// processing).
struct TraceOptions {
  // TraceFilterConfig is a filtering config written by the user for filtering
  // trace events.
  std::optional<tensorflow::profiler::TraceFilterConfig> trace_filter_config;
  // Visualize tensor core overlays. This is low-level debug information for
  // XLA and xprof devs. so don't show by default.
  bool overlay = false;

  // Visualize intermediate trace entries (Nf, NfIci and IciPacket entries)
  // for a DMA. By default only DMA endpoints are visualized.
  bool full_dma = false;

  // For gpu, whether to show derived timelines such as tf op, hlo module etc.
  bool derived_timeline = true;

  // For gpu, whether to show raw gpu streams (generated by CUPTI).
  bool raw_gpu_streams = true;

  // If selected_group_ids is not empty, only the trace events with group_id
  // listed inside selected_group_ids are visible. For instance, if
  // selected_group_ids is {78 91}, only trace events with group_id equal to 78
  // or 91 are visible.
  absl::flat_hash_set<int64_t> selected_group_ids;

  // Maximum number of host threads to display.
  int max_host_threads = 1000;

  bool perfetto_trace = false;

  // Whether to show Hlo Cost Model FLOPS and Bytes in trace viewer.
  bool show_hlo_cost_model = false;

  // Whether to show MPMD view in trace viewer.
  bool mpmd_pipeline_view = false;

  // Whether to enable legacy DCN view.
  // Note: This is deprecated and will be removed in Q2 2026.
  bool enable_legacy_dcn = false;
};

// Returns TraceOptions for use when the generated events will be visualized.
// Parses the tool options from a request.
TraceOptions TraceOptionsFromToolOptions(const ToolOptions& options);

// Converts TraceOptions to JsonTraceOptions::Details format.
JsonTraceOptions::Details TraceOptionsToDetails(TraceDeviceType device_type,
                                                const TraceOptions& options);

// Returns a trace events filter configured to filter events according to the
// given options.
std::unique_ptr<tensorflow::profiler::TraceEventsFilterInterface>
CreateTraceEventsFilterFromTraceOptions(const TraceOptions& options);

inline bool IsTpuTrace(const Trace& trace) {
  for (const auto& [device_id, device] : trace.devices()) {
    if (IsTpuCoreDeviceName(device.name())) {
      return true;
    }
  }
  return false;
}

// Returns true if event is an intermediate event in a flow.
inline bool IsFlowMid(const TraceEvent& event) {
  return event.has_flow_id() && event.flow_entry_type() == TraceEvent::FLOW_MID;
}

}  // namespace profiler
}  // namespace tensorflow

#endif  // THIRD_PARTY_XPROF_CONVERT_TRACE_OPTIONS_H_
