#include "CodeGen.h"

CodeGen::CodeGen()
{
}

CodeGen::~CodeGen()
{
    delete _globalScope;
}

void CodeGen::GenerateCode(TreeNode *root)
{
    rootNodeAST = root;

    _globalScope = new Scope();
    _currentScope = _globalScope;
    compileNode(rootNodeAST);

    _currentEmitIndex = 0;
    _endEmitIndex = 0;
}

void CodeGen::Emit(const std::string &code)
{
    _emittedCode.push_back(code);
}

void CodeGen::BeginEmit()
{
    _currentEmitIndex = _emittedCode.size();
}

void CodeGen::EndEmit()
{
    _endEmitIndex = _emittedCode.size() - 1;
}

void CodeGen::PasteEmittedCodeHere()
{
    std::vector<std::string> new_emittedCodeList;

    for (size_t i = 0; i < _emittedCode.size(); ++i)
    {
        if (i >= _currentEmitIndex && i <= _endEmitIndex)
        {
            continue;
        }
        new_emittedCodeList.push_back(_emittedCode[i]);
    }

    for (int i = _currentEmitIndex; i <= _endEmitIndex; ++i)
    {
        new_emittedCodeList.push_back(_emittedCode[i]);
    }

    _emittedCode = new_emittedCodeList;
}

void CodeGen::PrintGeneratedCode()
{
    for (size_t i = 0; i < _emittedCode.size(); ++i)
    {
        std::cout << _emittedCode[i] << std::endl;
    }
}

