/*========================== begin_copyright_notice ============================
Copyright (C) 2021-2025 Intel Corporation

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

/*!
 * @file Common GTPin tool utilities
 */

#ifndef GTPIN_TOOL_UTILS_H_
#define GTPIN_TOOL_UTILS_H_

#include <string.h>
#include <sstream>
#include <fstream>
#include <algorithm>

#include "gtpin_api.h"

#include "gtpin_crt_utils.h"
#include "kernel_exec_config.h"
#include "knob_parser.h"

using namespace gtpin;

/* ===================================================================================================== */
// Knobs defining ranges of BBLs and instructions to be instrumented
/* ===================================================================================================== */
extern Knob<int>  knobMinInstrumentBbl;
extern Knob<int>  knobMaxInstrumentBbl;
extern Knob<int>  knobMinInstrumentIns;
extern Knob<int>  knobMaxInstrumentIns;

/* ===================================================================================================== */
// Class GtTool
/* ===================================================================================================== */
/*!
 * Basic implementation of the IGtTool interface
 */
class GtTool : public IGtTool
{
public:
    const char* Name()          const                   override { return "GTPin tool"; }
    uint32_t    ApiVersion()    const                   override { return GTPIN_API_VERSION; }
    void        OnKernelBuild(IGtKernelInstrument&)     override {}
    void        OnKernelRun(IGtKernelDispatch&)         override {}
    void        OnKernelComplete(IGtKernelDispatch&)    override {}

    /// Register this tool with the GTPin core
    bool Register()
    {
        IGtCore* gtpinCore = GTPin_GetCore();
        if (!gtpinCore->RegisterTool(*this))
        {
            GTPIN_ERROR_MSG(gtpinCore->LastError().ToString());
            return false;
        }
        return true;
    }

    /*!
     * @brief Handle common instrumentation knobs for kernel instrumentation
     *
     *        This function processes common instrumentation configuration knobs that affect
     *        how the kernel is instrumented. Currently handles the max_orig_ins knob which
     *        limits the number of original kernel instructions that are executed by inserting
     *        an early termination (EOT) instruction.
     *
     *        When max_orig_ins is set to a value less than MAX_INST_DEAFULT, the function
     *        will locate the first instruction with an ID greater than the specified limit
     *        and instrument it with an EOT instruction to terminate execution early.
     *
     * @param[in] instrumentor    Object that implements kernel instrumentation interface
     * @return true - success, false - failure
     *
     * @note The function assumes that the kernel has at least one exit basic block with
     *       a valid EOT instruction that can be duplicated for early termination.
     */
    bool HandleCommonInstrumnetationKnobs(IGtKernelInstrument& instrumentor);

    /*!
     * @brief Insert the specified instrumentation procedure before/after the specified instruction
     *
     *        This function can be used to filter out instruction-level instrumentation sites based
     *        on configured ranges of BBLs [knobMinInstrumentBbl, knobMaxInstrumentBbl] and instructions
     *        [knobMinInstrumentIns, knobMaxInstrumentIns]
     *
     * @param[in] instrumentor    Object that implements kernel instrumentation interface
     * @param[in] ins             Instruction in the original code of the kernel
     * @param[in] ipoint          Location in the original code, relative to 'ins'
     * @param[in] proc            Instrumentation procedure to be inserted at the specified location
     * @return true - success, false - failure
     */
    bool InstrumentInstruction(IGtKernelInstrument& instrumentor, const IGtIns& ins, GtIpoint ipoint, const IGtGenProcedure& proc);

    /*!
     * @brief Insert the specified instrumentation procedure before/after the specified BBL
     *
     *        This function can be used to filter out BBL-level instrumentation sites based
     *        on configured ranges of BBLs [knobMinInstrumentBbl, knobMaxInstrumentBbl] and
     *        instructions [knobMinInstrumentIns, knobMaxInstrumentIns]
     *
     * @param[in] instrumentor    Object that implements kernel instrumentation interface
     * @param[in] bbl             BBL in the original code of the kernel
     * @param[in] ipoint          Location in the original code, relative to 'bbl'
     * @param[in] proc            Instrumentation procedure to be inserted at the specified location
     *
     * @note If 'ipoint' specifies location after BBL, the insertion position is as following:
     *       - before the last BBL's instruction, if BBL ends with a control flow or EOT instruction
     *       - after the last BBL's instruction, if BBL falls through to the next block
     * @return true - success, false - failure
     */
    bool InstrumentBbl(IGtKernelInstrument& instrumentor, const IGtBbl& bbl, GtIpoint ipoint, const IGtGenProcedure& proc);

    /*!
     * @brief Insert the specified instrumentation procedure to the dstBbl of the edge and execute only when passing through the specified edge
     *
     * @param[in] instrumentor    Object that implements kernel instrumentation interface
     * @param[in] edge            Edge between two BBL's in the original code of the kernel
     * @param[in] proc            Instrumentation procedure to be inserted at the specified location
     *
     * @return true - success, false - failure
     */
    bool InstrumentEdge(IGtKernelInstrument& instrumentor, const IGtEdge& edge, const IGtGenProcedure& proc);
};

/* ===================================================================================================== */
// Struct ExecDescAlignedFormat
/* ===================================================================================================== */
/*!
 *  Formatter that converts GtKernelExecDesc object into aligned string
 */
