#include <stdlib.h>
#include <stdio.h>
#include <vector>
#include <set>

#include "../Utils/SgUtils.h"
#include "../Utils/CommonBlock.h"
#include "../GraphCall/graph_calls.h"
#include "../ExpressionTransform/expr_transform.h"

#include "dvm.h"
#include "IR.h"
#include "CFGraph.h"

using namespace std;
using namespace SAPFOR;

typedef SAPFOR::BasicBlock BBlock;
typedef SAPFOR::Argument BArgument;

static void printBlock(BBlock* block) {
    cout << "block - " << block->getNumber() << endl;
    cout << "next -";
    for (auto i : block->getNext())
    {
        cout << " " << i->getNumber();
    }
    cout << endl << "prev -";
    for (auto i : block->getPrev())
    {
        cout << " " << i->getNumber();
    }
    cout << endl;

    for (auto i : block->getInstructions())
    {
        string resValue = "";
        string arg1Value = "";
        string arg2Value = "";
        if (i->getInstruction()->getResult() != nullptr && i->getInstruction()->getResult()->getType() == CFG_ARG_TYPE::VAR) {
            resValue = i->getInstruction()->getResult()->getValue();
            i->getInstruction()->getResult()->setValue(i->getInstruction()->getResult()->getValue() + to_string(i->getInstruction()->getResult()->getNumber()));
        } 
        if (i->getInstruction()->getArg1() != nullptr && i->getInstruction()->getArg1()->getType() == CFG_ARG_TYPE::VAR) {
            arg1Value = i->getInstruction()->getArg1()->getValue();
            i->getInstruction()->getArg1()->setValue(i->getInstruction()->getArg1()->getValue() + to_string(i->getInstruction()->getArg1()->getNumber()));
        }
        if (i->getInstruction()->getArg2() != nullptr && i->getInstruction()->getArg2()->getType() == CFG_ARG_TYPE::VAR) {
            arg2Value = i->getInstruction()->getArg2()->getValue();
            i->getInstruction()->getArg2()->setValue(i->getInstruction()->getArg2()->getValue() + to_string(i->getInstruction()->getArg2()->getNumber()));
        }

        cout << i->getNumber() << " " << i->getInstruction()->dump() << endl;// << (i->getInstruction()->getResult() != nullptr ? "   1 " : "   0 ") << (i->getInstruction()->getArg1() != nullptr ? "1 " : "0 ") << (i->getInstruction()->getArg2() != nullptr ? "1 " : "0 ") << endl;

        /*if (i->getInstruction()->getResult() != nullptr) {
            cout << "num - " << i->getInstruction()->getResult()->getNumber() << endl;
        })*/

        /*if (i->getInstruction()->getOperation() == CFG_OP::F_CALL) {
            cout << endl;
            cout << "arg1 - " << (i->getInstruction()->getArg1()->getType() == CFG_ARG_TYPE::FUNC ? "1" : "0") << "arg2 - " << (i->getInstruction()->getArg2()->getType() == CFG_ARG_TYPE::CONST ? "1" : "0") << "res - " << i->getInstruction()->getResult() << endl;
        }*/
        
        if (i->getInstruction()->getResult() != nullptr && i->getInstruction()->getResult()->getType() == CFG_ARG_TYPE::VAR) {
            i->getInstruction()->getResult()->setValue(resValue);
        }
        if (i->getInstruction()->getArg1() != nullptr && i->getInstruction()->getArg1()->getType() == CFG_ARG_TYPE::VAR) {
            i->getInstruction()->getArg1()->setValue(arg1Value);
        }
        if (i->getInstruction()->getArg2() != nullptr && i->getInstruction()->getArg2()->getType() == CFG_ARG_TYPE::VAR) {
            i->getInstruction()->getArg2()->setValue(arg2Value);
        }
    }

    cout << endl;
}
template <typename T>
static bool compareVectors(const vector<T>* vec1, const vector<T>* vec2) {
    if (vec1 == vec2) {
        return true;
    }

    if (!vec1 || !vec2) {
        return false;
    }

    vector<T> sortedVec1 = *vec1;
    vector<T> sortedVec2 = *vec2;

    sort(sortedVec1.begin(), sortedVec1.end());
    sort(sortedVec2.begin(), sortedVec2.end());

    return sortedVec1 == sortedVec2;
}

