/*
 * Copyright (C) 2024-2025 Intel Corporation
 *
 * SPDX-License-Identifier: MIT
 *
 */

// IWYU pragma: no_include "perfetto.h"

#include "trace_perfetto.hpp"

#include <fstream>
#include <stdlib.h>
#include <string>
#include <vector>

PERFETTO_TRACK_EVENT_STATIC_STORAGE();

[[maybe_unused]] static TracePerfetto trace_perfetto;

TracePerfetto::TracePerfetto() {
    const char *env = getenv("ZE_INTEL_NPU_LOGPERFETTO");
    enable = env == nullptr || env[0] == '0' || env[0] == '\0' ? false : true;

    if (enable) {
        bool isSystemBackend = std::string(env) == "system";

        perfetto::TracingInitArgs args;
        if (!isSystemBackend) {
            args.backends |= perfetto::kInProcessBackend;
        } else {
            args.backends |= perfetto::kSystemBackend;
        }
        perfetto::Tracing::Initialize(args);
        perfetto::TrackEvent::Register();

        perfetto::protos::gen::TrackEventConfig track_event_cfg;
        track_event_cfg.add_disabled_categories("*");
        track_event_cfg.add_enabled_categories("API");
        track_event_cfg.add_enabled_categories("SYS");
        track_event_cfg.add_enabled_categories("NPU_ELF");
        track_event_cfg.add_enabled_categories("NPU_COMPILER");

        if (!isSystemBackend) {
            perfetto::TraceConfig cfg;
            cfg.add_buffers()->set_size_kb(64 * 1024);
            auto *ds_cfg = cfg.add_data_sources()->mutable_config();
            ds_cfg->set_name("track_event");
            ds_cfg->set_track_event_config_raw(track_event_cfg.SerializeAsString());

            record = perfetto::Tracing::NewTrace();
            record->Setup(cfg);
            record->StartBlocking();
        }
    }
}

TracePerfetto::~TracePerfetto() {
    if (!record)
        return;

    record->StopBlocking();
    std::vector<char> trace_data(record->ReadTraceBlocking());

    std::ofstream out;
    out.open("intel-npu-umd.perfetto-trace", std::ios::out | std::ios::binary);
    out.write(&trace_data[0], (long)trace_data.size());
    out.close();
}
