#include <map>
#include <set>
#include <vector>
#include <iostream>
#include <algorithm>
#include <cstring>
#include <cstdlib>

#include "../../Utils/errors.h"
#include "../../Utils/SgUtils.h"
#include "../../GraphCall/graph_calls.h"
#include "../../GraphCall/graph_calls_func.h"
#include "../../CFGraph/CFGraph.h"
#include "../../CFGraph/IR.h"
#include "../../GraphLoop/graph_loops.h"
#include "move_operators.h"

using namespace std;

set<int> loop_tags = {FOR_NODE};
set<int> control_tags = {IF_NODE, ELSEIF_NODE, DO_WHILE_NODE, WHILE_NODE, LOGIF_NODE};
set<int> control_end_tags = {CONTROL_END};


static vector<SAPFOR::IR_Block*> findInstructionsFromStatement(SgStatement* st, const vector<SAPFOR::BasicBlock*>& blocks)
{
    vector<SAPFOR::IR_Block*> result;
    if (!st)
        return result;

    const int stmtId = st->id();
    for (auto* bb : blocks)
    {
        if (!bb)
            continue;

        for (auto* ir : bb->getInstructions())
        {
            if (!ir || !ir->getInstruction())
                continue;

            SgStatement* op = ir->getInstruction()->getOperator();
            if (op && op->id() == stmtId)
                result.push_back(ir);
        }
    }

    sort(result.begin(), result.end(),
         [](const SAPFOR::IR_Block* a, const SAPFOR::IR_Block* b) { return a->getNumber() < b->getNumber(); });
    return result;
}

vector<SAPFOR::BasicBlock*> findFuncBlocksByFuncStatement(SgStatement *st, const map<FuncInfo*, vector<SAPFOR::BasicBlock*>>& FullIR) {
    vector<SAPFOR::BasicBlock*> result;
    if (!st)
        return result;

    Statement* forSt = (Statement*)st;
    const string stmtFile = st->fileName();
    const int stmtLine = st->lineNumber();

    for (auto& func: FullIR) {
        if (!func.first || !func.first->funcPointer)
            continue;

        const string funcFile = func.first->fileName;
        const int funcLine = func.first->funcPointer->lineNumber();

        // Important: select CFG blocks only for the same file and function header.
        if (funcFile == stmtFile && funcLine == stmtLine)
        {
            result = func.second;
            break;
        }
    }
    return result;
}

map<SgForStmt*, vector<SAPFOR::BasicBlock*>> findAndAnalyzeLoops(SgStatement *st, const vector<SAPFOR::BasicBlock*>& blocks) {
    map<SgForStmt*, vector<SAPFOR::BasicBlock*>> result;
    SgStatement *lastNode = st->lastNodeOfStmt();
    
    while (st && st != lastNode) {
        if (loop_tags.find(st -> variant()) != loop_tags.end()) {
            SgForStmt *forSt = (SgForStmt*)st;
            SgStatement *loopBody = forSt -> body();
            SgStatement *lastLoopNode = st->lastNodeOfStmt();
            set<int> blocks_nums;
            
            while (loopBody && loopBody != lastLoopNode) {
                vector<SAPFOR::IR_Block*> irBlocks = findInstructionsFromStatement(loopBody, blocks);
                if (!irBlocks.empty()) {
                    SAPFOR::IR_Block* IR = irBlocks.front();
                    if (IR && IR->getBasicBlock()) {
                        if (blocks_nums.find(IR -> getBasicBlock() -> getNumber()) == blocks_nums.end()) {
                            result[forSt].push_back(IR -> getBasicBlock());
                            blocks_nums.insert(IR -> getBasicBlock() -> getNumber());
                        }
                    }
                }
                loopBody = loopBody -> lexNext();
            }
            sort(result[forSt].begin(), result[forSt].end());
        }
        st = st -> lexNext();
    }
    return result;
}


vector<SAPFOR::BasicBlock*> findBlocksInLoopsByFullIR(
    SgStatement* funcStmt,
    const map<FuncInfo*, vector<SAPFOR::BasicBlock*>>& FullIR)
{
    vector<SAPFOR::BasicBlock*> result;
    if (!funcStmt)
        return result;

    const vector<SAPFOR::BasicBlock*> funcBlocks = findFuncBlocksByFuncStatement(funcStmt, FullIR);
    const auto loopsMapping = findAndAnalyzeLoops(funcStmt, funcBlocks);

    set<SAPFOR::BasicBlock*> uniq;
    for (const auto& kv : loopsMapping)
        for (auto* bb : kv.second)
            if (bb)
                uniq.insert(bb);

    result.assign(uniq.begin(), uniq.end());
    sort(result.begin(), result.end(),
         [](const SAPFOR::BasicBlock* a, const SAPFOR::BasicBlock* b)
         {
             if (!a || !b)
                 return a < b;
             return a->getNumber() < b->getNumber();
         });

    return result;
}

