#include "leak_detector.h"

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cstdint>

#include "dvm.h"
#include "directive_omp_parser.h"
#include "directive_parser.h"

#include "SgUtils.h"

using std::vector;
using std::map;
using std::set;
using std::string;


void removeOmpDir(SgStatement* st)
{
    char* lineS = st->comments();
    if (!lineS)
        return;

    vector<string> split;
    splitString(lineS, '\n', split);

    int idx = 0;
    for (auto& elem : split)
    {
        string line = elem;
        convertToLower(line);
        if (line.substr(0, 5) == "!$omp")
            lineS[idx + 1] = '_';
        else if (line.substr(0, 3) == "!$ ")
            lineS[idx + 1] = '_';
        idx += line.size() + 1; // with '\n'
    }
}

static inline void addToAttribute(SgStatement* st, int var, vector<SgExpression*> list)
{
    if (list.size())
    {
        SgExprListExp* ex = new SgExprListExp();
        ex->setLhs(new SgExpression(var, makeExprList(list), NULL));
        SgStatement* toAdd = new SgStatement(SPF_ANALYSIS_DIR, NULL, NULL, ex, NULL, NULL);
        toAdd->setlineNumber(st->lineNumber());
        toAdd->setLocalLineNumber(SPF_OMP_DIR);

        //filter 
        if (var == ACC_PRIVATE_OP)
        {
            vector<SgExpression*> list_new;

            auto attributes = getAttributes<SgStatement*, SgStatement*>(st, set<int>{SPF_ANALYSIS_DIR});
            set<string> privates;
            for (auto& attr : attributes)
                fillPrivatesFromComment(new Statement(attr), privates);

            if (privates.size())
            {
                for (auto& elem : list)
                    if (privates.find(elem->unparse()) == privates.end())
                        list_new.push_back(elem);
                list = list_new;

                if (!list.size())
                {
                    __spf_print(1, "-- skip privates on line %d from OMP dir\n%s", st->lineNumber(), toAdd->unparse());
                    return;
                }
            }
        }
        else if (var == REDUCTION_OP)
        {
            auto attributes = getAttributes<SgStatement*, SgStatement*>(st, set<int>{SPF_ANALYSIS_DIR});
            map<string, set<string>> reduction;
            for (auto& attr : attributes)
                fillReductionsFromComment(new Statement(attr), reduction);

            map<string, set<string>> reductionToAdd;
            fillReductionsFromComment(new Statement(toAdd), reductionToAdd);

            vector<SgExpression*> list_new;
            if (reduction.size())
            {
                if (reduction == reductionToAdd)
                {
                    __spf_print(1, "-- skip reduction on line %d from OMP dir\n%s", st->lineNumber(), toAdd->unparse());
                    return;
                }

                map<string, set<string>> reductionToAddNew;
                for (auto& redPair : reductionToAdd)
                {
                    auto it = reduction.find(redPair.first);
                    if (it == reduction.end())
                        reductionToAddNew[redPair.first] = redPair.second;
                    else
                    {
                        set<string> newVar;
                        for (auto& var : redPair.second)
                        {
                            auto itVar = it->second.find(var);
                            if (itVar == it->second.end())
                                reductionToAddNew[redPair.first].insert(var);
                        }
                    }
                }

                if (!reductionToAddNew.size())
                {
                    __spf_print(1, "-- skip reduction on line %d from OMP dir\n%s", st->lineNumber(), toAdd->unparse());
                    return;
                }

                if (reductionToAddNew != reductionToAdd)
                {
                    list.clear();
                    for (auto& redPair : reductionToAddNew)
                        for (auto& var : redPair.second)
                            list.push_back(new SgExpression(ARRAY_OP,
                                new SgKeywordValExp(redPair.first.c_str()),
                                new SgVarRefExp(findSymbolOrCreate(current_file, var, NULL, getFuncStat(st)))));
                }
            }
        }

        ex = new SgExprListExp();
        ex->setLhs(new SgExpression(var, makeExprList(list), NULL));
        toAdd = new SgStatement(SPF_ANALYSIS_DIR, NULL, NULL, ex, NULL, NULL);

        st->addAttribute(SPF_ANALYSIS_DIR, toAdd, sizeof(SgStatement));

        if (var == ACC_PRIVATE_OP)
            __spf_print(1, "-- set private attribute to line %d from OMP dir\n%s", st->lineNumber(), toAdd->unparse());
        else if (var == REDUCTION_OP)
            __spf_print(1, "-- set reduction attribute to line %d from OMP dir\n%s", st->lineNumber(), toAdd->unparse());
    }
}

static bool is_write_in_do(SgStatement* st, const string& var)
{
    checkNull(st, convertFileName(__FILE__).c_str(), __LINE__);

    if (st->variant() != FOR_NODE)
        return false;

    SgStatement* lastNode = st->lastNodeOfStmt();
    for (SgStatement* op = st->lexNext(); st != lastNode; st = st->lexNext())
    {
        if (st->variant() == ASSIGN_STAT)
        {
            SgExpression* ex = st->expr(0);
            if (ex->variant() == ARRAY_REF || ex->variant() == VAR_REF)
                if (var == ex->symbol()->identifier())
                    return true;
        }
        else if (st->variant() == FOR_NODE)
        {
            if (var == isSgForStmt(st)->doName()->identifier())
                return true;
        }
    }
    return false;
}

