/*========================== begin_copyright_notice ============================
Copyright (C) 2026 Intel Corporation

SPDX-License-Identifier: MIT
============================= end_copyright_notice ===========================*/

/*!
 * @file Skeleton tool demonstrating basic HLI instrumentation using HLIF interface
 */

#include <fstream>
#include <iomanip>
#include <iostream>
#include <map>
#include <sstream>
#include <string>
#include <utility>
#include <vector>

#include "gtpin_api.h"
#include "gtpin_tool_utils.h"
#include "hli_skeleton.h"

using namespace gtpin;
using namespace std;

Knob<bool> KNOB_NO_COUT("no_cout", false, "Do not send results to standard output device");

/* ============================================================================================= */
// Class HliSkeleton
/* ============================================================================================= */
class HliSkeleton : public GtTool
{
public:
    // Implementation of the IGtTool interface
    // This is the minimal required interface that all GTPin tools must implement
    const char* Name()          const                   override { return "HLI Skeleton"; }
    void        OnKernelBuild(IGtKernelInstrument&)     override; ///< Callback invoked when a kernel is built by the driver
    void        OnKernelRun(IGtKernelDispatch&)         override; ///< Callback invoked when a kernel is launched by the driver
    void        OnKernelComplete(IGtKernelDispatch&)    override; ///< Callback invoked when a kernel has completed execution

    void                LoadHliLibrary();                ///< Compile and load library of HLI functions
    static HliSkeleton* Instance();                      ///< Return single instance of this class
    static void         OnFini();                        ///< Termination handler registered with atexit()

private:
    HliSkeleton() : _exampleFunc("ExampleFunc"), _hliModule(nullptr) {} ///< Constructor
    ExampleHliFunc                               _exampleFunc;          ///< Analysis routine signature, defined in hli_skeleton.h
    IGtHliModuleHandle                           _hliModule;            ///< Handle of the compiled HLI module

    using KernelProfile = std::vector<uint64_t>;        ///< Vector of counters
    std::map<GtKernelId, KernelProfile> _profiles;      ///< Collection of kernel profiles indexed by kernel id
    std::string                         _profileStr;    ///< Accumulated profiling string
};

/* ============================================================================================= */
// HliSkeleton implementation
/* ============================================================================================= */
HliSkeleton* HliSkeleton::Instance()
{
    static HliSkeleton instance;
    return &instance;
}

// Called during tool initialization (from GTPin_Entry) to compile and load the HLI module
void HliSkeleton::LoadHliLibrary()
{
    // Get installation directory from command line
    std::string installDir = GetKnobValue<std::string>("installDir");
    // Compile and load the HLI module from file
    _hliModule = GTPin_GetCore()->HliLibrary().CompileModuleFromFile(JoinPath(installDir, "Examples", "hli_skeleton.cl").c_str());
    GTPIN_ASSERT(_hliModule != nullptr);
}

// Called after the application exits
// This is the appropriate place to dump accumulated profiling results
void HliSkeleton::OnFini()
{
    // Dump results of all kernel runs
    ofstream fs(JoinPath(GTPin_GetCore()->ProfileDir(), "hli_skeleton.txt"));
    GTPIN_ASSERT(fs.is_open());
    // Write accumulated profile string to file and optionally to standard output
    const std::string& profileStr = Instance()->_profileStr;
    fs << profileStr;
    fs.close();
    if (!KNOB_NO_COUT)
    {
        std::cout << profileStr;
    }
}

/* ============================================================================================= */
// Callback implementations
/* ============================================================================================= */

