#include <map>
#include <set>
#include <vector>
#include <string>
#include <iostream>
#include <stack>

#include "../CFGraph/IR.h"
#include "GraphCall/graph_calls.h"
#include "implicit_loops_analyzer.h"

using std::map;
using std::set;
using std::vector;
using std::pair;
using std::string;
using std::cout;
using std::endl;
using std::make_pair;
using std::to_string;

enum VisitState { UNVISITED = 0, VISITING = 1, VISITED = 2 };

void dfs(SAPFOR::BasicBlock* block, map<int, int>& visit, vector<pair<SAPFOR::BasicBlock*, SAPFOR::BasicBlock*>>& startAndEnd, SAPFOR::BasicBlock* prev) {
    if (!block) return;

    if (visit[block->getNumber()] == VISITED) {
        cout << "error";
        return;
    }

    if (visit[block->getNumber()] == VISITING) {
        visit[block->getNumber()] = VISITED;
        startAndEnd.push_back(make_pair(prev, block));
        return;
    }

    visit[block->getNumber()] = VISITING;
    for (auto i : block->getNext()) {
        dfs(i, visit, startAndEnd, block);
    }
}

static void printBlock(SAPFOR::BasicBlock* 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() == SAPFOR::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() == SAPFOR::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() == SAPFOR::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;

        if (i->getInstruction()->getResult() != nullptr && i->getInstruction()->getResult()->getType() == SAPFOR::CFG_ARG_TYPE::VAR) {
            i->getInstruction()->getResult()->setValue(resValue);
        }
        if (i->getInstruction()->getArg1() != nullptr && i->getInstruction()->getArg1()->getType() == SAPFOR::CFG_ARG_TYPE::VAR) {
            i->getInstruction()->getArg1()->setValue(arg1Value);
        }
        if (i->getInstruction()->getArg2() != nullptr && i->getInstruction()->getArg2()->getType() == SAPFOR::CFG_ARG_TYPE::VAR) {
            i->getInstruction()->getArg2()->setValue(arg2Value);
        }
    }

    cout << endl;
}

void getLoopBody(SAPFOR::BasicBlock* loopHeader, const set<SAPFOR::BasicBlock*>& loopExits, std::vector<SAPFOR::BasicBlock*>& loopBody) {
    set<SAPFOR::BasicBlock*> visited;
    std::stack<SAPFOR::BasicBlock*> stack;

    stack.push(loopHeader);

    while (!stack.empty()) {
        auto block = stack.top();
        stack.pop();

        if (visited.count(block)) continue;
        visited.insert(block);

        for (auto succ : block->getNext()) {
            if (loopExits.count(succ)) continue;
            if (!visited.count(succ)) {
                stack.push(succ);
            }
        }
    }

    set<SAPFOR::BasicBlock*> backReachable;
    std::stack<SAPFOR::BasicBlock*> reverseStack;
    reverseStack.push(loopHeader);

    while (!reverseStack.empty()) {
        auto block = reverseStack.top();
        reverseStack.pop();

        if (backReachable.count(block)) continue;
        backReachable.insert(block);

        for (auto pred : block->getPrev()) {
            if (visited.count(pred) && !backReachable.count(pred)) {
                reverseStack.push(pred);
            }
        }
    }

    for (auto block : visited) {
        if (backReachable.count(block)) {
            loopBody.push_back(block);
        }
    }
}


set<SAPFOR::Argument*> findRegisterSourceVariables(const std::vector<SAPFOR::BasicBlock*>& blocks, SAPFOR::Argument* var)
{
    set<SAPFOR::Argument*> result;
    set<SAPFOR::Argument*> visited;
    std::stack<SAPFOR::Argument*> workStack;
    workStack.push(var);

    auto isBinaryOp = [](SAPFOR::CFG_OP op) {
        return op == SAPFOR::CFG_OP::ADD || op == SAPFOR::CFG_OP::SUBT ||
            op == SAPFOR::CFG_OP::MULT || op == SAPFOR::CFG_OP::DIV ||
            op == SAPFOR::CFG_OP::POW ||
            op == SAPFOR::CFG_OP::GE || op == SAPFOR::CFG_OP::LE ||
            op == SAPFOR::CFG_OP::GT || op == SAPFOR::CFG_OP::LT ||
            op == SAPFOR::CFG_OP::EQ || op == SAPFOR::CFG_OP::NEQV ||
            op == SAPFOR::CFG_OP::EQV || op == SAPFOR::CFG_OP::EMPTY ||
            op == SAPFOR::CFG_OP::OR || op == SAPFOR::CFG_OP::AND;
        };

    auto isUnaryOp = [](SAPFOR::CFG_OP op) {
        return op == SAPFOR::CFG_OP::UN_ADD || op == SAPFOR::CFG_OP::UN_MINUS ||
            op == SAPFOR::CFG_OP::NOT || op == SAPFOR::CFG_OP::ASSIGN;
        };

    while (!workStack.empty()) {
        auto variable = workStack.top();
        workStack.pop();
        if (!variable || visited.count(variable))
            continue;

        visited.insert(variable);

        for (auto block : blocks) {
            for (auto instrWrapper : block->getInstructions()) {
                auto instr = instrWrapper->getInstruction();
                if (!instr || instr->getResult() != variable)
                    continue;

                auto op = instr->getOperation();
                auto arg1 = instr->getArg1();
                auto arg2 = instr->getArg2();

                if (isBinaryOp(op) && arg1 && arg2) {
                    if (arg1->getType() == SAPFOR::CFG_ARG_TYPE::VAR)
                        result.insert(arg1);
                    else if (arg1->getType() == SAPFOR::CFG_ARG_TYPE::REG)
                        workStack.push(arg1);

                    if (arg2->getType() == SAPFOR::CFG_ARG_TYPE::VAR)
                        result.insert(arg2);
                    else if (arg2->getType() == SAPFOR::CFG_ARG_TYPE::REG)
                        workStack.push(arg2);
                }
                else if (isUnaryOp(op) && arg1) {
                    if (arg1->getType() == SAPFOR::CFG_ARG_TYPE::VAR)
                        result.insert(arg1);
                    else if (arg1->getType() == SAPFOR::CFG_ARG_TYPE::REG)
                        workStack.push(arg1);
                }
            }
        }
    }

    return result;
}

