#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;
    Statement* forSt = (Statement*)st;

    for (auto& func: FullIR) {
        if (func.first->funcPointer->getCurrProcessFile() == forSt->getCurrProcessFile()
            && func.first->funcPointer->lineNumber() == forSt->lineNumber())
        {
            result = func.second;
        }
    }
    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 || !ex->unparse())
            return string();
        auto normalizeExprText = [](const string& raw) -> string
        {
            string t;
            t.reserve(raw.size());
            for (unsigned char c : raw)
                if (!std::isspace(c))
                    t.push_back((char)c);

            auto stripOneLayer = [](const string& x) -> string
            {
                if (x.size() < 2 || x.front() != '(' || x.back() != ')')
                    return x;
                int bal = 0;
                for (size_t i = 0; i + 1 < x.size(); ++i)
                {
                    if (x[i] == '(') bal++;
                    else if (x[i] == ')') bal--;
                    if (bal == 0 && i + 1 < x.size() - 1)
                        return x;
                }
                return x.substr(1, x.size() - 2);
            };
            while (true)
            {
                const string stripped = stripOneLayer(t);
                if (stripped == t)
                    break;
                t = stripped;
            }
            return t;
        };
        return "MEMEX#" + normalizeExprText(ex->unparse());
    };

    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::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)
{
    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;

    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 buildPos = [&](const vector<SgStatement*>& v)
    {
        map<SgStatement*, int> pos;
        for (int i = 0; i < (int)v.size(); ++i)
            pos[v[i]] = i;
        return pos;
    };

    const int maxIterations = (int)order.size() * (int)order.size() + 10;
    bool anyMove = false;
    for (int iter = 0; iter < maxIterations; ++iter)
    {
        bool movedThisIter = false;
        const auto pos = buildPos(order);

        for (int i = 0; i < (int)order.size(); ++i)
        {
            SgStatement* curSt = order[i];
            auto it = depsMap.find(curSt);
            if (it == depsMap.end() || it->second.empty())
                continue;

            int lastDepIdx = -1;
            for (SgStatement* dep : it->second)
            {
                auto itP = pos.find(dep);
                if (itP != pos.end())
                    lastDepIdx = max(lastDepIdx, itP->second);
            }
            if (lastDepIdx < 0)
                continue;

            int target = lastDepIdx + 1;
            if (target == i)
                continue;

            SgStatement* moved = order[i];
            order.erase(order.begin() + i);
            if (target > i)
                target -= 1;
            if (target < 0)
                target = 0;
            if (target > (int)order.size())
                target = (int)order.size();
            order.insert(order.begin() + target, moved);

            movedThisIter = true;
            anyMove = true;
            break;
        }

        if (!movedThisIter)
            break;
    }

    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(const map<FuncInfo*, vector<SAPFOR::BasicBlock*>>& FullIR, int& countOfTransform) {
    const int funcNum = current_file->numberOfFunctions();
    for (int i = 0; i < funcNum; ++i)
    {
        SgStatement* st = current_file->functions(i);

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