// MIT License
//
// Copyright (c) 2017-2025 Advanced Micro Devices, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

// Header file for CmdBuilder and CmdBuffer interfaces

#ifndef SRC_PM4_CMD_BUILDER_H_
#define SRC_PM4_CMD_BUILDER_H_

#include <stdint.h>
#include <string.h>

#include <iomanip>
#include <iostream>
#include <sstream>
#include <type_traits>
#include <vector>
#include <atomic>
#include "lib/aqlprofile/util/reg_offsets.h"

#define APPEND_COMMAND_WRAPPER(cmdbuf, ...)                                                        \
    do                                                                                             \
    {                                                                                              \
        PrintCommand(__FUNCTION__, __VA_ARGS__);                                                   \
        cmdbuf->Append(__VA_ARGS__);                                                               \
    } while(false);

namespace pm4_builder
{
namespace
{
#if defined(DEBUG_TRACE)
template <typename T, typename std::enable_if<sizeof(T) % sizeof(uint32_t) == 0, int>::type = 0>
constexpr size_t
PacketsSize(T&& packet)
{
    return sizeof(packet) / sizeof(uint32_t);
}
template <typename T, typename... Ts>
constexpr size_t
PacketsSize(T&& packet, Ts&&... rest)
{
    return PacketsSize(std::forward<T>(packet)) + PacketsSize(std::forward<Ts>(rest)...);
}

template <typename T, typename std::enable_if<sizeof(T) % sizeof(uint32_t) == 0, int>::type = 0>
void
PrintPackets(std::ostream& stream, T&& packet)
{
    for(size_t i = 0; i < sizeof(packet) / sizeof(uint32_t); ++i)
        stream << " " << std::hex << std::right << std::setw(8) << std::setfill('0')
               << reinterpret_cast<const uint32_t*>(&packet)[i];
}
template <typename T, typename... Ts>
void
PrintPackets(std::ostream& stream, T&& packet, Ts&&... rest)
{
    PrintPackets(stream, std::forward<T>(packet));
    PrintPackets(stream, std::forward<Ts>(rest)...);
}
#endif

template <typename... Ts>
void
PrintCommand(const char* function_name, Ts&&... packets)
{
#if defined(DEBUG_TRACE)
    std::ostringstream oss;
    oss << "'" << function_name << "' size(" << std::dec
        << PacketsSize(std::forward<Ts>(packets)...) << ")";
    std::clog << std::setw(40) << std::left << oss.str() << ":";
    PrintPackets(std::clog, std::forward<Ts>(packets)...);
    std::clog << std::setfill(' ') << std::endl;
#endif
}

}  // namespace

/// @brief Implements the interface CmdBuffer and thus can be used to
/// translate various Gpu commands as byte stream.
/// @note: The Api does not require implementations to be thread safe.
/// Users are therefore required to be access in a serialized manner.
class CmdBuffer
{
public:
    /// @brief Append the command into the underlying buffer
    /// @param packet Buffer containing one or more instances of Gpu commands
    template <typename T, typename std::enable_if<sizeof(T) % sizeof(uint32_t) == 0, int>::type = 0>
    void Append(T&& packet)
    {
        size_t pos = data_.size();
        data_.resize(pos + sizeof(packet) / sizeof(uint32_t));
        memcpy(&data_[pos], &packet, sizeof(T));
    }
    template <typename... Ts>
    void Append(Ts&&... packets)
    {
        using expander = int[];
        (void) expander{0, (Append(std::forward<Ts>(packets)), 0)...};
    }

    /// @brief Append num_dwords of data to the buffer
    /// @param data_pointer points to the data to append
    /// @param num_dwords number of dwords
    void Append(uint32_t* data_pointer, uint32_t num_dwords)
    {
        size_t pos = data_.size();
        data_.resize(pos + num_dwords);
        memcpy(&data_[pos], data_pointer, num_dwords);
    }

    /// @brief Return size of Gpu commands in bytes in the underlying buffer
    size_t Size() const { return data_.size() * sizeof(data_[0]); }

    /// @brief Return size of Gpu commands in DWord in the underlying buffer
    size_t DwSize() const { return data_.size(); }