std::vector<SAPFOR::Instruction*> getPhiArguments(SAPFOR::BasicBlock* block, SAPFOR::Instruction* phiInstr) {
    std::vector<SAPFOR::Instruction*> result;

    auto& instructions = block->getInstructions();
    bool collecting = false;
    for (int i = instructions.size() - 1; i >= 0; --i) {
        auto instr = instructions[i]->getInstruction();

        if (collecting) {
            if (instr->getOperation() == SAPFOR::CFG_OP::PARAM) {
                auto arg = instr->getArg1();
                if (arg) {
                    result.push_back(instr);
                }
            }
            else {
                break;
            }
        }

        if (!instr) continue;

        if (instr == phiInstr) {
            collecting = true;
            continue;
        }
    }

    std::reverse(result.begin(), result.end());
    return result;
}

SAPFOR::BasicBlock* findInstructionBlock(SAPFOR::Instruction* targetInstr, const std::vector<SAPFOR::BasicBlock*>& blocks) {
    for (auto block : blocks) {
        for (auto instrWrapper : block->getInstructions()) {
            auto instr = instrWrapper->getInstruction();
            if (instr == targetInstr) {
                return block;
            }
        }
    }
    return nullptr;
}

SAPFOR::BasicBlock* findInstructionBlockByNumber(int number, const std::vector<SAPFOR::BasicBlock*>& blocks) {
    for (auto block : blocks) {
        for (auto instrWrapper : block->getInstructions()) {
            auto instr = instrWrapper->getInstruction();
            if (instr->getNumber() == number) {
                return block;
            }
        }
    }
    return nullptr;
}


void findInductiveVars(const std::vector<SAPFOR::BasicBlock*>& blocks, const std::vector<SAPFOR::BasicBlock*>& Loopblocks, SAPFOR::BasicBlock* loopHeader, const set<SAPFOR::BasicBlock*>& loopExits) {
    set<string> inductiveVars;
    set<SAPFOR::BasicBlock*> relevantBlocks = { loopHeader };

    for (auto block : relevantBlocks) {

        for (auto instrWrapper : block->getInstructions()) {
            auto instr = instrWrapper->getInstruction();
            if (!instr) continue;

            auto op = instr->getOperation();
            auto res = instr->getResult();
            auto arg1 = instr->getArg1();
            auto arg2 = instr->getArg2();

            if (op == SAPFOR::CFG_OP::JUMP_IF) {
                if (arg1 && arg1->getType() == SAPFOR::CFG_ARG_TYPE::VAR) {
                    inductiveVars.insert(arg1->getValue());
                }
                if (arg1 && arg1->getType() == SAPFOR::CFG_ARG_TYPE::REG) {
                    auto foundVariables = findRegisterSourceVariables(blocks, arg1);
                    for (auto var : foundVariables) {
                        inductiveVars.insert(var->getValue());
                    }
                }
            }
        }
    }

    set<string> finalInductiveVars;

    for (auto instrWrapper : loopHeader->getInstructions()) {
        auto instr = instrWrapper->getInstruction();
        if (!instr || instr->getOperation() != SAPFOR::CFG_OP::F_CALL || !instr->getArg1() || instr->getArg1()->getValue() != "FI_FUNCTION") continue;

        auto phiRes = instr->getResult();
        if (!phiRes || !inductiveVars.count(phiRes->getValue())) continue;

        auto currentBlock = findInstructionBlock(instr, blocks);
        if (!currentBlock) continue;

        auto phiArgs = getPhiArguments(currentBlock, instr);

        bool hasInLoopDefinition = false;

        for (const auto& argInstr : phiArgs) {
            if (!argInstr) continue;

            int definitionInstrNumber = stoi(argInstr->getArg1()->getValue());
            if (definitionInstrNumber == -1) continue;

            auto phiBlock = findInstructionBlockByNumber(definitionInstrNumber, blocks);
            if (!phiBlock) continue;

            if (std::find(Loopblocks.begin(), Loopblocks.end(), phiBlock) != Loopblocks.end()) {
                hasInLoopDefinition = true;
            }
        }

        if (hasInLoopDefinition) {
            finalInductiveVars.insert(phiRes->getValue());
        }
    }

    for (auto i : finalInductiveVars) {
        std::cout << "Confirmed inductive variable: " << i << std::endl;
    }

    if (finalInductiveVars.empty()) {
        std::cout << "No confirmed inductive variables found." << std::endl;
    }
}