vector<OmpDir> parseOmpInStatement(SgStatement* st, const set<string>& globalPriv, bool forDo)
{
    vector<OmpDir> resultAll;

    const char* lineS = st->comments();
    if (!lineS)
        return resultAll;

    string comment(lineS);
    convertToLower(comment);

    vector<string> split;
    splitString(comment, '\n', split);

    for (int z = split.size() - 1; z >= 0; z--)
    {
        string line = split[z];

        if (line.substr(0, 6) == "!$omp&")
        {
            if (z - 1 < 0)
                break;
            split[z - 1] += line.substr(6);
            split[z] = "";
        }
    }

    for (auto& line : split)
    {
        if (line.substr(0, 5) == "!$omp")
        {
            OmpDir result;

            string line1 = "";
            int space = 0;
            int brake = 0;
            for (int z = 0; z < line.size(); ++z)
            {
                if (brake < 0)
                    return vector<OmpDir>(); // error

                if (brake == 0)
                {
                    if (line[z] == ' ')
                        space++;
                    else
                        space = 0;
                    if ((line[z] == ' ' && space <= 1) || line[z] != ' ')
                        line1 += line[z];
                }
                else
                {
                    if (line[z] != ' ')
                        line1 += line[z];
                }

                if (line[z] == '(')
                {
                    while (line1.size() > 2 && line1[line1.size() - 2] == ' ')
                        line1 = line1.erase(line1.size() - 2, 1);
                    brake++;
                    space = 0;
                }
                else if (line[z] == ')')
                    brake--;
            }
            vector<string> lexems;
            splitString(line1, ' ', lexems);
            bool doLexem = false;
            bool end = false;
            bool parallel = false;
            bool privat = false;

            for (auto& lexem : lexems)
            {
                if (lexem == "do")
                {
                    doLexem = true;
                    result.keys.insert(lexem);
                }
                if (lexem == "end")
                {
                    end = true;
                    result.keys.insert(lexem);
                }
                if (lexem == "parallel")
                {
                    parallel = true;
                    result.keys.insert(lexem);
                }
                if (lexem == "private")
                {
                    privat = true;
                    result.keys.insert(lexem);
                }
            }

            if (privat == false)
            {
                if (forDo && doLexem)
                {
                    vector<SgExpression*> list;
                    for (auto& var : globalPriv)
                        if (is_write_in_do(st, var))
                            list.push_back(new SgVarRefExp(findSymbolOrCreate(current_file, var, NULL, getFuncStat(st))));

                    if (list.size())
                        addToAttribute(st, ACC_PRIVATE_OP, list);
                }
            }

            for (auto& lexem : lexems)
            {
                bool priv = lexem.substr(0, strlen("private(")) == "private(";
                bool threadpriv = lexem.substr(0, strlen("threadprivate(")) == "threadprivate(";
                bool red = lexem.substr(0, strlen("reduction(")) == "reduction(";
                if (priv || threadpriv)
                {
                    vector<string> sublex;
                    splitString(lexem, '(', sublex);
                    if (sublex.size() == 2 && lexem.back() == ')')
                    {
                        splitString(sublex[1].erase(sublex[1].size() - 1), ',', sublex);

                        vector<SgExpression*> list;
                        set<string> uniqList;
                        for (auto& varG : globalPriv)
                            uniqList.insert(varG);

                        for (auto& var : sublex)
                            uniqList.insert(var);

                        for (auto& var : uniqList)
                        {
                            if (priv)
                            {
                                result.privVars.insert(var);
                                list.push_back(new SgVarRefExp(findSymbolOrCreate(current_file, var, NULL, getFuncStat(st))));
                            }
                            else
                                result.threadPrivVars.insert(var);
                        }

                        if (forDo && doLexem && priv)
                            addToAttribute(st, ACC_PRIVATE_OP, list);
                    }
                }
                else if (red)
                {
                    vector<string> sublex;
                    splitString(lexem, '(', sublex);
                    if (sublex.size() == 2 && lexem.back() == ')')
                    {
                        splitString(sublex[1].erase(sublex[1].size() - 1), ':', sublex);

                        vector<string> vars;
                        vector<SgExpression*> list;
                        splitString(sublex[1], ',', vars);
                        string op = "";

                        if (sublex[0] == "+")
                            op = "sum";
                        else if (sublex[0] == "*")
                            op = "prod";
                        else if (sublex[0] == "max")
                            op = "max";
                        else if (sublex[0] == "min")
                            op = "min";
                        else if (sublex[0] == ".or." || sublex[0] == "or")
                            op = "or";
                        else if (sublex[0] == ".and." || sublex[0] == "and")
                            op = "and";
                        else if (sublex[0] == ".eqv." || sublex[0] == "eqv")
                            op = "eqv";
                        else if (sublex[0] == ".neqv." || sublex[0] == "neqv")
                            op = "neqv";

                        if (op != "")
                        {
                            for (auto& var : vars)
                            {
                                result.redVars[sublex[0]].insert(var);
                                list.push_back(new SgExpression(ARRAY_OP, new SgKeywordValExp(op.c_str()), new SgVarRefExp(findSymbolOrCreate(current_file, var, NULL, getFuncStat(st)))));
                            }
                        }

                        if (forDo && doLexem && op != "")
                            addToAttribute(st, REDUCTION_OP, list);
                    }
                }
            }

            resultAll.push_back(result);
        }
    }

    return resultAll;
}