template <typename T>
static vector<T>* getCommonElements(const vector<vector<T>>* vectors) {
    if (!vectors || vectors->empty()) {
        return new vector<T>(); // Return an empty vector if input is null or empty
    }

    // Start with the first vector
    vector<T>* commonElements = new vector<T>((*vectors)[0]);

    for (size_t i = 1; i < vectors->size(); ++i) {
        vector<T> tempCommon;

        // Sort the current vector and common result for intersection
        vector<T> sortedVec = (*vectors)[i];
        sort(commonElements->begin(), commonElements->end());
        sort(sortedVec.begin(), sortedVec.end());

        // Find the intersection
        set_intersection(
            commonElements->begin(), commonElements->end(),
            sortedVec.begin(), sortedVec.end(),
            back_inserter(tempCommon)
        );

        // Update common result
        *commonElements = tempCommon;

        // If no common elements left, break early
        if (commonElements->empty()) {
            break;
        }
    }

    return commonElements;
}


static map<BBlock*, vector<BBlock*>> findDominators(vector<BBlock*> blocks) {
    map<BBlock*, vector<BBlock*>> result;

    bool changed = true;
    while (changed) {
        changed = false;

        for (auto currentBlock : blocks) {
            auto pred = currentBlock->getPrev();
            auto prevDominators = new vector<vector<BBlock*>>();

            for (auto predBlock : pred)
                prevDominators->push_back(result.find(predBlock) != result.end() ? result[predBlock] : blocks);

            auto currentBlockResult = getCommonElements(prevDominators);
            currentBlockResult->push_back(currentBlock);

            if (result.find(currentBlock) == result.end() || !compareVectors(currentBlockResult, &result[currentBlock])) {
                result[currentBlock] = *currentBlockResult;
                changed = true;
            }
        }
    }

    return result;
}

static void RenumberBlocks(BBlock* current, int* n, map<int, int>* res, set<BBlock*>* visited) {
    if (visited->find(current) != visited->end()) {
        return;
    }

    visited->insert(current);

    vector<BBlock*> nextBlocks = current->getNext();

    sort(nextBlocks.begin(), nextBlocks.end(), [](BBlock* a, BBlock* b) {
        return a->getInstructions()[0]->getInstruction()->getOperator()->lineNumber() > b->getInstructions()[0]->getInstruction()->getOperator()->lineNumber();
    });

    for (auto i : nextBlocks) {
        RenumberBlocks(i, n, res, visited);
    }

    (*res)[current->getNumber()] = *n;
    *n -= 1;
}

static map<BBlock*, vector<BBlock*>> findDominatorBorders(vector<BBlock*>& blocks, map<BBlock*, BBlock*>& iDominators) {
    map<BBlock*, vector<BBlock*>> result;

    for (auto block : blocks) {
        result[block] = *(new vector<BBlock*>());
    }

    for (auto block : blocks) {
        if (block->getPrev().size() > 1) {
            for (auto prev : block->getPrev()) {
                auto tmpBlock = prev;

                while (tmpBlock != iDominators[block]) {
                    result[tmpBlock].push_back(block);
                    tmpBlock = iDominators[tmpBlock];
                }
            }
        }
    }

    return result;
}

static BBlock* findImmediateDominatorsDfsHelper(BBlock* block, BBlock* currentBlock, BBlock* currentImmediateDominator, vector<BBlock*> &visited, map<BBlock*, vector<BBlock*>>& dominators) {
    if (block == currentBlock) {
        return currentImmediateDominator;
    }
    
    if (find(visited.begin(), visited.end(), currentBlock) != visited.end()) {
        return nullptr;
    }
    
    visited.push_back(currentBlock);

    if (find(dominators[block].begin(), dominators[block].end(), currentBlock) != dominators[block].end()) {
        currentImmediateDominator = currentBlock;
    }

    for (auto nextBlock : currentBlock->getNext()) {
        auto result = findImmediateDominatorsDfsHelper(block, nextBlock, currentImmediateDominator, visited, dominators);

        if (result) {
            return result;
        }
    }

    return nullptr;
}