SAPFOR::Instruction* findInstructionAfterLoop(const std::vector<SAPFOR::BasicBlock*>& loopBody) {
    set<SAPFOR::BasicBlock*> loopSet(loopBody.begin(), loopBody.end());

    for (auto block : loopBody) {
        for (auto succ : block->getNext()) {
            if (!loopSet.count(succ)) {
                // Нашли выход из цикла — возьмём первую инструкцию
                auto instructions = succ->getInstructions();
                for (auto wrapper : instructions) {
                    if (auto instr = wrapper->getInstruction()) {
                        return instr;
                    }
                }
            }
        }
    }

    return nullptr; // не нашли
}

bool isEqual(const char* cstr, const std::string& str) {
    return str == cstr;
}

void findImplicitLoops(const map<FuncInfo*, vector<SAPFOR::BasicBlock*>>& fullIR_SSA, const char* fileName)
{
    for (auto& i : fullIR_SSA)
    {
        //for (auto j : i.second)
        //    printblock(j);

        if (!isEqual(fileName, i.first->fileName))
            continue;

        map<int, int> visited;
        for (auto i : i.second)
            visited[i->getNumber()] = UNVISITED;

        //vector<int> visited(i.second.size(), UNVISITED);
        vector<pair<SAPFOR::BasicBlock*, SAPFOR::BasicBlock*>> startAndEnd;
        dfs(i.second[0], visited, startAndEnd, NULL);

        vector<LoopGraph*> loops;

        for (auto& [tail, header] : startAndEnd) {
            set<SAPFOR::BasicBlock*> loopExits;

            for (auto succ : tail->getNext()) {
                if (succ != header) {
                    loopExits.insert(succ);
                }
            }

            vector<SAPFOR::BasicBlock*> loopBody;
            getLoopBody(header, loopExits, loopBody);

            cout << "LOOP DETECTED:" << endl;
            cout << "  Header: " << header->getNumber() << endl;
            cout << "  Tail: " << tail->getNumber() << endl;
            cout << "  Body blocks: ";
            for (auto block : loopBody) {
                cout << block->getNumber() << " ";
            }
            cout << endl;

            findInductiveVars(i.second, loopBody, header, loopExits);


            SAPFOR::Instruction* instructionAfterLoop = findInstructionAfterLoop(loopBody);

            if (instructionAfterLoop == NULL) {
                cout << "Warning: instruction after loop not found!" << endl;
                continue;
            }

            auto firstInstruction = header->getInstructions()[0]->getInstruction();
            auto lastInstruction = tail->getInstructions().back()->getInstruction();

            cout << "first - " << firstInstruction->getNumber() << " last - " << lastInstruction->getNumber() << " after - " << instructionAfterLoop->getNumber() << endl;
            auto x = firstInstruction->getOperator();

            auto tmpLoop = new LoopGraph();
            tmpLoop->isFor = true;
            tmpLoop->lineNum = firstInstruction->getOperator()->lineNumber();
            tmpLoop->lineNumAfterLoop = instructionAfterLoop->getOperator()->lineNumber();

            if (firstInstruction->getOperator()->variant() == FOR_NODE) {
                SgForStmt* stmt = isSgForStmt(firstInstruction->getOperator());

                cout << "for loop" << endl;// << stmt->sunparse() << endl;
            }
            else if (firstInstruction->getOperator()->variant() == WHILE_NODE) {
                SgWhileStmt* stmt = isSgWhileStmt(firstInstruction->getOperator());

                cout << (stmt->conditional() == NULL ? "infinit" : "") << "while loop" << endl;//<< stmt->sunparse() << endl;
            }
            else if (firstInstruction->getOperator()->variant() == DO_WHILE_NODE) {
                SgWhileStmt* stmt = isSgDoWhileStmt(firstInstruction->getOperator());

                cout << "do while loop" << endl;// << stmt->sunparse() << endl;
            }
            else if (firstInstruction->getOperator()->variant() == LOOP_NODE) {
                cout << "not known loop" << endl;// << firstInstruction->getOperator()->sunparse() << endl;
            }
            else {

                cout << "goto loop" << endl;// firstInstruction->getOperator()->sunparse() << endl;
            }

            cout << "loop start line " << tmpLoop->lineNum << endl;
            cout << "after loop line " << tmpLoop->lineNumAfterLoop << endl << endl;

            loops.push_back(tmpLoop);
        }
    }
}