struct ExecDescAlignedFormat
{
    /// Implementation of the GtKernelExecDescFormat interface
    std::string Device(uint32_t device) const
    {
        std::ostringstream ostr; ostr << device; return ostr.str();
    }
    std::string Enqueue(uint32_t enqueue) const
    {
        std::ostringstream ostr; ostr << std::setw(10) << enqueue; return ostr.str();
    }
    std::string Draw(uint32_t draw) const
    {
        std::ostringstream ostr; ostr << std::setw(10) << draw; return ostr.str();
    }
    std::string CommandList(uint32_t cl) const
    {
        std::ostringstream ostr; ostr << std::hex << std::showbase << std::setw(15) << cl; return ostr.str();
    }
    std::string Bus(uint32_t bus) const
    {
        std::ostringstream ostr; ostr << std::setw(10) << bus; return ostr.str();
    }
    std::string Function(uint32_t function) const
    {
        std::ostringstream ostr; ostr << std::setw(10) << function; return ostr.str();
    }
    std::string Pso(uint32_t pso) const
    {
        std::ostringstream ostr; ostr << std::hex << std::showbase << std::setw(15) << pso; return ostr.str();
    }
    std::string Frame(uint32_t frame) const
    {
        std::ostringstream ostr; ostr << std::setw(10) << frame; return ostr.str();
    }
    std::string Runtime(const GtKernelDispatchOclRuntime&) const
    {
        return "";
    }
    std::string Runtime(const GtKernelDispatchL0Runtime&) const
    {
        return "";
    }
    std::string Separator()  const { return ""; }
};

/*!
 *  Formatter that converts GtKernelExecDesc object into file name
 */
struct ExecDescFileNameFormat
{
    /// Implementation of the GtKernelExecDescFormat interface
    std::string Device(uint32_t device) const
    {
        std::ostringstream ostr; ostr << "d_" << device; return ostr.str();
    }
    std::string Enqueue(uint32_t enqueue) const
    {
        std::ostringstream ostr; ostr << "e_" << enqueue; return ostr.str();
    }
    std::string Draw(uint32_t draw) const
    {
        std::ostringstream ostr; ostr << "dw_" << draw; return ostr.str();
    }
    std::string CommandList(uint32_t cl) const
    {
        std::ostringstream ostr; ostr << "cl_" << cl; return ostr.str();
    }
    std::string Bus(uint32_t) const
    {
        return "";
    }
    std::string Function(uint32_t) const
    {
        return "";
    }
    std::string Pso(uint32_t pso) const
    {
        std::ostringstream ostr; ostr << "pso_" << pso; return ostr.str();
    }
    std::string Frame(uint32_t frame) const
    {
        std::ostringstream ostr; ostr << "f_" << frame; return ostr.str();
    }
    std::string Runtime(const GtKernelDispatchOclRuntime&) const
    {
        return "";
    }
    std::string Runtime(const GtKernelDispatchL0Runtime&) const
    {
        return "";
    }
    std::string Separator()  const { return "_"; }
};

/*!
 * @return String that represents name and basic descriptive properties of the specified kernel.
 *         The string includes kernel's hash value, so it can be used as a unique identifier of the kernel.
 */
inline std::string ExtendedKernelName(const IGtKernel& kernel)
{
    std::ostringstream extName;
    extName << GlueString(kernel.Name()) << "___" <<  kernel.UniqueName();
    return extName.str();
}

/*!
 * @return Assembly text of the code specified by its CFG
 */
inline std::string CfgAsmText(const IGtCfg& cfg)
{
    std::ostringstream asmText;
    for (auto bblPtr : cfg.Bbls())
    {
        asmText << "// BBL" << bblPtr->Id() << std::endl;
        for (auto insPtr : bblPtr->Instructions())
        {
            asmText << "[" << std::setw(3) << insPtr->Id() << "] " << insPtr->ToString() << std::endl;
        }
        asmText << std::endl;
    }
    return asmText.str();
}

/*!
 * Write kernel's assembly text 'asmText' into <kernelName>.asm file in the ISA directory 
 * @return success/failure status
 */
inline bool DumpKernelAsmText(const std::string& kernelName, const std::string& fileName, const std::string& asmText)
{
    std::string dir = JoinPath(std::string(GTPin_GetCore()->ProfileDir()), "ISA");
    if (!MakeDirectory(dir))
    {
        GTPIN_WARNING("Could not create directory " + dir);
        return false;
    }

    std::string filePath = JoinPath(dir, NormalizeFilename(fileName, ".asm"));
    std::ofstream fs(filePath);
    if (!fs)
    {
        GTPIN_WARNING("Could not create file " + filePath);
        return false;
    }

    fs << "// kernel name: " << kernelName << std::endl << std::endl << asmText;
    return true;
}

/* ===================================================================================================== */
// Class LockableFile
/* ===================================================================================================== */
/*!
 * Implementation for a safe multi-process file object for reading and writing
 */
class LockableFile
{
public:

    LockableFile(const std::string& filePath) : filePath(filePath){};

    /*!
     * @brief Write the specified string to the file associated with this object.
     *        Lock the file during the write operation, so that no other processes can
     *        read or write to this file concurrently.
     *
     * @param[in] content         Content to write to the file
     * @param[in] append          Indicates whether to open the file on append mode (true), or delete previous content before writing (false)
     * @return true - success, false - failure
     */
    bool Write(const std::string& content, bool append = false);

    /*!
     * @brief Read the amount of bytes requested from the file associated with this object at the offset given.
     *        Lock the file during the read operation, so that no other processes can
     *        read or write to this file concurrently..
     *
     * @param[in] bytes         Amount of bytes to read
     * @param[in] offset        Offset from the beginning of the file to start reading from
     * @return string with the content that has read from the offset given on success, or an
     *         empty string on failure
     */
    std::string Read(int bytes, int offset = 0) const;

    /*!
     * @return String of the file path for this object
     */
    std::string FilePath() const;

private:
    std::string filePath;
};

/*!
 * Create a new file or removing the content of an existing file in the path given
 * @return success/failure status
 */
bool CreateCleanFile(const std::string& filePath);

#endif