// Called when a kernel is being built by the driver - this is where we add instrumentation
void HliSkeleton::OnKernelBuild(IGtKernelInstrument& instrumentor)
{
    // Link the kernel with the library of HLI functions
    instrumentor.LinkHliModule(_hliModule);
    // Num of elements in shared vector
    // In this example, we create one counter per basic block
    // TO CUSTOMIZE: Change KernelProfile to match your profiling needs (e.g., more complex data structures)
    //               Also update the corresponding code in ExampleFunc (hli_skeleton.cl) and the function signature
    //               in hli_skeleton.h to match your data collection requirements
    const uint32_t vector_size = instrumentor.Cfg().NumBbls();
    // Create profile for this kernel
    auto result = _profiles.emplace(instrumentor.Kernel().Id(), vector_size); // Construct vector with vector_size elements
    // result is a pair of an iterator to the inserted element and a bool value which is set to true if and only if the insertion took place.
    // Assert that the kernel profile was newly inserted (not already present)    
    GTPIN_ASSERT(result.second);
    // Use a reference to the created KernelProfile
    KernelProfile& profile = result.first->second;
    // Share profile data between host and device memory
    instrumentor.MemoryMapper().Map(profile);

    // Insert instrumentation at each basic block
    for (auto bblPtr : instrumentor.Cfg().Bbls())
    {
        int bblId = bblPtr->Id(); // Basic block ID for pointer arithmetic
        if (!bblPtr->IsEmpty())
        {
            // Insert instrumentation APIs are defined at Src/api/gt_hli_function.h, e.g: InsertCallAtInstruction
            // In this example, we insert a call to ExampleFunc at the beginning of each basic block
            _exampleFunc.InsertCallAtBbl(instrumentor, *bblPtr, GtIpoint::Before(),
                                        NullReg(),                       // Unused return value
                                        profile.data() + bblId);         // arg[0]: Pointer to counter at KernelProfile[bblId]
        }
    }

    // Alternative example: insert instrumentation at each instruction (commented out)
    // Note: This example shows per-instruction instrumentation, but uses the same single-argument function.
    // If you need multiple arguments, modify ExampleFunc signature in hli_skeleton.h and hli_skeleton.cl
    /*for (auto bblPtr : instrumentor.Cfg().Bbls()) {
        for (auto insPtr : bblPtr->Instructions())
        {
            const IGtIns& ins = *insPtr;
            int bblId = bblPtr->Id();
            // Insert instrumentation APIs are defined at Src/api/gt_hli_function.h, e.g: InsertCallAtInstruction
            // In this example, we insert a call to ExampleFunc at each instruction
            _exampleFunc.InsertCallAtInstruction(instrumentor, ins, GtIpoint::Before(),
                                        NullReg(),                      // Unused return value
                                        profile.data() + bblId);        // arg[0]: Pointer to counter for this basic block
        }
    }*/
}

void HliSkeleton::OnKernelRun(IGtKernelDispatch& dispatcher)
{
    // Get the kernel profile for this specific kernel (guaranteed to exist from OnKernelBuild)
    KernelProfile& profile = _profiles.at(dispatcher.Kernel().Id());
    // Zero-initialize profile before this kernel dispatch
    std::fill(profile.begin(), profile.end(), 0);
    // Tell GTPin to run instrumented code, false will run original kernel code
    dispatcher.SetProfilingMode(true);
}

void HliSkeleton::OnKernelComplete(IGtKernelDispatch& dispatcher)
{
    // Translate profile data of this kernel dispatch into text
    // Note: Profile is guaranteed to exist (created in OnKernelBuild)
    KernelProfile& profile = _profiles.at(dispatcher.Kernel().Id());

    // Pretty print the results using string stream to accumulate output
    // Results are written to file at once in OnFini()
    // In this example, we print the number of times each basic block was executed.
    // NOTE: You can deduce instruction execution count by counting the number of instructions
    //       in each basic block during OnKernelBuild (using bblPtr->Instructions().size())
    //       and multiplying by the execution count shown here.
    std::ostringstream os;
    os << std::string(100, '-') << std::endl;
    os << "Kernel Id: " << dispatcher.DispatchId() << "  " << ExtendedKernelName(dispatcher.Kernel()) << std::endl;
    os << std::string(100, '-') << std::endl;
    os << setw(15) << "Record ID" << setw(15) << "Num" << std::endl << std::endl;
    for (uint32_t bblId = 0; bblId != profile.size(); ++bblId)
    {
        auto count = profile[bblId];
        if (count != 0)
        {
            os << setw(15) << bblId << setw(15) << count << std::endl;
        }
    }
    os << std::endl;

    // Add profile data of this kernel dispatch to _profileStr
    // Note: Assumes serial kernel completion or external synchronization
    _profileStr += os.str();
}

/* ============================================================================================= */
// GTPin_Entry - REQUIRED entry point for all GTPin tools
// This function is called by GTPin when the tool is loaded
/* ============================================================================================= */
EXPORT_C_FUNC void GTPin_Entry(int argc, const char *argv[])
{
    // Parse command line and configure GTPin
    ConfigureGTPin(argc, argv);

    // Register the tool (callbacks) with the GTPin core
    HliSkeleton::Instance()->Register();

    // Compile and load library of HLI functions
    HliSkeleton::Instance()->LoadHliLibrary();

    // Register the termination function
    atexit(HliSkeleton::OnFini);
}
