// TODO: store all of values, allocated registers, stack offsets in single allocator
// to be able to find out which kind a specific temp register represent

#pragma once
#include <cstddef>
#include "string.hpp"
#include "ir.hpp"

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

struct StackSlot
{
    size_t offset;
    StringView addr;
};

class StackAllocator : public Allocator<StackSlot>
{
public:
    StackAllocator() = default;
    ~StackAllocator() override = default;
public:
    const StackSlot& Allocate(const StringView& addr)
    {
        m_offset_counter += 4;
        m_slots.Push(StackSlot { m_offset_counter, addr });
        return m_slots.data[m_slots.size - 1];
    }

    const StackSlot& Resolve(const StringView& addr)
    {
        for (size_t i = 0; i < m_slots.size; ++i)
        {
            if (strcmp(m_slots.data[i].addr.c_str(), addr.c_str()) == 0)
            {
                return m_slots.data[i];
            }
        }

        assert(0 && "could not resolve stack offset for specified address");
    }
private:
    size_t m_offset_counter = 0;
    Builder<StackSlot> m_slots;
};

// struct ConstSlot
// {
//     long value;
//     StringView addr;
// };

// class ConstAllocator : public Allocator<ConstSlot>
// {
// public:
//     ConstAllocator() = default;
//     ~ConstAllocator() override = default;
// public:
//     const ConstSlot& Allocate(const StringView& addr)
//     {
//         m_slots.Push(ConstSlot { 0, addr });
//         return m_slots.data[m_slots.size - 1];
//     }

//     const ConstSlot& StoreValue(const StringView& addr, long value)
//     {
//         for (size_t i = 0; i < m_slots.size; ++i)
//         {
//             if (strcmp(m_slots.data[i].addr.c_str(), addr.c_str()) == 0)
//             {
//                 m_slots.data[i].value = value;
//                 return m_slots.data[i];
//             }
//         }

//         assert(0 && "could not resolve const under specified address");
//     }

//     const ConstSlot& Resolve(const StringView& addr)
//     {
//         for (size_t i = 0; i < m_slots.size; ++i)
//         {
//             if (strcmp(m_slots.data[i].addr.c_str(), addr.c_str()) == 0)
//             {
//                 return m_slots.data[i];
//             }
//         }

//         assert(0 && "could not resolve const under specified address");
//     }
// private:
//     Builder<ConstSlot> m_slots;
// };

// struct RegisterSlot
// {
//     const StringView& reg;
//     StringView addr;
// };

// class RegisterAllocator : public Allocator<RegisterSlot>
// {
// public:
//     RegisterAllocator()
//     {
//         m_regs.Push(std::move(StringView("eax")));
//         m_regs.Push(std::move(StringView("ecx")));
//     }
//     ~RegisterAllocator() override = default;
// public:
//     const RegisterSlot& Allocate(const StringView& addr)
//     {
//         assert(m_slots.size < m_regs.size && "no space available for allocating to register");
//         m_slots.Push(RegisterSlot { m_regs.data[m_slots.size], addr });
//         return m_slots.data[m_slots.size - 1];
//     }

//     const RegisterSlot& Resolve(const StringView& addr)
//     {
//         for (size_t i = 0; i < m_slots.size; ++i)
//         {
//             if (strcmp(m_slots.data[i].addr.c_str(), addr.c_str()) == 0)
//             {
//                 return m_slots.data[i];
//             }
//         }

//         assert(0 && "could not resolve register for specified address");
//     }

//     void Clear()
//     {
//         m_slots.size = 0;
//     }
// private:
//     Builder<RegisterSlot> m_slots;
//     Builder<StringView> m_regs;
// };

class CodeGenerator
{
public:
    virtual ~CodeGenerator() {};

    virtual void Generate(const char* filename, View<IR::Op*> ops) = 0;
};

class StackFasmX86_64Generator : public CodeGenerator
{
public:
    ~StackFasmX86_64Generator() override = default;
private:
    int GetStackSize(const IR::OpView ops)
    {
        int stackSize = 0;
        for (auto &op : ops)
        {
            switch (op->GetType())
            {
            case IR::OpType::STORE:
            case IR::OpType::LOAD:
            case IR::OpType::LOAD_CONST:
            case IR::OpType::ADD:
            case IR::OpType::CALL:
                stackSize += 4;
            }
        }
        return stackSize;
    }

