#pragma once
#include "codegen/codegen.hpp"
#include "ir/op.hpp"
#include "ir/ops.hpp"
#include "ir/value.hpp"
#include "prelude/linkedlist.hpp"
#include "prelude/string.hpp"
#include <unordered_map>

class FasmX86_64Generator : public CodeGenerator
{
public:
    FasmX86_64Generator() = default;

public:
    bool Generate(const DoubleLinkedList<IR::Op*>* ops) override
    {
        m_ops = ops;

        output().Extend("format ELF64\n");
        output().Extend("section '.text' executable\n");

        for (m_current = m_ops->Begin(); m_current != nullptr; node_next())
        {
            GenerateStatement();
        }

        return true;
    }

private:
    void GenerateExtern()
    {
        auto extrn = current<IR::ExternOp>();
        // TODO: instead of __<symbol, use some random UID or hash string
        // for safety and collision considerations
        output().AppendFormat("extrn '%s' as __%s\n", extrn->symbol().c_str(), extrn->symbol().c_str());
        output().AppendFormat("%s = PLT __%s\n", extrn->symbol().c_str(), extrn->symbol().c_str());
    }

    void GenerateFunction()
    {
        auto fn = current<IR::FnOp>();

        m_slots.clear();
        m_stackCounter = 0;

        output().AppendFormat("public %s\n", fn->name().c_str());
        output().AppendFormat("%s:\n", fn->name().c_str());
        output().Extend("  push rbp\n");
        output().Extend("  mov rbp, rsp\n");

        auto fnNode = node<IR::Op>();
        auto ops = m_ops;
        m_ops = &fn->body().ops();
        for (m_current = m_ops->Begin(); m_current != nullptr; node_next())
        {
            GenerateStatement();
        }
        m_ops = ops;
        m_current = fnNode;

        output().Extend("  leave\n");
        output().Extend("  ret\n");
    }

    void GenerateCall()
    {
        auto call = current<IR::CallOp>();
        // TODO: support stack spilled arguments
        assert(call->args().size < 7 && "stack arguments not supported yet");
        const char *regs[6] = {"edi", "esi", "edx", "ecx","e8", "e9"};
        for (size_t i = 0; i < call->args().size; ++i)
        {
            auto arg = call->args().data[i];
            if (arg->HasId()) {
                auto sp = EnsureSlot(arg);
                // TODO:
                auto size = "dword";
                output().AppendFormat("  mov %s, %s [rbp-%d]\n", regs[i], size, sp);
            } else {
                output().AppendFormat("  mov %s, %d\n", regs[i], reinterpret_cast<IR::ConstantInt*>(arg)->GetValue());
            }
        }
        output().AppendFormat("  call %s\n", call->callee().c_str());
        auto sp = EnsureSlot(call->result());
        auto size = "dword";
        output().AppendFormat("  mov %s [rbp-%d], eax\n", size, sp);
    }

    void GenerateAllocate()
    {
        auto totalAllocSize = 0;
        while (current<IR::Op>()->GetType() == IR::OpType::ALLOCATE)
        {
            auto alloc = current<IR::AllocateOp>();
            // TODO: support other types
            assert(alloc->Type()->kind == IR::ValueHandle::Type::Kind::Int);
            // TODO: dynamic size + alignment
            auto allocSize = 4;
            totalAllocSize += allocSize;
            m_stackCounter += allocSize;
            m_slots.insert(std::make_pair(alloc->result()->GetId(), m_stackCounter));

            if (seek<IR::Op>() && seek<IR::Op>()->get()->GetType() == IR::OpType::ALLOCATE) node_next();
            else break;
        };
        output().AppendFormat("  sub rsp, %d\n", totalAllocSize);
    }

    void GenerateStore()
    {
        auto store = current<IR::StoreOp>();
        auto sp = EnsureSlot(store->dst());

        // TODO: support other types
        assert(store->src()->GetType()->kind == IR::ValueHandle::Type::Kind::Int);
        if (!store->src()->HasId()) {
            auto value = reinterpret_cast<const IR::ConstantInt*>(store->src());
            auto size = "dword";
            output().AppendFormat("  mov %s [rbp-%d], %d\n", size, sp, value->GetValue());
        } else {
            auto ssp = EnsureSlot(store->src());
            auto size = "dword";
            output().AppendFormat("  mov eax, %s [rbp-%d]\n", size, ssp);
            output().AppendFormat("  mov %s [rbp-%d], eax\n", size, sp);
        }
    }