static map<BBlock*, BBlock*> findImmediateDominators(map<BBlock*, vector<BBlock*>>& dominators, BBlock* fistBlock) {
    map<BBlock*, BBlock*> iDominators;

    for (const auto& domPair : dominators) {
        BBlock* block = domPair.first;
        const vector<BBlock*>& domBlocks = domPair.second;
        vector<BBlock*> visited;

        if (block == fistBlock) {
            continue;
        }

        iDominators[block] = findImmediateDominatorsDfsHelper(block, fistBlock, fistBlock, visited, dominators);
    }

    return iDominators;
}

static vector<BArgument*> getDefForBlock(const BBlock& block) {
    vector<BArgument*> def;
    const auto& instructions = block.getInstructions();

    for (const auto& irBlock : instructions) {
        if (irBlock) {
            Instruction* instr = irBlock->getInstruction();

            if (instr) {
                BArgument* result = instr->getResult();
                if (result) {
                    def.push_back(result);
                }
            }
        }
    }

    return def;
}

static pair<set<BArgument*>, map<BArgument*, set<BBlock*>>> getGlobalsAndVarBlocks(vector<BBlock*> blocks) {
    set<BArgument*> globals;
    map<BArgument*, set<BBlock*>> varBlocks;

    for (auto block : blocks) {
        set<BArgument*> def;
        const auto& instructions = block->getInstructions();

        for (const auto& irBlock : instructions) {
            if (irBlock) {
                Instruction* instr = irBlock->getInstruction();

                if (instr) {
                    auto arg1 = instr->getArg1();
                    auto arg2 = instr->getArg2();
                    auto res = instr->getResult();

                    if (arg1 && arg1->getType() == CFG_ARG_TYPE::VAR && find(def.begin(), def.end(), arg1) == def.end()) {
                        globals.insert(arg1);
                    }
                    if (arg2 && arg2->getType() == CFG_ARG_TYPE::VAR && find(def.begin(), def.end(), arg2) == def.end()) {
                        globals.insert(arg2);
                    }

                    if (res && res->getType() == CFG_ARG_TYPE::VAR) {
                        def.insert(res);
                        varBlocks[res].insert(block);
                    }
                }
            }
        }
    }

    return make_pair(globals, varBlocks);
}

static vector<BBlock*> getBlocksWithFiFunctions(vector<BBlock*> blocks, set<BArgument*> globals, map<BArgument*, set<BBlock*>> varBlocks, map<BBlock*, vector<BBlock*>> dominatorBorders) {
    vector<BBlock*> blocksWithFiFunctions;
    auto fiFunc = new BArgument(CFG_ARG_TYPE::FUNC, CFG_MEM_TYPE::NONE_, "FI_FUNCTION");
    auto paramCount = new BArgument(CFG_ARG_TYPE::CONST, CFG_MEM_TYPE::LOCAL_, "0");

    for (auto var : globals) {
        auto worklist = varBlocks[var];
        set<BBlock*> hasFiFunction;

        while (!worklist.empty()) {
            auto block = *worklist.begin();
            worklist.erase(block);

            for (auto dfBlock : dominatorBorders[block]) {
                if (hasFiFunction.find(dfBlock) == hasFiFunction.end()) {
                    hasFiFunction.insert(dfBlock);

                    Instruction* phiInstruction = new Instruction(CFG_OP::F_CALL, fiFunc, paramCount, var);

                    //cout << "insert fi functio9n in " << dfBlock->getNumber() << " with var " << var->getValue() << endl;

                    IR_Block* phiBlock = new IR_Block(phiInstruction);
                    dfBlock->addInstructionInFront(phiBlock);

                    blocksWithFiFunctions.push_back(dfBlock);
                }
            }
        }
    }

    return blocksWithFiFunctions;
}




string ToString(CFG_ARG_TYPE type) {
    switch (type) {
    case CFG_ARG_TYPE::NONE: return "NONE";
    case CFG_ARG_TYPE::REG: return "REG";
    case CFG_ARG_TYPE::VAR: return "VAR";
    case CFG_ARG_TYPE::ARRAY: return "ARRAY";
    case CFG_ARG_TYPE::CONST: return "CONST";
    case CFG_ARG_TYPE::FUNC: return "FUNC";
    case CFG_ARG_TYPE::LAB: return "LAB";
    case CFG_ARG_TYPE::INSTR: return "INSTR";
    case CFG_ARG_TYPE::CONST_STR: return "CONST_STR";
    case CFG_ARG_TYPE::RECORD: return "RECORD";
    case CFG_ARG_TYPE::CONSTR_REF: return "CONSTR_REF";
    default: return "UNKNOWN";
    }
}