static map<SgStatement*, vector<SgStatement*>> analyzeBasicBlockIntraDependencies(const SAPFOR::BasicBlock* bb)
{
    map<SgStatement*, vector<SgStatement*>> result;
    if (!bb)
        return result;

    auto isCompoundStmt = [](SgStatement* st) -> bool
    {
        if (!st)
            return true;
        const int v = st->variant();
        return loop_tags.count(v) || control_tags.count(v) || control_end_tags.count(v);
    };

    auto isTrackable = [](const SAPFOR::Argument* a) -> bool
    {
        if (!a)
            return false;
        const auto t = a->getType();
        return t == SAPFOR::CFG_ARG_TYPE::VAR || t == SAPFOR::CFG_ARG_TYPE::REG;
    };

    auto argKey = [&](const SAPFOR::Argument* a) -> string
    {
        if (!a)
            return string();
        return to_string((int)a->getType()) + "#" + to_string((int)a->getMemType()) + "#" + a->getValue();
    };

    auto memKeyFromInstr = [&](const SAPFOR::Instruction* instr) -> string
    {
        if (!instr)
            return string();
        SgExpression* ex = instr->getExpression();
        if (!ex)
            return string();

        auto exprKey = [&](auto&& self, SgExpression* e) -> string
        {
            if (!e)
                return string("_");

            if (auto* ar = isSgArrayRefExp(e))
            {
                SgSymbol* sym = ar->symbol() ? OriginalSymbol(ar->symbol()) : nullptr;
                string key = string("A(") + (sym ? sym->identifier() : "?");
                const int n = ar->numberOfSubscripts();
                for (int i = 0; i < n; ++i)
                {
                    key += ",";
                    key += self(self, ar->subscript(i));
                }
                key += ")";
                return key;
            }

            if (e->variant() == VAR_REF || e->variant() == CONST_REF)
            {
                SgSymbol* sym = e->symbol() ? OriginalSymbol(e->symbol()) : nullptr;
                return string((e->variant() == VAR_REF) ? "V(" : "C(") + (sym ? sym->identifier() : "?") + ")";
            }

            if (auto* v = isSgValueExp(e))
            {
                if (e->variant() == INT_VAL)
                    return string("I(") + to_string(v->intValue()) + ")";
                if (e->variant() == BOOL_VAL)
                    return string("B(") + (v->boolValue() ? "1" : "0") + ")";
                if (e->variant() == CHAR_VAL)
                    return string("CH(") + string(1, v->charValue()) + ")";
                if (e->variant() == FLOAT_VAL)
                    return string("F(") + (v->floatValue() ? v->floatValue() : "") + ")";
                if (e->variant() == DOUBLE_VAL)
                    return string("D(") + (v->doubleValue() ? v->doubleValue() : "") + ")";
                if (e->variant() == STRING_VAL)
                    return string("S(") + (v->stringValue() ? v->stringValue() : "") + ")";
            }

            string key = string("N(") + to_string(e->variant());
            if (e->lhs())
                key += ",L=" + self(self, e->lhs());
            if (e->rhs())
                key += ",R=" + self(self, e->rhs());
            key += ")";
            return key;
        };

        return "MEMEX#" + exprKey(exprKey, ex);
    };

    auto isBarrier = [&](const SAPFOR::Instruction* instr) -> bool
    {
        if (!instr)
            return true;
        const auto op = instr->getOperation();
        switch (op)
        {
        case SAPFOR::CFG_OP::F_CALL:
        case SAPFOR::CFG_OP::IO_PARAM:
        case SAPFOR::CFG_OP::DVM_DIR:
        case SAPFOR::CFG_OP::SPF_DIR:
        case SAPFOR::CFG_OP::POINTER_ASS:
        case SAPFOR::CFG_OP::EXIT:
            return true;
        default:
            return false;
        }
    };

    auto isDef = [&](const SAPFOR::Instruction* instr) -> bool
    {
        if (!instr)
            return false;
        SAPFOR::Argument* r = instr->getResult();
        if (!isTrackable(r))
            return false;

        const auto op = instr->getOperation();
        if (op == SAPFOR::CFG_OP::STORE || op == SAPFOR::CFG_OP::REC_REF_STORE)
            return false;

        return true;
    };

    // Reaching definitions inside the BasicBlock in straight-line order:
    // lastDef[var] = last operator in this block that defined it.
    map<string, pair<SgStatement*, const SAPFOR::Argument*>> lastDef;
    map<string, pair<SgStatement*, const SAPFOR::Argument*>> lastMemDef;
    map<SgStatement*, set<SgStatement*>> depsSets;

    for (auto* ir : bb->getInstructions())
    {
        if (!ir || !ir->getInstruction())
            continue;

        const SAPFOR::Instruction* instr = ir->getInstruction();
        SgStatement* opStmt = instr->getOperator();
        if (!opStmt)
            continue;
        if (isCompoundStmt(opStmt))
            continue;

        if (isBarrier(instr))
        {
            for (auto it = lastDef.begin(); it != lastDef.end();)
            {
                const SAPFOR::Argument* a = it->second.second;
                if (!a || a->isMemGlobal() || a->isParameter())
                    it = lastDef.erase(it);
                else
                    ++it;
            }
            for (auto it = lastMemDef.begin(); it != lastMemDef.end();)
            {
                const SAPFOR::Argument* a = it->second.second;
                if (!a || a->isMemGlobal() || a->isParameter())
                    it = lastMemDef.erase(it);
                else
                    ++it;
            }
        }

        if (!result.count(opStmt))
            result[opStmt] = {};

        auto addDep = [&](SAPFOR::Argument* use)
        {
            if (!isTrackable(use))
                return;
            const string k = argKey(use);
            auto it = lastDef.find(k);
            if (it == lastDef.end())
                return;
            if (it->second.first && it->second.first != opStmt)
                depsSets[opStmt].insert(it->second.first);
        };

        auto addMemDep = [&](const string& key)
        {
            if (key.empty())
                return;
            auto it = lastMemDef.find(key);
            if (it == lastMemDef.end())
                return;
            if (it->second.first && it->second.first != opStmt)
                depsSets[opStmt].insert(it->second.first);
        };

        addDep(instr->getArg1());
        addDep(instr->getArg2());
        if (instr->getOperation() == SAPFOR::CFG_OP::RANGE)
            addDep(instr->getResult());

        if (instr->getOperation() == SAPFOR::CFG_OP::STORE || instr->getOperation() == SAPFOR::CFG_OP::REC_REF_STORE)
            addDep(instr->getResult());

        if (instr->getOperation() == SAPFOR::CFG_OP::LOAD || instr->getOperation() == SAPFOR::CFG_OP::REC_REF_LOAD)
        {
            const string memKey = memKeyFromInstr(instr);
            addMemDep(memKey);
        }

        if (isDef(instr))
        {
            const string dk = argKey(instr->getResult());
            lastDef[dk] = { opStmt, instr->getResult() };
        }

        if (instr->getOperation() == SAPFOR::CFG_OP::STORE || instr->getOperation() == SAPFOR::CFG_OP::REC_REF_STORE)
        {
            const string k = memKeyFromInstr(instr);
            SAPFOR::Argument* base = instr->getArg1();
            if (!k.empty() && base)
                lastMemDef[k] = { opStmt, base };
            addMemDep(k);
        }
    }

    for (auto& kv : result)
    {
        SgStatement* op = kv.first;
        auto it = depsSets.find(op);
        if (it == depsSets.end())
            continue;

        kv.second.assign(it->second.begin(), it->second.end());
        sort(kv.second.begin(), kv.second.end(),
             [](SgStatement* a, SgStatement* b)
             {
                 const int la = a ? a->lineNumber() : -1;
                 const int lb = b ? b->lineNumber() : -1;
                 if (la != lb)
                     return la < lb;
                 return a < b;
             });
    }

    return result;
}

