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

namespace fxn
{
    struct Vector2
    {
        float x;
        float y;
    };

    struct Vector3
    {
        float x;
        float y;
        float z;
    };

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

    struct NativeContext
    {
    private:
        struct scrVector3
        {
            alignas(8) float x;
            alignas(8) float y;
            alignas(8) float z;
        };

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

    private:
        void*            m_ReturnBuffer;
        std::uint32_t    m_ArgumentCount;
        void*            m_Arguments;
        std::uint32_t    m_DataCount;

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

        scrVector3      m_Vectors[4];

    private:
        enum
        {
            MAX_ARGUMENTS = 32,
            ARGUMENT_SIZE = 8
        };

        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);
        }

    public:
        template<typename R>
        inline R GetResult()
        {
            return *reinterpret_cast<R*>(m_ReturnBuffer);
        }

        template<>
        inline Vector3 GetResult<Vector3>()
        {
            scrVector3* vec = reinterpret_cast<scrVector3*>(m_ReturnBuffer);
            return Vector3{ vec->x, vec->y, vec->z };
        }

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

        inline std::uint32_t GetArgumentCount() const
        {
            return m_ArgumentCount;
        }

        inline void* GetArguments() const
        {
            return m_Arguments;
        }

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

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

    public:
        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]);
        }
    };
}