#include <map>
#include <set>
#include <vector>
#include <iostream>
#include <algorithm>
#include <cstring>
#include <cstdlib>
#include <cmath>
#include <climits>
#include <functional>
#include <cctype>

#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, WHILE_NODE, DO_WHILE_NODE};
set<int> control_tags = {IF_NODE, ELSEIF_NODE, DO_WHILE_NODE, WHILE_NODE, LOGIF_NODE};
set<int> control_end_tags = {CONTROL_END};


static bool isParentStmt(SgStatement* stmt, SgStatement* parent)
{
    for (; stmt; stmt = stmt->controlParent())
        if (stmt == parent)
            return true;
    return false;
}

static bool isStatementInFile(SgStatement* st, const string& fileName)
{
    if (!st)
        return false;

    const char* stmtFile = st->fileName();
    return stmtFile && fileName == stmtFile;
}

static vector<SgStatement*> findLoopStatementsInFunction(SgStatement* funcStmt, const string& fileName)
{
    vector<SgStatement*> loops;
    if (!funcStmt)
        return loops;

    SgStatement* lastNode = funcStmt->lastNodeOfStmt();
    for (SgStatement* st = funcStmt; st && st != lastNode; st = st->lexNext())
        if (loop_tags.count(st->variant()) && isStatementInFile(st, fileName))
            loops.push_back(st);

    return loops;
}

static SgStatement* firstSignificantOperatorInBlock(const SAPFOR::BasicBlock* bb, const string& fileName)
{
    if (!bb)
        return nullptr;

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

        SgStatement* st = ir->getInstruction()->getOperator();
        if (isStatementInFile(st, fileName))
            return st;
    }

    return nullptr;
}

static SgStatement* lastSignificantOperatorInBlock(const SAPFOR::BasicBlock* bb, const string& fileName)
{
    if (!bb)
        return nullptr;

    const auto& instructions = bb->getInstructions();
    for (auto it = instructions.rbegin(); it != instructions.rend(); ++it)
    {
        auto* ir = *it;
        if (!ir || !ir->getInstruction())
            continue;

        SgStatement* st = ir->getInstruction()->getOperator();
        if (isStatementInFile(st, fileName))
            return st;
    }

    return nullptr;
}

static bool isBasicBlockInAnyLoop(const SAPFOR::BasicBlock* bb, const vector<SgStatement*>& loops, const string& fileName)
{
    SgStatement* first = firstSignificantOperatorInBlock(bb, fileName);
    SgStatement* last = lastSignificantOperatorInBlock(bb, fileName);
    if (!first || !last)
        return false;



    for (auto* loop : loops)
        if (loop && isParentStmt(first, loop) && isParentStmt(last, loop))
            return true;

    return false;
}

