#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 IR::OpView* ops) override
    {
        output().Extend("format ELF64\n");
        output().Extend("section '.text' executable\n");

        for (size_t i = 0; i < ops->size; ++i)
        {
            GenerateStatement(ops->data[i]);
        }

        return true;
    }

private:
    void GenerateExtern(IR::ExternOp* extrn)
    {
        // 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(IR::FnOp* fn)
    {
        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");

        for (auto cur = fn->body().ops().Begin(); cur != nullptr; cur = cur->next)
        {
            GenerateStatement(cur->value);
        }

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

    void GenerateCall(IR::CallOp* call)
    {
        // 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(IR::AllocateOp* alloc)
    {
        // TODO: support other types
        assert(alloc->Type()->kind == IR::ValueHandle::Type::Kind::Int);

        EnsureSlot(alloc->result());
    }

    void GenerateStore(IR::StoreOp* store)
    {
        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(IR::LoadOp* load)
    {
        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(IR::MathOp* math)
    {
        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(IR::Op* op)
    {
    switch(op->GetType())
    {
    case IR::OpType::EXTERN:
        return GenerateExtern(reinterpret_cast<IR::ExternOp*>(op));
    case IR::OpType::FN:
        return GenerateFunction(reinterpret_cast<IR::FnOp*>(op));
    case IR::OpType::CALL:
        return GenerateCall(reinterpret_cast<IR::CallOp*>(op));
    case IR::OpType::ALLOCATE:
        return GenerateAllocate(reinterpret_cast<IR::AllocateOp*>(op));
    case IR::OpType::STORE:
        return GenerateStore(reinterpret_cast<IR::StoreOp*>(op));
    case IR::OpType::LOAD:
        return GenerateLoad(reinterpret_cast<IR::LoadOp*>(op));
    case IR::OpType::ADD:
    case IR::OpType::SUB:
    case IR::OpType::MUL:
    case IR::OpType::DIV:
        return GenerateMath(reinterpret_cast<IR::MathOp*>(op));
    // 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:
    std::unordered_map<uint32_t, uint32_t> m_slots;
    uint32_t m_stackCounter = 0;
};