#pragma once
#include <cstdint>
#include <cstddef>
#include <stdexcept>
#include <atomic>

namespace fxn
{
    struct Vector4
    {
        float x;
        float y;
        float z;
        float w;
    };

    struct alignas(16) Vector3 // input
    {
        float x;
        float y;
        float z;
    };

    struct scrVector3 // output
    {
        alignas(8) float x;
        alignas(8) float y;
        alignas(8) float z;
    };

    struct Vector2
    {
        float x;
        float y;
    };

    struct alignas(16) scrVector3N
    {
        float x;
        float y;
        float z;
    };

    struct VectorSpace
    {
        scrVector3* outVectors[4];
        scrVector3N  inVectors[4];
    };

    struct NativeContext
    {
    private:
        void*            m_ReturnBuffer;
        std::uint32_t    m_ArgumentCount;
        void*            m_Arguments;
        std::uint32_t    m_DataCount;
 
        VectorSpace      m_VectorSpace;
        scrVector3       m_Vectors[4];

        enum
        {
            MAX_ARGUMENTS = 32,
            ARGUMENT_SIZE = sizeof(void*)
        };

        std::uint8_t     m_Stack[MAX_ARGUMENTS * ARGUMENT_SIZE];
        std::uint64_t    m_NativeHash;
        std::atomic_bool m_Executed;

    public:
        inline NativeContext()
            : m_ReturnBuffer(nullptr), m_ArgumentCount(0), m_Arguments(nullptr), m_DataCount(0), m_NativeHash(0), m_VectorSpace{}, m_Vectors{}, m_Stack{}, m_Executed(false)
        {
            for (std::size_t i = 0; i < sizeof(m_Stack); i++)
            {
                m_Stack[i] = 0;
            }

            m_Arguments = &m_Stack;
            m_ReturnBuffer = &m_Stack;

            m_ArgumentCount = 0;
            m_DataCount = 0;
        }

        inline NativeContext(std::uint64_t hash)
            : NativeContext()
        {
            m_NativeHash = hash;
        }

    public:
        template<typename T>
        inline void PushArgument(const T& arg)
        {
            if (m_ArgumentCount >= MAX_ARGUMENTS)
            {
                throw std::runtime_error("Exceeded maximum number of arguments.");
            }

            if constexpr (sizeof(T) < ARGUMENT_SIZE)
            {
                // Zero out the memory to avoid garbage values
                *reinterpret_cast<std::uintptr_t*>(m_Stack + (m_ArgumentCount * ARGUMENT_SIZE)) = 0;
            }

            *reinterpret_cast<T*>(m_Stack + (m_ArgumentCount * ARGUMENT_SIZE)) = arg;
            m_ArgumentCount++;
        }

        template<>
        inline void PushArgument<Vector3>(const Vector3& arg)
        {
            PushArgument(arg.x);
            PushArgument(arg.y);
            PushArgument(arg.z);
        }

        template<typename T>
        inline T GetResult()
        {
            return *reinterpret_cast<T*>(m_ReturnBuffer);
        }

        template<>
        inline Vector3 GetResult<Vector3>()
        {
            auto vec = *reinterpret_cast<scrVector3*>(m_ReturnBuffer);
            return Vector3{ vec.x, vec.y, vec.z }; // fix for alignment issues
        }

        inline void SetExecuted(bool executed)
        {
            m_Executed.store(executed, std::memory_order_release);
        }

        inline bool IsExecuted() const
        {
            return m_Executed.load(std::memory_order_acquire);
        }

        inline void SetVectorResults()
        {
            for (size_t i = 0; i < m_DataCount; i++)
            {
                auto outVector = m_VectorSpace.outVectors[i];
                const auto& inVector = m_VectorSpace.inVectors[i];

                outVector->x = inVector.x;
                outVector->y = inVector.y;
                outVector->z = inVector.z;
            }
        }

        template<typename T>
        inline const T& GetArgument(std::size_t index) const
        {
            if (index >= m_ArgumentCount)
            {
                throw std::out_of_range("Argument index out of range.");
            }

            auto functionData = reinterpret_cast<uintptr_t*>(m_Arguments);
            return *reinterpret_cast<T*>(&functionData[index]);
        }

        inline std::uint64_t GetNativeHash() const
        {
            return m_NativeHash;
        }
    };
}