    void Assign(uint32_t index, uint32_t value) { data_[index] = value; }

    /// @brief Return address of the start of accumulated commands.
    const void* Data() const { return &data_[0]; }

    /// @brief Clear buffer.
    void Clear() { return data_.clear(); }

private:
    /// @brief Defines Gpu command buffer as a vector of uint32_t
    std::vector<uint32_t> data_;
};

enum ChipletId
{
    CHIPLET_XCD0 = 0,
    CHIPLET_XCD1 = 1,
    CHIPLET_XCD2 = 2,
    CHIPLET_XCD3 = 3,
    CHIPLET_XCD4 = 4,
    CHIPLET_XCD5 = 5,
    CHIPLET_XCD6 = 6,
    CHIPLET_XCD7 = 7,
    CHIPLET_AID0 = 8,
    CHIPLET_AID1 = 9,
    CHIPLET_AID2 = 10,
    CHIPLET_AID3 = 11,
    CHIPLET_MID0 = 12,
    CHIPLET_MID1 = 13,
};

/// @brief Specifies the public interface of CmdBuilder for use by
/// clients to build Gpu command streams.
class CmdBuilder
{
public:
    CmdBuilder(const reg_base_offset_table* _table)
    : ip_offset_table(_table)
    {}

    /// @brief Build and copy a WaitIdle Gpu command into command buffer
    /// @param cmdbuf Pointer to command buffer to be appended
    virtual void BuildWriteWaitIdlePacket(CmdBuffer* cmdbuf) = 0;

    /// @brief Builds a Gpu command to wait until condition is realized
    /// @param cmdbuf command buffer to be appended with launch command
    /// @param mem_space if the address is in memory or is a register offset
    /// @param wait_addr address to wait on
    /// @param func_eq true means equal, false means not-equal
    /// @param mask_val Mask to apply on value from addr in comparison
    /// @param wait_val value to apply for the func given above
    virtual void BuildWaitRegMemCommand(CmdBuffer* cmdbuf,
                                        bool       mem_space,
                                        uint64_t   wait_addr,
                                        bool       func_eq,
                                        uint32_t   mask_val,
                                        uint32_t   wait_val) = 0;

    /// @brief Build CP command to program a Gpu register
    /// @param cmdbuf Pointer to command buffer to be appended
    /// @param addr Register to be programmed
    /// @param value Value to write into register
    virtual void BuildWriteUConfigRegPacket(CmdBuffer* cmdbuf, uint32_t addr, uint32_t value) = 0;

    /// @brief Build and copy WriteShReg command
    /// @param cmdbuf Pointer to command buffer to be appended
    /// @param addr Offset of the register
    /// @param value Value to write into register
    virtual void BuildWriteShRegPacket(CmdBuffer* cmdbuf, uint32_t addr, uint32_t value) = 0;

    /// @brief Build a Gpu command that copies data from a specified
    /// source register to destination
    /// @param cmdbuf Pointer to command buffer to be appended
    /// @param src_sel Source register selector
    /// @param src_reg_add 32-bit Source register address of the data to read from
    /// @param dst_addr Destination address for the data to be written to
    /// @param size Size of the data to be written
    /// @param  wait True if Gpu command should confirm the write operation
    /// operation has completed successfully
    virtual void BuildCopyRegDataPacket(CmdBuffer*  cmdbuf,
                                        uint32_t    src_reg_addr,
                                        const void* dst_addr,
                                        uint32_t    size,
                                        bool        wait) = 0;

    /// @brief Same as the above function but with 64-bit src_reg_addr.
    virtual void BuildCopyRegDataPacket(CmdBuffer*  cmdbuf,
                                        uint64_t    src_reg_addr,
                                        const void* dst_addr,
                                        uint32_t    size,
                                        bool        wait)
    {
        BuildCopyRegDataPacket(cmdbuf, uint32_t(src_reg_addr), dst_addr, size, wait);
    }

    /// @brief Builds the Gpu command to reference indirectly a stream
    /// of other Gpu commands. The launch command is then copied into
    /// the command buffer parameter.
    /// @param cmdBuf command buffer to be appended with launch command
    /// @param cmd_addr Address of command buffer carrying command stream
    /// @param cmd_size Size of dispatch command stream in bytes
    virtual void BuildIndirectBufferCmd(CmdBuffer*  cmdbuf,
                                        const void* cmd_addr,
                                        std::size_t cmd_size) = 0;