    StringView GetTempAddr(IR::Reg reg)
    {
        return std::move((StringBuilder() << 't' << reg).view());
    }
private:
    void GenerateOp(const IR::Op *op)
    {
        switch (op->GetType())
        {
        case IR::OpType::EXTERN:
        {
            auto extrn = reinterpret_cast<const IR::ExternOp *>(op);
            auto symbol = extrn->symbol();
            printf("extrn '%s' as __%s\n", symbol.c_str(), symbol.c_str());
            printf("%s = PLT __%s\n", symbol.c_str(), symbol.c_str());
        }
        break;
        case IR::OpType::FN:
        {
            auto fn = reinterpret_cast<const IR::FnOp *>(op);
            auto name = fn->name();
            printf("public %s\n", name.c_str());
            printf("%s:\n", name.c_str());
            printf("push rbp\n");
            printf("mov rbp, rsp\n");
            m_stack = new StackAllocator();
            int stackSize = GetStackSize(fn->ops());
            printf("sub rsp, %d\n", stackSize);
            if (fn->params().size > 0)
            {
                auto param_slot = m_stack->Allocate(fn->params().data[0]);
                printf("mov [rbp-%d], edi\n", param_slot.offset);
            }
            for(auto &fOp : fn->ops())
            {
                GenerateOp(fOp);
            }
            m_stack = nullptr;
            printf("leave\nret\n");
        }
        break;
        case IR::OpType::CALL:
        {
            auto call = reinterpret_cast<const IR::CallOp *>(op);
            // TODO: support several arguments
            if (call->args().size == 1)
            {
                auto slot = m_stack->Resolve(GetTempAddr(call->args().data[0]));
                printf("mov edi, [rbp-%d]\n", slot.offset);
            }
            printf("call %s\n", call->callee().c_str());
            auto result_slot = m_stack->Allocate(GetTempAddr(call->result()));
            printf("mov dword [rbp-%d], eax\n", result_slot.offset);
        }
        break;
        case IR::OpType::LOAD_CONST:
        {
            auto lc = reinterpret_cast<const IR::LoadConstOp *>(op);
            auto addr = GetTempAddr(lc->result());
            auto slot = m_stack->Allocate(addr);
            printf("mov dword [rbp-%d], %ld\n", slot.offset, lc->value());
        }
        break;
        case IR::OpType::STORE:
        {
            auto s = reinterpret_cast<const IR::StoreOp *>(op);
            printf("; DEBUG: resolving stack slot at %s\n", s->addr().c_str());
            auto slot = m_stack->Allocate(s->addr());
            auto value_slot = m_stack->Resolve(GetTempAddr(s->src()));
            printf("mov eax, [rbp-%d]\n", value_slot.offset);
            printf("mov dword [rbp-%d], eax\n", slot.offset);
        }
        break;
        case IR::OpType::LOAD:
        {
            auto l = reinterpret_cast<const IR::LoadOp *>(op);
            auto value_slot = m_stack->Allocate(GetTempAddr(l->result()));
            auto variable_slot = m_stack->Resolve(l->addr());
            printf("mov eax, [rbp-%d]\n", variable_slot.offset);
            printf("mov dword [rbp-%d], eax\n", value_slot.offset);
        }
        break;
        case IR::OpType::ADD:
        {
            auto expr = reinterpret_cast<const IR::AddOp *>(op);
            auto lhs_slot = m_stack->Resolve(GetTempAddr(expr->lhs()));
            printf("mov eax, [rbp-%d]\n", lhs_slot.offset);
            auto rhs_slot = m_stack->Resolve(GetTempAddr(expr->rhs()));
            printf("add eax, [rbp-%d]\n", rhs_slot.offset);
            auto result_slot = m_stack->Allocate(GetTempAddr(expr->result()));
            printf("mov dword [rbp-%d], eax\n", result_slot.offset);
        }
        break;
        default: printf("; NOT HANDLED\n; %s\n", op->Format(0).c_str()); break;
        }
    }
public:
    void Generate(const char* filename, View<IR::Op*> ops) override
    {
        printf("; fasm x86_64 linux generated assembly using pl\n");
        printf("format ELF64\n");
        printf("section '.text' executable\n");

        for (auto& op : ops)
        {
            GenerateOp(op);
        }
    }
public:
    // TODO: handle sub-blocks
    StackAllocator* m_stack = nullptr;
};