void restoreConnections(const vector<BBlock*>& originalBlocks, vector<BBlock*>& copiedBlocks) {
    // Создаем отображение оригинальных блоков в их копии
    map<BBlock*, BBlock*> blockMapping;
    for (size_t i = 0; i < originalBlocks.size(); ++i) {
        blockMapping[originalBlocks[i]] = copiedBlocks[i];
    }

    // Восстанавливаем связи между копиями
    for (size_t i = 0; i < originalBlocks.size(); ++i) {
        BBlock* originalBlock = originalBlocks[i];
        BBlock* copiedBlock = copiedBlocks[i];

        for (auto j : copiedBlock->getPrev()) {
            copiedBlock->removePrev(j);
        }
        for (auto j : copiedBlock->getNext()) {
            copiedBlock->removeNext(j);
        }

        // Восстанавливаем связи succ (следующих блоков)
        for (auto* succ : originalBlock->getNext()) {
            copiedBlock->addNext(blockMapping[succ]);
        }

        // Восстанавливаем связи prev (предыдущих блоков)
        for (auto* prev : originalBlock->getPrev()) {
            copiedBlock->addPrev(blockMapping[prev]);
        }
    }
}


BArgument* NewName(BArgument* var, map<string, int>& counter, map<string, stack<BArgument*>>& stack) {
    int index = counter[var->getValue()];
    counter[var->getValue()]++;

    BArgument* newName = new BArgument(var->getType(), var->getMemType(), var->getValue(), index + 1);
    stack[var->getValue()].push(newName);
    return newName;
}

void RenameFiFunctionResultVar(BBlock* block, map<string, int>& counter, map<string, stack<BArgument*>>& stack) {
    for (auto irBlock : block->getInstructions()) {
        auto instruction = irBlock->getInstruction();
        if (instruction->getOperation() == CFG_OP::F_CALL && instruction->getArg1() != nullptr && instruction->getArg1()->getValue() == "FI_FUNCTION" && instruction->getResult() != nullptr) {
            instruction->setResult(NewName(instruction->getResult(), counter, stack));
        }
    }
}

void RenameInstructionVars(BBlock* block, map<string, int>& counter, map<string, stack<BArgument*>>& stack) {
    for (auto irBlock : block->getInstructions()) {
        auto instruction = irBlock->getInstruction();

        if (instruction->getArg1() != nullptr && instruction->getArg1()->getType() == CFG_ARG_TYPE::VAR) {
            instruction->setArg1(stack[instruction->getArg1()->getValue()].top());
        }

        if (instruction->getArg2() != nullptr && instruction->getArg2()->getType() == CFG_ARG_TYPE::VAR) {
            instruction->setArg2(stack[instruction->getArg2()->getValue()].top());
        }

        if (instruction->getResult() != nullptr && instruction->getResult()->getType() == CFG_ARG_TYPE::VAR) {
            instruction->setResult(NewName(instruction->getResult(), counter, stack));
        }
    }
}

void RenameFiFunctionArgsVar(BBlock* block, map<string, stack<BArgument*>>& stack) {
    cout << "test" << endl;
    auto size = block->getInstructions().size();
    for (auto i = 0; i < size; i++) {
        auto irBlock = block->getInstructions()[i];
        auto instruction = irBlock->getInstruction();
        if (instruction->getOperation() == CFG_OP::F_CALL && instruction->getArg1() != nullptr && instruction->getArg1()->getValue() == "FI_FUNCTION" && instruction->getResult() != nullptr) {
            //cout << "test" << endl;

            auto paramInstruction = new Instruction(CFG_OP::PARAM, stack[instruction->getResult()->getValue()].top());
            block->addInstructionBeforeInstruction(new IR_Block(paramInstruction), instruction);

            i++;
        }
    }
}

vector<BBlock*> findBlocksWithValue(map<BBlock*, BBlock*>& iDominators, BBlock* x) {
    vector<BBlock*> result;

    // Проходим по всем элементам map
    for (auto& pair : iDominators) {
        // Если значение равно x, добавляем ключ в результат
        if (pair.second == x) {
            result.push_back(pair.first);
        }
    }

    return result;
}