    virtual void BuildMutexAcquirePacket(CmdBuffer* cmdbuf, size_t addr) = 0;
    virtual void BuildMutexReleasePacket(CmdBuffer* cmdbuf, size_t addr) = 0;

    virtual uint32_t MakeMutexSlot()
    {
        static std::atomic<size_t> slot{1};
        return uint32_t(slot.fetch_add(1)) | (1u << 31);
    }

    /// @brief Builds a NOP packet
    /// @param cmdBuf command buffer to be appended with launch command
    /// @param num_dwords number of dwords of the whole NOP packet, including the Header
    virtual void BuildNopPacket(CmdBuffer* cmdbuf, uint32_t num_dwords) = 0;

    /// @brief Builds a TT Finish Event
    /// @param cmdBuf  command buffer to be appended with launch command
    virtual void BuildThreadTraceEventFinish(CmdBuffer* cmdBuf) = 0;

    /// @brief Prime pages of ATCL2
    /// @param cmdBuf  command buffer to be appended with launch command
    virtual void BuildPrimeL2(CmdBuffer* cmdBuf, uint64_t addr) = 0;

    /// @brief Generates RT packets into thread trace buffer (gfx9 only)
    /// @param cmdBuf  command buffer to be appended with launch command
    /// @param dst  where gpu clock data is r/w. Must persist during packet dispatch
    /// @param reg  userdata register address
    /// @param header  SQTT packet header
    virtual void BuildGPUClockPacket(CmdBuffer*      cmdBuf,
                                     uint64_t*       dst,
                                     const Register& reg,
                                     uint32_t        header){};

    /// @brief Release resources used by CmdBuilder
    virtual ~CmdBuilder(){};

    bool bUsePerfCounterMode{true};
    // is this an MI300 command builder?

    uint32_t get_addr(const Register& reg)
    {
        if(!ip_offset_table)
        {
            std::cerr << "CmdBuilder::get_addr(): IP offset table is NULL!" << std::endl;
            std::abort();
        }
        return (*ip_offset_table)[reg];
    }

    virtual void BuildWriteConfigRegPacket(CmdBuffer* cmdbuf, uint32_t addr, uint32_t value) = 0;

    virtual void BuildWritePConfigRegPacketToChiplet(CmdBuffer* cmdbuf,
                                                     uint32_t   addr,
                                                     uint32_t   value,
                                                     ChipletId  chiplet,
                                                     bool       write_to_aid = true)
    {}
    virtual void BuildWritePConfigRegPacketToChiplet(CmdBuffer*      cmdbuf,
                                                     const Register& reg,
                                                     uint32_t        value,
                                                     ChipletId       chiplet,
                                                     bool            write_to_aid = true)
    {}

    virtual uint32_t BuildCopyCounterDataPacketFromChiplet(CmdBuffer*      cmdbuf,
                                                           const Register& reg_lo,
                                                           const Register& reg_hi,
                                                           const void*     dst_addr,
                                                           uint32_t        dw_mask,
                                                           ChipletId       chiplet,
                                                           bool            copy_from_aid = true)
    {
        return 0;
    }

private:
    const reg_base_offset_table* const ip_offset_table;
};

/// @brief Returns the lower 32-bits of a value
constexpr uint32_t
Low32(uint64_t u)
{
    return static_cast<uint32_t>(u);
}

/// @brief Returns the upper 32-bits of a value
constexpr uint32_t
High32(uint64_t u)
{
    return static_cast<uint32_t>(u >> 32);
}

/// @brief Returns the lower 32-bits of an address
inline uint32_t
PtrLow32(const void* p)
{
    return Low32(reinterpret_cast<uintptr_t>(p));
}

/// @brief Returns the upper 32-bits of an address
inline uint32_t
PtrHigh32(const void* p)
{
    return High32(reinterpret_cast<uintptr_t>(p));
}

}  // namespace pm4_builder

#endif  // SRC_PM4_CMD_BUILDER_H_