// Copyright Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT

#pragma once

#include <rocRoller/KernelArguments.hpp>

#include <rocRoller/DataTypes/DataTypes_Utils.hpp>
#include <rocRoller/Operations/CommandArgument_fwd.hpp>

#include <concepts>
#include <cstring>
#include <sstream>
#include <string>
#include <unordered_map>
#include <vector>

#include <rocRoller/Utilities/Logging.hpp>

namespace rocRoller
{

    struct AppendKernelArgumentsVisitor
    {
        KernelArguments*   args;
        std::string const& argName;

        template <CCommandArgumentValue T>
        void operator()(T const& value)
        {
            args->append(argName, value);
        }

        void call(CommandArgumentValue const& value)
        {
            std::visit(*this, value);
        }
    };

    template <typename T>
    inline void KernelArguments::append(std::string const& argName, T value)
    {
        if constexpr(std::is_same_v<T, CommandArgumentValue>)
        {
            AppendKernelArgumentsVisitor visitor{this, argName};
            visitor.call(value);
        }
        else
        {
            append(argName, value, true);
        }
    }

    template <typename T>
    inline void KernelArguments::appendUnbound(std::string const& argName)
    {
        append(argName, static_cast<T>(0), false);
    }

    template <typename T>
    inline void KernelArguments::bind(std::string const& argName, T value)
    {
        if(!m_log)
        {
            throw std::runtime_error("Binding is not supported without logging.");
        }

        auto it = m_argRecords.find(argName);
        if(it == m_argRecords.end())
        {
            throw std::runtime_error("Attempt to bind unknown argument " + argName);
        }

        auto& record = it->second;

        if(std::get<ArgBound>(record))
        {
            throw std::runtime_error("Attempt to bind already bound argument " + argName);
        }

        if(sizeof(T) != std::get<ArgSize>(record))
        {
            throw std::runtime_error("Size mismatch in binding argument " + argName);
        }

        size_t offset = std::get<ArgOffset>(record);

        if(offset % alignof(T) != 0)
        {
            throw std::runtime_error("Alignment error in argument " + argName + ": type mismatch?");
        }

        writeValue(offset, value);

        std::get<ArgString>(record) = stringForValue(value, true);
        std::get<ArgBound>(record)  = true;
    }

    template <typename T>
    inline std::string KernelArguments::stringForValue(T value, bool bound) const
    {
        if(!m_log)
            return "";

        if(!bound)
            return "<unbound>";

        std::ostringstream msg;

        // Some pointer types (char, int8_t, uint8_t, etc) are interpreted
        // as a string, we just want the pointer value.
        if constexpr(std::is_pointer_v<T>)
        {
            msg << static_cast<void const*>(value);
        }
        else
        {
            msg << value;
        }
        return msg.str();
    }

    template <CScaleType T>
    inline void KernelArguments::append(std::string const& argName, T value, bool bound)
    {
        auto constexpr PACKED_SCALE_ALIGN = alignof(typename PackedTypeOf<T>::type);
        auto constexpr PACKED_SCALE_SIZE  = sizeof(typename PackedTypeOf<T>::type);

        alignTo(PACKED_SCALE_ALIGN);
        size_t offset  = m_data.size();
        size_t argSize = PACKED_SCALE_SIZE;

        if(m_log)
        {
            std::string valueString = stringForValue(value, bound);
            appendRecord(argName, Arg(offset, argSize, bound, valueString));
        }

        m_data.insert(m_data.end(), argSize, 0);
        writeValue(offset, value);
    }

    template <CScaleType T>
    inline void KernelArguments::writeValue(size_t offset, T value)
    {
        auto constexpr PACKED_SCALE_SIZE = sizeof(typename PackedTypeOf<T>::type);

        size_t argSize = PACKED_SCALE_SIZE;
        AssertFatal(offset + argSize <= m_data.size(),
                    "Value exceeds allocated bounds.",
                    ShowValue(offset),
                    ShowValue(argSize),
                    ShowValue(m_data.size()));

        std::memset(&m_data[offset], value.scale, argSize);
    }

    template <typename T>
    inline void KernelArguments::append(std::string const& argName, T value, bool bound)
    {
        alignTo(alignof(T));

        size_t offset  = m_data.size();
        size_t argSize = sizeof(T);

        if(m_log)
        {
            std::string valueString = stringForValue(value, bound);
            appendRecord(argName, Arg(offset, argSize, bound, valueString));
        }

        m_data.insert(m_data.end(), argSize, 0);
        writeValue(offset, value);
    }

    template <typename T>
    inline void KernelArguments::writeValue(size_t offset, T value)
    {
        AssertFatal(offset + sizeof(T) <= m_data.size(),
                    "Value exceeds allocated bounds.",
                    ShowValue(offset),
                    ShowValue(sizeof(T)),
                    ShowValue(m_data.size()));

        std::memcpy(&m_data[offset], &value, sizeof(T));
    }

    inline void KernelArguments::alignTo(size_t alignment)
    {
        size_t extraElements = m_data.size() % alignment;
        size_t padding       = (alignment - extraElements) % alignment;

        m_data.insert(m_data.end(), padding, 0);
    }

    inline void KernelArguments::appendRecord(std::string const&   argName,
                                              KernelArguments::Arg record)
    {
        auto it = m_argRecords.find(argName);
        if(it != m_argRecords.end())
        {
            throw std::runtime_error("Duplicate argument name: " + argName);
        }

        m_argRecords[argName] = record;
        m_names.push_back(argName);
    }

    template <typename T>
    KernelArguments::const_iterator::operator T() const
    {
        if(sizeof(T) != m_value.second)
        {
            throw std::bad_cast();
        }
        return *reinterpret_cast<T*>(const_cast<void*>(m_value.first));
    }

    inline bool KernelArguments::log() const
    {
        return m_log;
    }

    inline std::vector<uint8_t> const& KernelArguments::dataVector() const
    {
        return m_data;
    }

    inline RuntimeArguments KernelArguments::runtimeArguments() const
    {
        return RuntimeArguments(m_data.data(), m_data.size());
    }
}