#pragma once

#include "codegen/codegen.hpp"
#include "ir/slot.hpp"
#include "ir/value.hpp"

class StackFasmX86_64Generator : public CodeGenerator
{
public:
    ~StackFasmX86_64Generator() override = default;
private:
    StringView GetTempAddr(IR::Value reg)
    {
        return std::move((StringBuilder() << reg.Format()).view());
    }

    StringView GetSlotAddr(const IR::IRSlot& slot)
    {
        switch (slot.GetType())
        {
        case IR::IRSlot::Type::REGISTRY:
        {
            StringBuilder sb;
            sb.AppendFormat("r%d", slot.GetSlot() + 10); // for r10, r11, r12 etc.
            return sb.view();
        } break;
        case IR::IRSlot::Type::STACK:
        {
            StringBuilder sb;
            sb.AppendFormat("[rbp-%d]", slot.GetSlot()); // for r10, r11, r12 etc.
            return sb.view();
        } break;
        default:
            assert(0 && "TODO: either unreachable or handle properly");
        }
    }
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();
            appendf("extrn '%s' as __%s\n", symbol.c_str(), symbol.c_str());
            appendf("%s = PLT __%s\n", symbol.c_str(), symbol.c_str());
        }
        break;
        case IR::OpType::FN:
        {
            m_allocator = new SlotAllocator();

            auto fn = reinterpret_cast<const IR::FnOp *>(op);
            auto name = fn->name();
            appendf("public %s\n", name.c_str());
            appendf("%s:\n", name.c_str());
            appendf("push rbp\n");
            appendf("mov rbp, rsp\n");

            StringBuilder fnOutput;
            StringBuilder* backup = m_output;
            m_output = &fnOutput;
            
            if (fn->params().size > 0)
            {
                // TODO: support multiple parameters
                auto param_slot = m_allocator->Allocate(fn->params().data[0]);
                appendf("mov %s, rdi\n", GetSlotAddr(param_slot).c_str());
            }
            for(auto &fOp : fn->ops())
            {
                GenerateOp(fOp);
            }

            int stackSize = m_allocator->GetStackSize();
            m_output = backup;
            appendf("sub rsp, %d\n", stackSize);
            *m_output << fnOutput.c_str();

            appendf("leave\nret\n");

            m_allocator = nullptr;
        }
        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_allocator->Resolve(GetTempAddr(call->args().data[0]));
                appendf("mov rdi, %s\n", GetSlotAddr(slot).c_str());
            }
            appendf("call %s\n", call->callee().c_str());
            auto result_slot = m_allocator->Allocate(GetTempAddr(call->result()));
            appendf("mov %s, rax\n", GetSlotAddr(result_slot).c_str());
        }
        break;
        case IR::OpType::LOAD_CONST:
        {
            auto lc = reinterpret_cast<const IR::LoadConstOp *>(op);
            auto addr = GetTempAddr(lc->result());
            auto slot = m_allocator->Allocate(addr);
            appendf("mov %s, %ld\n", GetSlotAddr(slot).c_str(), lc->value());
        }
        break;
        case IR::OpType::STORE:
        {
            auto s = reinterpret_cast<const IR::StoreOp *>(op);
            appendf("; DEBUG: resolving stack slot at %s\n", s->addr().c_str());
            auto slot = m_allocator->Allocate(s->addr());
            auto value_slot = m_allocator->Resolve(GetTempAddr(s->src()));
            appendf("mov rax, %s\n", GetSlotAddr(value_slot).c_str());
            appendf("mov %s, rax\n", GetSlotAddr(slot).c_str());
        }
        break;
        case IR::OpType::LOAD:
        {
            auto l = reinterpret_cast<const IR::LoadOp *>(op);
            auto value_slot = m_allocator->Allocate(GetTempAddr(l->result()));
            auto variable_slot = m_allocator->Resolve(l->addr());
            appendf("mov rax, %s\n", GetSlotAddr(variable_slot).c_str());
            appendf("mov %s, rax\n", GetSlotAddr(value_slot).c_str());
        }
        break;
        case IR::OpType::ADD:
        {
            auto expr = reinterpret_cast<const IR::AddOp *>(op);
            auto lhs_slot = m_allocator->Resolve(GetTempAddr(expr->lhs()));
            appendf("mov rax, %s\n", GetSlotAddr(lhs_slot).c_str());
            auto rhs_slot = m_allocator->Resolve(GetTempAddr(expr->rhs()));
            appendf("add rax, %s\n", GetSlotAddr(rhs_slot).c_str());
            auto result_slot = m_allocator->Allocate(GetTempAddr(expr->result()));
            appendf("mov %s, rax\n", GetSlotAddr(result_slot).c_str());
        }
        break;
        default: appendf("; NOT HANDLED\n; %s\n", op->Format(0).c_str()); break;
        }
    }

    void appendf(const char * fmt, ...)
    {
        assert(m_output != nullptr && "nowhere to write");
        va_list args;
        va_start(args, fmt);
        m_output->VAppendFormat(fmt, args);
        va_end(args);
    }
public:
    StringView GetOutput() { return m_output->view(); }
private:
    StringBuilder* m_output = nullptr;
public:
    void Generate(const char* filename, View<IR::Op*> ops) override
    {
        m_output = new StringBuilder();
        appendf("; fasm x86_64 linux generated assembly using pl\n");
        appendf("format ELF64\n");
        appendf("section '.text' executable\n");

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