    void GenerateLoad()
    {
        auto load = current<IR::LoadOp>();
        auto sp = EnsureSlot(load->Ptr());

        // TODO: support other types
        auto size = "dword";
        output().AppendFormat("  mov eax, %s [rbp-%d]\n", size, sp);

        sp = EnsureSlot(load->result());
        output().AppendFormat("  mov dword [rbp-%d], eax\n", sp);
    }

    void GenerateMath()
    {
        auto math = current<IR::MathOp>();
        StringBuilder sb;

        switch(math->GetType())
        {
        case IR::OpType::ADD:
            sb << "add";
            break;
        case IR::OpType::SUB:
            sb << "sub";
            break;
        case IR::OpType::MUL:
            sb << "imul";
            break;
        case IR::OpType::DIV:
            sb << "div";
            break;
        default: assert(false && "unreachable or not implemented");
        }

        auto op = sb.view();

        // TODO:
        auto size = "dword";
        if (!math->lhs()->HasId()) {
            output().AppendFormat("  mov eax, %d\n", reinterpret_cast<IR::ConstantInt*>(math->lhs())->GetValue());
        } else {
            auto lsp = EnsureSlot(math->lhs());
            output().AppendFormat("  mov eax, %s [rbp-%d]\n", size, lsp);
        }
        if (!math->rhs()->HasId()) {
            output().AppendFormat("  %s eax, %d\n", op.c_str(), reinterpret_cast<IR::ConstantInt*>(math->rhs())->GetValue());
        } else {
            auto lsp = EnsureSlot(math->rhs());
            output().AppendFormat("  %s eax, %s [rbp-%d]\n", op.c_str(), size, lsp);
        }

        auto sp = EnsureSlot(math->result());
        output().AppendFormat("  mov %s [rbp-%d], eax\n", size, sp);
    }

    void GenerateStatement()
    {
        auto op = current<IR::Op>();
        switch(op->GetType())
        {
        case IR::OpType::EXTERN:
            return GenerateExtern();
        case IR::OpType::FN:
            return GenerateFunction();
        case IR::OpType::CALL:
            return GenerateCall();
        case IR::OpType::ALLOCATE:
            return GenerateAllocate();
        case IR::OpType::STORE:
            return GenerateStore();
        case IR::OpType::LOAD:
            return GenerateLoad();
        case IR::OpType::ADD:
        case IR::OpType::SUB:
        case IR::OpType::MUL:
        case IR::OpType::DIV:
            return GenerateMath();
        // TODO:
        default: output().AppendFormat("  ; %d not implemented\n", op->GetType());
        }
    }

private:
    uint32_t EnsureSlot(const IR::ValueHandle* value)
    {
        assert(value->HasId());
        auto stackPointer = m_slots.find(value->GetId());
        if (stackPointer != m_slots.end())
        {
            return stackPointer->second;
        }

        // TODO: dynamic size based on type
        auto allocSize = 4;
        output().AppendFormat("  sub rsp, %d\n", allocSize);
        m_stackCounter += allocSize;
        m_slots.insert(std::make_pair(value->GetId(), m_stackCounter));
        return m_stackCounter;
    }

private:
    template<typename T>
    ListNode<T*>* node() { return reinterpret_cast<ListNode<T*>*>(m_current); }
    template<typename T>
    ListNode<T*>* seek() { return m_current && m_current->next ? reinterpret_cast<ListNode<T*>*>(m_current->next) : nullptr; }

    void node_next() { assert(m_current); m_current = m_current->next; }

    template<typename T>
    T* current() const { return reinterpret_cast<T*>(m_current->value); }

private:
    const DoubleLinkedList<IR::Op*>* m_ops;
    ListNode<IR::Op*>* m_current;
    std::unordered_map<uint32_t, uint32_t> m_slots;
    uint32_t m_stackCounter = 0;
};