static vector<SAPFOR::BasicBlock*> findBlocksInLoopsByFullIR(
    SgFile* file,
    const map<FuncInfo*, vector<SAPFOR::BasicBlock*>>& FullIR)
{
    vector<SAPFOR::BasicBlock*> result;
    if (!file || !file->filename())
        return result;

    const string fileName = file->filename();
    set<SAPFOR::BasicBlock*> uniq;

    for (const auto& func : FullIR)
    {
        FuncInfo* funcInfo = func.first;
        if (!funcInfo || !funcInfo->funcPointer)
            continue;

        if (funcInfo->fileName != fileName)
            continue;

        const vector<SgStatement*> loops = findLoopStatementsInFunction(funcInfo->funcPointer, fileName);
        if (loops.empty())
            continue;

        for (auto* bb : func.second)
            if (isBasicBlockInAnyLoop(bb, loops, fileName))
                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*, set<SgStatement*>> analyzeBasicBlockIntraDependencies(const SAPFOR::BasicBlock* bb)
{
    if (!bb)
        return {};

    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 normalizeSageString = [](const string& src) -> string
    {
        string normalized;
        bool prevSpace = false;

        for (char ch : src)
        {
            const bool isSpace = std::isspace(static_cast<unsigned char>(ch));
            if (isSpace)
            {
                if (!normalized.empty())
                    prevSpace = true;
                continue;
            }

            if (prevSpace && !normalized.empty())
                normalized += ' ';
            normalized += ch;
            prevSpace = false;
        }

        return normalized;
    };

    auto sageExprToString = [&](SgExpression* expr) -> string
    {
        if (!expr)
            return string();

        char* raw = expr->unparse();
        return normalizeSageString(raw ? string(raw) : string());
    };

    auto arrayElementKey = [&](SgArrayRefExp* arrayRef) -> string
    {
        if (!arrayRef)
            return string();

        string key;
        SgSymbol* sym = arrayRef->symbol() ? OriginalSymbol(arrayRef->symbol()) : nullptr;
        if (sym && sym->identifier())
            key = sym->identifier();
        else
            return sageExprToString(arrayRef);

        if (key.empty())
            return string();

        key += "(";
        for (int i = 0; i < arrayRef->numberOfSubscripts(); ++i)
        {
            if (i)
                key += ", ";
            key += sageExprToString(arrayRef->subscript(i));
        }
        key += ")";

        return key;
    };

    function<void(SgExpression*, set<string>&)> collectUsedKeysFromExpression =
        [&](SgExpression* expr, set<string>& usedKeys)
    {
        if (!expr)
            return;

        if (isSgValueExp(expr) || expr->variant() == CONST_REF)
            return;

        if (auto* arrayRef = isSgArrayRefExp(expr))
        {
            const string key = arrayElementKey(arrayRef);
            if (!key.empty())
                usedKeys.insert(key);

            for (int i = 0; i < arrayRef->numberOfSubscripts(); ++i)
                collectUsedKeysFromExpression(arrayRef->subscript(i), usedKeys);
            return;
        }

        if (auto* call = isSgFunctionCallExp(expr))
        {
            for (int i = 0; i < call->numberOfArgs(); ++i)
                collectUsedKeysFromExpression(call->arg(i), usedKeys);
            return;
        }

        if (expr->variant() == VAR_REF)
        {
            SgSymbol* sym = expr->symbol() ? OriginalSymbol(expr->symbol()) : nullptr;
            if (sym && sym->identifier())
                usedKeys.insert(sym->identifier());
            return;
        }

        collectUsedKeysFromExpression(expr->lhs(), usedKeys);
        collectUsedKeysFromExpression(expr->rhs(), usedKeys);
    };

    auto collectUsedKeysFromArraySubscripts = [&](SgArrayRefExp* arrayRef, set<string>& usedKeys)
    {
        if (!arrayRef)
            return;

        for (int i = 0; i < arrayRef->numberOfSubscripts(); ++i)
            collectUsedKeysFromExpression(arrayRef->subscript(i), usedKeys);
    };

    auto addOperatorDependencies = [&](SgStatement* stmt,
                                       const set<string>& usedKeys,
                                       const map<string, SgStatement*>& varDeclarations,
                                       map<SgStatement*, set<SgStatement*>>& operatorsDependencies)
    {
        if (!stmt)
            return;

        for (const string& key : usedKeys)
        {
            auto it = varDeclarations.find(key);
            if (it != varDeclarations.end() && it->second && it->second != stmt)
                operatorsDependencies[stmt].insert(it->second);
        }
    };

    auto declarationKeyFromLeftPart = [&](SgStatement* stmt) -> string
    {
        if (!stmt || stmt->variant() != ASSIGN_STAT)
            return string();

        SgExpression* lhs = stmt->expr(0);
        if (!lhs)
            return string();

        if (auto* arrayRef = isSgArrayRefExp(lhs))
            return arrayElementKey(arrayRef);

        if (lhs->variant() == VAR_REF)
        {
            SgSymbol* sym = lhs->symbol() ? OriginalSymbol(lhs->symbol()) : nullptr;
            if (sym && sym->identifier())
                return sym->identifier();
        }

        return string();
    };

    vector<SgStatement*> operatorsOrder;
    set<SgStatement*> seenOperators;

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

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

        if (seenOperators.insert(stmt).second)
            operatorsOrder.push_back(stmt);
    }

    map<string, SgStatement*> varDeclarations;
    map<SgStatement*, set<SgStatement*>> operatorsDependencies;

    for (SgStatement* stmt : operatorsOrder)
    {
        set<string> usedKeys;
        if (stmt)
        {
            collectUsedKeysFromExpression(stmt->expr(1), usedKeys);
            collectUsedKeysFromArraySubscripts(isSgArrayRefExp(stmt->expr(0)), usedKeys);
        }
        addOperatorDependencies(stmt, usedKeys, varDeclarations, operatorsDependencies);

        const string declarationKey = declarationKeyFromLeftPart(stmt);
        if (!declarationKey.empty())
            varDeclarations[declarationKey] = stmt;
    }

    return operatorsDependencies;
}

static bool reorderOperatorsInBasicBlockUsingDeps(
    SAPFOR::BasicBlock* bb,
    const map<SgStatement*, set<SgStatement*>>& operatorsDependencies)
{
    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;
    }

    map<SgStatement*, int> originalIndex;
    for (int i = 0; i < (int)ops.size(); ++i)
        originalIndex[ops[i]] = i;

    map<SgStatement*, vector<SgStatement*>> moveAfter;
    set<SgStatement*> moved;

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

        SgStatement* lastDep = nullptr;
        int lastDepIdx = -1;
        for (SgStatement* dep : itDeps->second)
        {
            auto itIdx = originalIndex.find(dep);
            if (itIdx != originalIndex.end() && itIdx->second > lastDepIdx)
            {
                lastDep = dep;
                lastDepIdx = itIdx->second;
            }
        }
        if (!lastDep || lastDep == st)
            continue;

        moveAfter[lastDep].push_back(st);
        moved.insert(st);
    }

    vector<SgStatement*> order;
    order.reserve(ops.size());
    set<SgStatement*> emitted;
    set<SgStatement*> active;
    bool invalidOrder = false;

    function<void(SgStatement*)> emitStatement = [&](SgStatement* st)
    {
        if (!st || invalidOrder || emitted.count(st))
            return;
        if (active.count(st))
        {
            invalidOrder = true;
            return;
        }

        active.insert(st);
        emitted.insert(st);
        order.push_back(st);

        auto itMovedAfter = moveAfter.find(st);
        if (itMovedAfter != moveAfter.end())
            for (SgStatement* dependent : itMovedAfter->second)
                emitStatement(dependent);

        active.erase(st);
    };

    for (SgStatement* st : ops)
        if (!moved.count(st))
            emitStatement(st);

    if (invalidOrder || order.size() != ops.size())
        return false;

    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 auto loopBlocks = findBlocksInLoopsByFullIR(file, FullIR);

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