void CodeGen::compileNode(TreeNode *node)
{
    if (node == nullptr)
    {
        return;
    }
    Symbol* symbol = nullptr;
    switch (node->m_type)
    {
    case NodeType::NODE_COMPILATION_UNIT:
        for (TreeNode *child : node->m_children)
        {
            compileNode(child);
        }
        break;
    case NodeType::NODE_FUNCTIONDECL:
        _currentScope->addSymbol(new Symbol(node->getChild(0)->getChildIdentifier(1), SymbolType::SYMBOL_TYPE_FUNCTION, node->getChild(0)->getChild(0)->m_type));
        Emit("_" + _currentScope->symbolTable->getSymbolByName(node->getChild(0)->getChildIdentifier(1))->_name + ":");

        // compile code block
        BeginEmit();
        _currentScope = _currentScope->createScope();
        compileNode(node->getChild(1));
        EndEmit();

        Emit("push ebp");
        Emit("mov ebp, esp");
        Emit("sub esp, " + std::to_string(4 * _currentScope->symbolTable->getSymbolCount()));

        PasteEmittedCodeHere();

        Emit("mov esp, ebp");
        Emit("pop ebp");
        Emit("ret");

        _currentScope = _currentScope->leaveScope();
        break;
    case NodeType::NODE_BLOCK:
        for (TreeNode *child : node->m_children)
        {
            compileNode(child);
        }
        break;
    case NodeType::NODE_DECLARATION:
        _currentScope->addSymbol(new Symbol(node->getChild(1)->m_token._identifier));
        for (TreeNode *child : node->m_children)
        {
            compileNode(child);
        }

        Emit("mov [ebp - " + std::to_string(_currentScope->symbolTable->getSymbolByName(node->getChild(1)->m_token._identifier)->_address) + "], eax");
        break;
    case NodeType::NODE_ASSIGN:
        if (node->getChild(0)->m_type == NodeType::NODE_FACTOR)
        {
            if (node->getChild(0)->getChild(0)->m_type == NodeType::NODE_IDENTIFIER)
            {
                Emit("mov eax, [ebp - " + std::to_string(_currentScope->symbolTable->getSymbolByName(node->getChild(0)->getChild(0)->m_token._identifier)->_address) + "]");
            }
        }
        else
        {
            compileNode(node->getChild(0));
        }
        break;
    case NodeType::NODE_IDENTIFIER:
        compileNode(node->getChild(0));
        break;
    case NodeType::NODE_ASSEMBLY_LINE:
        Emit(node->m_token._identifier);
        break;
    case NodeType::NODE_BINARY_ADD:
        if(node->getChild(0)->m_type == NodeType::NODE_FACTOR)
        {
            if(node->getChild(0)->getChild(0)->m_type == NodeType::NODE_NUMBER)
            {
                Emit("mov eax, " + std::to_string(node->getChild(0)->getChild(0)->m_token._value.uintvalue));
            } else if(node->getChild(0)->getChild(0)->m_type == NodeType::NODE_IDENTIFIER)
            {
                Emit("mov eax, [ebp - " + std::to_string(_currentScope->symbolTable->getSymbolByName(node->getChild(0)->getChild(0)->m_token._identifier)->_address) + "]");
            }
        } else {
            compileNode(node->getChild(0));
        }

        if(node->getChild(1)->m_type == NodeType::NODE_FACTOR)
        {
            if(node->getChild(1)->getChild(0)->m_type == NodeType::NODE_NUMBER)
            {
                Emit("mov ebx, " + std::to_string(node->getChild(1)->getChild(0)->m_token._value.uintvalue));
            } else if(node->getChild(1)->getChild(0)->m_type == NodeType::NODE_IDENTIFIER)
            {
                Emit("mov ebx, [ebp - " + std::to_string(_currentScope->symbolTable->getSymbolByName(node->getChild(1)->getChild(0)->m_token._identifier)->_address) + "]");
            }
        } else {
            compileNode(node->getChild(1));
        }

        Emit("add eax, ebx");
        break;
    case NodeType::NODE_BINARY_SUB:
        if(node->getChild(0)->m_type == NodeType::NODE_FACTOR)
        {
            if(node->getChild(0)->getChild(0)->m_type == NodeType::NODE_NUMBER)
            {
                Emit("mov eax, " + std::to_string(node->getChild(0)->getChild(0)->m_token._value.uintvalue));
            } else if(node->getChild(0)->getChild(0)->m_type == NodeType::NODE_IDENTIFIER)
            {
                Emit("mov eax, [ebp - " + std::to_string(_currentScope->symbolTable->getSymbolByName(node->getChild(0)->getChild(0)->m_token._identifier)->_address) + "]");
            }
        } else {
            compileNode(node->getChild(0));
        }

        if(node->getChild(1)->m_type == NodeType::NODE_FACTOR)
        {
            if(node->getChild(1)->getChild(0)->m_type == NodeType::NODE_NUMBER)
            {
                Emit("mov ebx, " + std::to_string(node->getChild(1)->getChild(0)->m_token._value.uintvalue));
            } else if(node->getChild(1)->getChild(0)->m_type == NodeType::NODE_IDENTIFIER)
            {
                Emit("mov ebx, [ebp - " + std::to_string(_currentScope->symbolTable->getSymbolByName(node->getChild(1)->getChild(0)->m_token._identifier)->_address) + "]");
            }
        } else {
            compileNode(node->getChild(1));
        }

        Emit("sub eax, ebx");
        break;
    case NodeType::NODE_MUTIPLICATION:
        if(node->getChild(0)->m_type == NodeType::NODE_FACTOR)
        {
            if(node->getChild(0)->getChild(0)->m_type == NodeType::NODE_NUMBER)
            {
                Emit("mov eax, " + std::to_string(node->getChild(0)->getChild(0)->m_token._value.uintvalue));
            } else if(node->getChild(0)->getChild(0)->m_type == NodeType::NODE_IDENTIFIER)
            {
                Emit("mov eax, [ebp - " + std::to_string(_currentScope->symbolTable->getSymbolByName(node->getChild(0)->getChild(0)->m_token._identifier)->_address) + "]");
            }
        } else {
            compileNode(node->getChild(0));
        }

        if(node->getChild(1)->m_type == NodeType::NODE_FACTOR)
        {
            if(node->getChild(1)->getChild(0)->m_type == NodeType::NODE_NUMBER)
            {
                Emit("mov ebx, " + std::to_string(node->getChild(1)->getChild(0)->m_token._value.uintvalue));
            } else if(node->getChild(1)->getChild(0)->m_type == NodeType::NODE_IDENTIFIER)
            {
                Emit("mov ebx, [ebp - " + std::to_string(_currentScope->symbolTable->getSymbolByName(node->getChild(1)->getChild(0)->m_token._identifier)->_address) + "]");
            }
        } else {
            compileNode(node->getChild(1));
        }

        Emit("imul eax, ebx");
        break;
    case NodeType::NODE_DIVISION:
        if(node->getChild(0)->m_type == NodeType::NODE_FACTOR)
        {
            if(node->getChild(0)->getChild(0)->m_type == NodeType::NODE_NUMBER)
            {
                Emit("mov eax, " + std::to_string(node->getChild(0)->getChild(0)->m_token._value.uintvalue));
            } else if(node->getChild(0)->getChild(0)->m_type == NodeType::NODE_IDENTIFIER)
            {
                Emit("mov eax, [ebp - " + std::to_string(_currentScope->symbolTable->getSymbolByName(node->getChild(0)->getChild(0)->m_token._identifier)->_address) + "]");
            }
        } else {
            compileNode(node->getChild(0));
        }

        if(node->getChild(1)->m_type == NodeType::NODE_FACTOR)
        {
            if(node->getChild(1)->getChild(0)->m_type == NodeType::NODE_NUMBER)
            {
                Emit("mov ebx, " + std::to_string(node->getChild(1)->getChild(0)->m_token._value.uintvalue));
            } else if(node->getChild(1)->getChild(0)->m_type == NodeType::NODE_IDENTIFIER)
            {
                Emit("mov ebx, [ebp - " + std::to_string(_currentScope->symbolTable->getSymbolByName(node->getChild(1)->getChild(0)->m_token._identifier)->_address) + "]");
            }
        } else {
            compileNode(node->getChild(1));
        }

        Emit("cdq");
        Emit("idiv ebx");
        Emit("mov eax, eax");
        break;
    case NodeType::NODE_FUNCARGS:
        for(TreeNode* child : node->m_children)
        {
            if(child->m_type == NodeType::NODE_IDENTIFIER)
            {
                symbol = _currentScope->symbolTable->getSymbolByName(child->m_token._identifier);
                Emit("mov eax, [ebp - " + std::to_string(symbol->_address) + "]");
                Emit("push eax");
            } else if(child->m_type == NodeType::NODE_NUMBER)
            {
                Emit("push " + std::to_string(child->m_token._value.uintvalue));
            }
        }
        symbol = nullptr;
        break;
    case NodeType::NODE_FUNCCALL:
        symbol = _currentScope->parent->symbolTable->getSymbolByName(node->getChildIdentifier(0));
        if(symbol == nullptr)
        {
            symbol = _currentScope->symbolTable->getSymbolByName(node->getChildIdentifier(0));
        }

        if(node->getChild(1)->m_type == NodeType::NODE_FUNCARGS)
        {
            compileNode(node->getChild(1));
        }

        if(symbol == nullptr)
        {
            std::cout << "Error: Function " << node->getChildIdentifier(0) << " not declared!" << std::endl;
            return;
        }

        Emit("call _" + symbol->_name);
        symbol = nullptr;
        break;
    case NodeType::NODE_IF:
        compileNode(node->getChild(0));
        Emit("cmp eax, 0");
        Emit("je " + node->getChildIdentifier(1));
        compileNode(node->getChild(2));
        break;
    case NodeType::NODE_WHILE:
        Emit("_" + node->getChildIdentifier(0) + ":");
        compileNode(node->getChild(1));
        Emit("cmp eax, 0");
        Emit("je _" + node->getChildIdentifier(2));
        Emit("jmp _" + node->getChildIdentifier(0));
        Emit("_" + node->getChildIdentifier(2) + ":");
        break;
    case NodeType::NODE_RETURN:
        compileNode(node->getChild(0));
        Emit("mov eax, ebx");
        Emit("jmp _" + node->getChildIdentifier(1));
        break;
    case NodeType::NODE_ELSE:
        compileNode(node->getChild(0));
        Emit("cmp eax, 0");
        Emit("je " + node->getChildIdentifier(1));
        compileNode(node->getChild(2));
        break;
    case NodeType::NODE_LOR:
        compileNode(node->getChild(0));
        Emit("cmp eax, 0");
        Emit("jne " + node->getChildIdentifier(1));
        compileNode(node->getChild(2));
        break;
    case NodeType::NODE_LAND:
        compileNode(node->getChild(0));
        Emit("cmp eax, 0");
        Emit("je " + node->getChildIdentifier(1));
        compileNode(node->getChild(2));
        break;
    case NodeType::NODE_LNOT:
        compileNode(node->getChild(0));
        Emit("cmp eax, 0");
        Emit("je _" + node->getChildIdentifier(1));
        Emit("mov eax, 0");
        Emit("jmp _" + node->getChildIdentifier(1));
        Emit("_" + node->getChildIdentifier(1) + ":");
        break;
    case NodeType::NODE_EQUALITY:
        compileNode(node->getChild(0));
        compileNode(node->getChild(1));
        Emit("cmp eax, ebx");
        Emit("je " + node->getChildIdentifier(2));
        break;
    case NodeType::NODE_INEQUALITY:
        compileNode(node->getChild(0));
        compileNode(node->getChild(1));
        Emit("cmp eax, ebx");
        Emit("jne " + node->getChildIdentifier(2));
        break;
    case NodeType::NODE_GREATER:
        compileNode(node->getChild(0));
        compileNode(node->getChild(1));
        Emit("cmp eax, ebx");
        Emit("jg " + node->getChildIdentifier(2));
        break;
    case NodeType::NODE_LESS:
        compileNode(node->getChild(0));
        compileNode(node->getChild(1));
        Emit("cmp eax, ebx");
        Emit("jl " + node->getChildIdentifier(2));
        break;
    case NodeType::NODE_GERATER_EQUAL:
        compileNode(node->getChild(0));
        compileNode(node->getChild(1));
        Emit("cmp eax, ebx");
        Emit("jge " + node->getChildIdentifier(2));
        break;
    case NodeType::NODE_LESS_EQUAL:
        compileNode(node->getChild(0));
        compileNode(node->getChild(1));
        Emit("cmp eax, ebx");
        Emit("jle " + node->getChildIdentifier(2));
        break;
    case NodeType::NODE_COMMENT:
        break;
    default:
        for(TreeNode* child : node->m_children)
        {
            compileNode(child);
        }
        break;
    }
}