#include "leak_detector.h"

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cstdint>
#include <assert.h>
#include <set>

#include "dvm.h"
#include "SgUtils.h"
#include "errors.h"

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

void preprocess_allocates(SgFile *file)
{
    const string currfile = file->filename();
    int funcNum = file->numberOfFunctions();

    set<SgStatement*> allStats;
    for (auto st = file->firstStatement(); st != NULL; st = st->lexNext())
        allStats.insert(st);

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

        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;

            // save SgStatement PROC call to declaration attribute
            if (st->variant() == ALLOCATE_STMT || st->variant() == DEALLOCATE_STMT)
            {
                SgExpression *list = st->expr(0);
                set<SgStatement*> decls;
                while (list)
                {
                    if (list->lhs()->variant() == ARRAY_REF)
                        decls.insert(declaratedInStmt(OriginalSymbol(isSgArrayRefExp(list->lhs())->symbol())));
                    else if (list->lhs()->variant() == RECORD_REF)
                    {
                        SgExpression* field = list->lhs()->lhs();
                        if (field->variant() == ARRAY_REF)
                            decls.insert(declaratedInStmt(OriginalSymbol(field->symbol())));
                        else if (field->variant() == VAR_REF)
                            decls.insert(declaratedInStmt(OriginalSymbol(field->symbol())));
                        else
                            printf(""); // TODO
                    }

                    list = list->rhs();
                }

                set<SgStatement*> altDecls;
                for (auto& decl : decls)
                {
                    if (decl->fileName() != currfile)
                    {
                        const string fIncl = decl->fileName();
                        for (auto& stat : allStats)
                            if (stat->lineNumber() == decl->lineNumber() && stat->fileName() == fIncl && decl != stat)
                                altDecls.insert(stat);
                    }
                }

                for (auto& alt : altDecls)
                    decls.insert(alt);

                char buf[256];
                for (auto& decl : decls)
                {
                    decl->addAttribute(st->variant(), st, sizeof(SgStatement*));
                    sprintf(buf, "  attribute (%d %s) is added to declaration on line %d of %s file\n", 
                                  st->lineNumber(), st->fileName(), decl->lineNumber(), decl->fileName());
                    addToGlobalBufferAndPrint(buf);
                }

                if (decls.size() == 0)
                {
                    sprintf(buf, "ERROR at line %d %s: Can not find declarations for allocation on line %d\n", __LINE__, convertFileName(__FILE__).c_str(), st->lineNumber());
                    addToGlobalBufferAndPrint(buf);
                    throw(-1);
                }
            }
            st = st->lexNext();
        }
    }
}

//TODO: extend ALLOC/DEALLOC moving
//TODO: add support of different files
// only one level of call is supported
void moveAllocatesInterproc(const map<DIST::Array*, set<DIST::Array*>> &arrayLinksByFuncCalls)
{
    for (auto &array : arrayLinksByFuncCalls)
    {
        if (array.first->GetLocation().first == DIST::l_PARAMETER)
        {
            //check one level of calls
            bool ok = true;
            for (auto &elem : array.second)
                if (elem->GetLocation().first == DIST::l_PARAMETER)
                    ok = false;

            //move to all real arrays
            if (ok)
            {
                auto allPlaces = array.first->GetDeclInfo();
                for (auto &place : allPlaces)
                {
                    if (SgFile::switchToFile(place.first) != -1)
                    {
                        SgStatement *decl = SgStatement::getStatementByFileAndLine(place.first, place.second);
                        checkNull(decl, convertFileName(__FILE__).c_str(), __LINE__);

                        map<SgStatement*, set<SgStatement*>> needToAdd;
                        for (auto &data : getAttributes<SgStatement*, SgStatement*>(decl, set<int>{ ALLOCATE_STMT }))
                        {
                            for (auto &realArray : array.second)
                            {
                                auto allPlacesR = realArray->GetDeclInfo();
                                for (auto &placeR : allPlacesR)
                                {
                                    //the same file
                                    if (placeR.first == place.first)
                                    {
                                        SgStatement* declR = SgStatement::getStatementByFileAndLine(placeR.first, placeR.second);
                                        checkNull(declR, convertFileName(__FILE__).c_str(), __LINE__);

                                        needToAdd[declR].insert(data);
                                    }
                                }
                            }
                        }

                        for (auto& toAdd : needToAdd)
                        {
                            for (auto& list : toAdd.second)
                            {
                                char buf[256];
                                toAdd.first->addAttribute(list->variant(), list, sizeof(SgStatement*));
                                sprintf(buf, "  [INTERPROC] attribute is added to declaration on line %d\n", toAdd.first->lineNumber());
                                addToGlobalBufferAndPrint(buf);
                            }
                        }
                    }
                }
            }
        }
    }

    map<string, map<int, set<SgStatement*>>> attrsForModules; // mod -> [ line -> decls]

    //get module declarations
    for (int z = 0; z < CurrentProject->numberOfFiles(); ++z)
    {
        SgFile* file = &(CurrentProject->file(z));

        vector<SgStatement*> modules;
        findModulesInFile(file, modules);

        for (auto& mod : modules)
        {
            SgStatement* st = mod;
            SgStatement* end = mod->lastNodeOfStmt();
            string modName = mod->symbol()->identifier();

            while (st != end)
            {
                if (st->variant() == CONTAINS_STMT)
                    break;

                for (auto& data : getAttributes<SgStatement*, SgStatement*>(st, set<int>{ ALLOCATE_STMT, DEALLOCATE_STMT }))
                    attrsForModules[modName][st->lineNumber()].insert(data);
                st = st->lexNext();
            }
        }
    }

    //synchronize module declarations
    for (int z = 0; z < CurrentProject->numberOfFiles(); ++z)
    {
        SgFile* file = &(CurrentProject->file(z));

        vector<SgStatement*> modules;
        findModulesInFile(file, modules);

        for (auto& mod : modules)
        {
            SgStatement* st = mod;
            SgStatement* end = mod->lastNodeOfStmt();
            string modName = mod->symbol()->identifier();

            if (attrsForModules.find(modName) == attrsForModules.end())
                continue;

            auto& info = attrsForModules[modName];
            while (st != end)
            {
                if (st->variant() == CONTAINS_STMT)
                    break;

                if (info.find(st->lineNumber()) != info.end())
                {
                    const set<SgStatement*>& decls = info[st->lineNumber()];
                    
                    auto attrs = getAttributes<SgStatement*, SgStatement*>(st, set<int>{ ALLOCATE_STMT, DEALLOCATE_STMT });
                    const set<SgStatement*> inAttr(attrs.begin(), attrs.end());

                    for (auto& data : decls)
                    {
                        if (inAttr.find(data) == inAttr.end())
                        {
                            char buf[256];
                            st->addAttribute(data->variant(), data, sizeof(SgStatement*));

                            sprintf(buf, "  [MODULE SYNC] attribute (%d %s) is added to declaration on line %d of %s file\n",
                                    data->lineNumber(), data->fileName(), st->lineNumber(), st->fileName());
                            addToGlobalBufferAndPrint(buf);
                        }
                    }
                }
                st = st->lexNext();
            }
        }
    }
}