static bool reorderOperatorsInBasicBlockUsingDeps(SAPFOR::BasicBlock* bb, const char* expectedFile)
{
    if (!bb)
        return false;

    auto isCompoundStmt = [](SgStatement* st) -> bool
    {
        if (!st)
            return true;
        const int v = st->variant();
        return loop_tags.count(v) || control_tags.count(v) || control_end_tags.count(v);
    };

    vector<SgStatement*> ops;
    ops.reserve(bb->getInstructions().size());
    set<SgStatement*> seen;
    for (auto* ir : bb->getInstructions())
    {
        if (!ir || !ir->getInstruction())
            continue;
        SgStatement* st = ir->getInstruction()->getOperator();
        if (!st || isCompoundStmt(st))
            continue;
        if (seen.insert(st).second)
            ops.push_back(st);
    }

    if (ops.size() < 2)
        return false;

    // Check that analyzed BB is in the same file as the expected file
    const char* bbFile = ops.front()->fileName();
    if (!bbFile)
        bbFile = "(unknown)";
    if (expectedFile && strcmp(expectedFile, bbFile) != 0)
        return false;
    for (auto* st : ops)
    {
        if (!st || !st->fileName() || strcmp(st->fileName(), bbFile) != 0)
            return false;
    }

    SgStatement* parent = ops.front()->controlParent();
    if (!parent)
        return false;
    for (auto* st : ops)
        if (!st || st->controlParent() != parent)
            return false;

    set<SgStatement*> opSet(ops.begin(), ops.end());
    {
        SgStatement* lastNode = parent->lastNodeOfStmt();
        SgStatement* cur = ops.front();
        size_t idx = 0;

        {
            SgStatement* scan = parent->lexNext();
            SgStatement* end = lastNode;
            bool found = false;
            set<SgStatement*> visited;
            while (scan && scan != end)
            {
                if (!visited.insert(scan).second)
                    return false;
                if (scan == ops.front())
                {
                    found = true;
                    break;
                }
                scan = scan->lexNext();
            }
            if (!found)
                return false;
        }

        while (cur && cur != lastNode && idx < ops.size())
        {
            if (cur == ops[idx])
                ++idx;
            else if (isSgExecutableStatement(cur) && !opSet.count(cur))
                return false;
            if (idx == ops.size())
                break;
            cur = cur->lexNext();
        }
        if (idx != ops.size())
            return false;
    }

    // Compute dependencies (inside BB) and build a new order by moving each statement
    // as close as possible after its last dependency (if any).
    const auto depsMap = analyzeBasicBlockIntraDependencies(bb);
    vector<SgStatement*> order = ops;

    auto indexIn = [](const vector<SgStatement*>& v, SgStatement* s) -> int
    {
        for (int i = 0; i < (int)v.size(); ++i)
            if (v[i] == s)
                return i;
        return -1;
    };

    for (SgStatement* s : ops)
    {
        auto itDeps = depsMap.find(s);
        if (itDeps == depsMap.end() || itDeps->second.empty())
            continue;

        int posS = indexIn(order, s);
        if (posS < 0)
            continue;

        int lastDepIdx = -1;
        for (SgStatement* dep : itDeps->second)
        {
            const int j = indexIn(order, dep);
            if (j >= 0)
                lastDepIdx = max(lastDepIdx, j);
        }
        if (lastDepIdx < 0)
            continue;

        if (posS == lastDepIdx + 1)
            continue;

        order.erase(order.begin() + posS);

        int lp = lastDepIdx;
        if (posS < lastDepIdx)
            lp = lastDepIdx - 1;

        const int insertAt = lp + 1;
        order.insert(order.begin() + insertAt, s);
    }

    bool changed = false;
    for (size_t i = 0; i < ops.size(); ++i)
        if (ops[i] != order[i])
        {
            changed = true;
            break;
        }
    if (!changed)
        return false;

    SgStatement* anchor = parent;
    {
        SgStatement* scan = parent->lexNext();
        SgStatement* end = parent->lastNodeOfStmt();
        while (scan && scan != end)
        {
            if (scan == ops.front())
                break;
            anchor = scan;
            scan = scan->lexNext();
        }
    }

    // apply AST reordering
    map<SgStatement*, int> savedLine;
    map<SgStatement*, char*> savedComments;
    map<SgStatement*, SgStatement*> extracted;

    for (SgStatement* st : ops)
    {
        if (!st)
            return false;

        savedLine[st] = st->lineNumber();
        savedComments[st] = st->comments() ? strdup(st->comments()) : nullptr;

        SgStatement* ex = st->extractStmt();
        if (!ex)
            return false;
        extracted[st] = ex;
    }

    SgStatement* insertAfter = anchor;
    for (SgStatement* st : order)
    {
        SgStatement* ex = extracted[st];
        if (!ex)
            continue;

        auto itC = savedComments.find(st);
        if (itC != savedComments.end() && itC->second)
            ex->setComments(itC->second);

        auto itL = savedLine.find(st);
        if (itL != savedLine.end())
            ex->setlineNumber(itL->second);

        insertAfter->insertStmtAfter(*ex, *parent);
        insertAfter = ex;
    }

    for (auto& kv : savedComments)
        if (kv.second)
            free(kv.second);

    return true;
}

void moveOperators(SgFile* file, const map<FuncInfo*, vector<SAPFOR::BasicBlock*>>& FullIR, int& countOfTransform) {
    if (!file)
        return;
    if (SgFile::switchToFile(file->filename()) == -1)
        printInternalError(convertFileName(__FILE__).c_str(), __LINE__);
    const int funcNum = file->numberOfFunctions();
    for (int i = 0; i < funcNum; ++i)
    {
        SgStatement* st = file->functions(i);

        const auto loopBlocks = findBlocksInLoopsByFullIR(st, FullIR);
        for (auto* bb : loopBlocks)
        {
            if (!bb)
                continue;
            if (reorderOperatorsInBasicBlockUsingDeps(bb, file->filename()))
                countOfTransform += 1;
        }
    }
}