void RenameIR(BBlock* block, map<BBlock*, BBlock*>& iDominators, map<string, int>& counter, map<string, stack<BArgument*>>& stack) {

    RenameFiFunctionResultVar(block, counter, stack);

    RenameInstructionVars(block, counter, stack);

    for (auto* successor : block->getNext()) {
        RenameFiFunctionArgsVar(successor, stack);
    }

    for (auto* child : findBlocksWithValue(iDominators, block)) {
        cout << "sec" << endl;
        RenameIR(child, iDominators, counter, stack);
    }

    for (auto& irBlock : block->getInstructions()) {
        auto instruction = irBlock->getInstruction();
        if (instruction->getResult() != nullptr && instruction->getResult()->getType() == CFG_ARG_TYPE::VAR) {
            string varName = instruction->getResult()->getValue();
            stack[varName].pop();
        }
    }

    for (auto& irBlock : block->getInstructions()) {
        auto instruction = irBlock->getInstruction();
        if (instruction->getOperation() == CFG_OP::F_CALL && instruction->getArg1() != nullptr && instruction->getArg1()->getValue() == "FI_FUNCTION" && instruction->getResult() != nullptr) {
            string varName = instruction->getResult()->getValue();
            stack[varName].pop();
        }
    }
}

map<FuncInfo*, vector<BBlock*>> buildIRSSAForm(map<FuncInfo*, vector<BBlock*>> fullIR) {

    map<FuncInfo*, vector<BBlock*>> result;

    for (auto item : fullIR) {
        auto funcinfo = item.first;
        auto funcIRConst = item.second;

        vector<BBlock*> funcIR;
        vector<BBlock*> funcResult;

        for (auto i : funcIRConst) {
            funcIR.push_back(new BBlock(*i));
        }
        restoreConnections(funcIRConst, funcIR);

        for (auto i : funcIR) {
            //printBlock(i);
        }

        auto dominators = findDominators(funcIR);

        /*cout << endl << endl << endl << "DOMINATORS" << endl << endl << endl;

        for (auto i : dominators) {
            cout << "block - " << i.first->getNumber() << endl;
            for (auto j : i.second) {
                cout << "dominators - " << j->getNumber() << endl;
            }

            cout << endl;
        }*/

        auto iDominators = findImmediateDominators(dominators, funcIR[0]);

        /*for (auto i : iDominators) {
            cout << "block - " << i.first->getNumber() << endl;
            cout << "Idominators - " << i.second->getNumber() << endl;
            cout << endl;
        }
        */
        auto dominatorBorders = findDominatorBorders(funcIR, iDominators);
        

        /*for (auto i : dominatorBorders) {
            cout << "block - " << i.first->getNumber() << endl;
            for (auto j : i.second) {
                cout << "border - " << j->getNumber() << endl;
            }

            cout << endl;
        }*/


        auto globalsAndVarBlocks = getGlobalsAndVarBlocks(funcIR);
        auto globals = globalsAndVarBlocks.first;
        auto varBlocks = globalsAndVarBlocks.second;

        /*for (auto i : globals) {
            cout << i->getValue() << " " << ToString(i->getType()) << " " << i->getNumber() << endl;

        }
        cout << endl;
        for (auto i : varBlocks) {
            cout << i.first->getValue() << " - ";
            for (auto j : i.second) {
                cout << j->getNumber() << ", ";
            }
            cout << endl;
        }
        cout << endl;
        */

        funcResult = getBlocksWithFiFunctions(funcIR, globals, varBlocks, dominatorBorders);

        map<string, int> count;
        map<string, stack<BArgument*>> varStack;

        for (auto var : globals) {
            count[var->getValue()] = 0;

            stack<BArgument*> tmp;
            varStack[var->getValue()] = tmp;
        }

        RenameIR(funcIR[0], iDominators, count, varStack);

        for (auto i : funcIR) {
            printBlock(i);
        }
        cout << endl << endl << endl << endl << endl;
        for (auto i : funcIRConst) {
            printBlock(i);
        }


        result[funcinfo] = funcResult;
    }

    return result;
}