//TODO: need to use IR and RD for checking
static void filterPrivates(OmpDir& dir)
{
    if (dir.privVars.size() == 0)
        return;

    for (auto st = dir.start; st != dir.end; st = st->lexNext())
    {
        vector<OmpDir> res;
        if (st != dir.start)
        {
            set<string> dummy;
            res = parseOmpInStatement(st, dummy);
        }

        bool hasParallelDo = false;
        for (auto& dir : res)
        {
            if (dir.keys.find("parallel") != dir.keys.end() ||
                dir.keys.find("do") != dir.keys.end())
            {
                hasParallelDo = true;
            }
        }

        if (res.size() == 0 || !hasParallelDo)
        {
            if (st->variant() == ASSIGN_STAT)
            {
                if (st->expr(0))
                {
                    string ref = st->expr(0)->symbol()->identifier();
                    dir.privVars.erase(ref);
                }
            }
        }
    }
}

static vector<OmpDir> findAllGlobalParallelRegions(SgStatement* stFunc)
{
    vector<OmpDir> sections;

    SgStatement* lastNode = stFunc->lastNodeOfStmt();
    for (auto st = stFunc; st != lastNode; st = st->lexNext())
    {
        if (st == NULL)
        {
            __spf_print(1, "internal error in analysis, parallel directives will not be generated for this file!\n");
            break;
        }

        if (st->variant() == CONTAINS_STMT)
            break;

        set<string> dummy;
        auto res = parseOmpInStatement(st, dummy);

        for (auto& dir : res)
        {
            auto end = dir.keys.end();
            if (dir.keys.find("parallel") != end
                && dir.keys.find("do") == end
                && dir.keys.find("end") == end)
            {
                if (sections.size() && sections.back().end == NULL) // has open parallel region
                {
                    __spf_print(1, "wrong omp directives placed on line %d\n", st->lineNumber());
                    printInternalError(convertFileName(__FILE__).c_str(), __LINE__);
                }

                sections.push_back(dir);
                sections.back().start = st;
            }
            else if (dir.keys.find("parallel") != end
                     && dir.keys.find("do") == end
                     && dir.keys.find("end") != end)
            {
                if (!sections.size())
                {
                    __spf_print(1, "wrong omp directives placed on line %d\n", st->lineNumber());
                    printInternalError(convertFileName(__FILE__).c_str(), __LINE__);
                }
                sections.back().end = st;
            }
        }
    }

    for (auto& dir : sections)
        filterPrivates(dir);

    return sections;
}

static set<string> getGlobalPrivate(SgStatement* st, const vector<OmpDir>& globalParallelRegions)
{
    set<string> globalPrivates;
    const int line = st->lineNumber();
    if (line > 0)
    {
        for (auto& reg : globalParallelRegions)
        {
            if (reg.start->lineNumber() <= line && line < reg.end->lineNumber())
            {
                if (reg.privVars.size())
                    return reg.privVars;
                else
                    return globalPrivates;
            }
        }
    }
    else
    {
        for (auto& reg : globalParallelRegions)
        {
            for (auto stF = reg.start; stF != reg.end; stF = stF->lexNext())
            {
                if (st == stF)
                {
                    if (reg.privVars.size())
                        return reg.privVars;
                    else
                        return globalPrivates;
                }
            }
        }
    }

    return globalPrivates;
}

void parseOmpDirectives(SgFile* file, vector<Messages>& currMessages)
{
    int funcNum = file->numberOfFunctions();

    for (int i = 0; i < funcNum; ++i)
    {
        SgStatement* st = file->functions(i);
        SgStatement* lastNode = st->lastNodeOfStmt();

        vector<OmpDir> globalParallelRegions = findAllGlobalParallelRegions(st);

        while (st != lastNode)
        {
            if (st == NULL)
            {
                __spf_print(1, "internal error in analysis, parallel directives will not be generated for this file!\n");
                break;
            }

            if (st->variant() == CONTAINS_STMT)
                break;
           
            if (st->variant() == FOR_NODE)
            {
                SgForStmt* currSt = (SgForStmt*)st;
                if (currSt->isEnddoLoop() == 0)
                {
                    __spf_print(1, "wrong omp directives placed\n");
                    printInternalError(convertFileName(__FILE__).c_str(), __LINE__);
                }
                else
                    parseOmpInStatement(st, getGlobalPrivate(st, globalParallelRegions), true);
            }
            st = st->lexNext();
        }
    }
}