#pragma once
#include "prelude/string.hpp"
#include "ir/slot.hpp"

#define AVAILABLE_REGISTERS 4
#define AVAILABLE_STACK_SIZE 4096

template<typename Slot, typename SlotAddr = const StringView&>
class Allocator
{
public:
    virtual ~Allocator() {};
public:
    virtual const Slot& Allocate(SlotAddr addr) = 0;
    virtual const Slot& Resolve(SlotAddr addr) = 0;
};

class SlotAllocator : public Allocator<IR::IRSlot>
{
public:
    SlotAllocator() = default;
    ~SlotAllocator() = default;
public:
    const IR::IRSlot& Allocate(const StringView& addr) override
    {
        if (m_regs < AVAILABLE_REGISTERS)
        {
            m_slots.Push(IR::IRSlot(addr, m_regs++, IR::IRSlot::Type::REGISTRY));
            return m_slots.data[m_slots.size - 1];
        }
        if (m_offset_counter + 8 <= AVAILABLE_STACK_SIZE)
        {
            m_offset_counter += 8;
            m_slots.Push(IR::IRSlot(addr, m_offset_counter, IR::IRSlot::Type::STACK));
            return m_slots.data[m_slots.size - 1];
        }
        // TODO: proper error handling (stack overflow etc.)
        assert(0 && "failed to allocate local");
    }

    const IR::IRSlot& Resolve(const StringView& addr) override
    {
        for (size_t i = 0; i < m_slots.size; ++i)
        {
            if (m_slots.data[i].GetAddr() == addr)
            {
                return m_slots.data[i];
            }
        }

        assert(0 && "could not resolve stack offset for specified address");
    }
public:
    View<IR::IRSlot> slots() const { return m_slots.view(); }
public:
    unsigned int GetStackSize() const { return m_offset_counter; }
public:
    int m_regs = 0;
    unsigned int m_offset_counter = 0;
    Builder<IR::IRSlot> m_slots;
};