#include "leak_detector.h"

#include <string>
#include <map>

#include "Utils/SgUtils.h"
#include "set_implicit_none.h"

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

static const char commonIntLetters[6] = { 'i', 'j', 'k', 'm', 'n', 'l' };

static void InitTypes(map<char, SgType*>& types)
{
    for (char letter = 'a'; letter <= 'z'; letter++)
        types[letter] = 0;
}

static void FillCommonTypes(map<char, SgType*>& types)
{
    for (char letter : commonIntLetters)
        if (types[letter] == 0)
            types[letter] = new SgType(T_INT);

    for (auto letter : types)
        if (letter.second == NULL)
            types[letter.first] = new SgType(T_FLOAT);
}

static bool isByUse(SgSymbol* s, SgStatement* checkScope) 
{
    auto scope = s->scope();
    auto isByUse_s = IS_BY_USE(s);

    if (!isByUse_s && scope == checkScope)
        return false;
    else
        return true;
}

static void FindAllVars(SgExpression* expr, set<SgSymbol*>& allVars, set<SgSymbol*>& allVarsConst, SgStatement* scope)
{
    if (expr == NULL)
        return;

    const int var = expr->variant();
    if (var == VAR_REF || var == ARRAY_REF || var == FUNC_CALL)
    {
        auto s = expr->symbol(); 
        const string ident(s->identifier());
        const int s_var = s->variant();

        if (var == FUNC_CALL /*(s->attributes() & EXTERNAL_BIT)*/)
        {
            if (!IS_BY_USE(s) && ident.find("::") == string::npos /* && s->scope() == scope*/)
                allVars.insert(s);
        }
        else if (s_var != CONSTRUCT_NAME ||
                 s_var == VARIABLE_NAME)
        {
            if (!IS_BY_USE(s) && ident.find("::") == string::npos && s->scope() == scope)
                allVars.insert(s);
        }
    }
    else if (var == CONST_REF)
    {
        auto s = expr->symbol();
        if (!isByUse(s, scope))
            allVarsConst.insert(s);
    }

    FindAllVars(expr->lhs(), allVars, allVarsConst, scope);
    FindAllVars(expr->rhs(), allVars, allVarsConst, scope);
}

static char getValue(SgExpression* ex)
{
    char charVal = 0;
    if (ex && ex->variant() == CHAR_VAL)
        charVal = isSgValueExp(ex)->charValue();
    return charVal;
}

static void AddLettersToMap(SgExpression* expr, SgType* type, map<char, SgType*>& types)
{
    while (expr)
    {
        if (expr->variant() != EXPR_LIST)
            printInternalError(convertFileName(__FILE__).c_str(), __LINE__);

        SgExpression* val = expr->lhs();
        if (val->variant() == DDOT)
        {
            char leftVal = getValue(val->lhs());
            char rightVal = getValue(val->rhs());

            if (leftVal == 0 || rightVal == 0)
                printInternalError(convertFileName(__FILE__).c_str(), __LINE__);

            for (char letter = leftVal; letter <= rightVal; letter++)
                types[letter] = type;
        }
        else
        {
            char charVal = getValue(val);
            if (charVal == 0)
                printInternalError(convertFileName(__FILE__).c_str(), __LINE__);

            types[charVal] = type;
        }

        expr = expr->rhs();
    }
}

static vector<SgSymbol*> getVars(const set<string>& functionSymbs, set<SgSymbol*>& toRename,
                                 const set<SgSymbol*>& allVars, const map<char, SgType*>& types, 
                                 SgStatement* scope)
{
    vector<SgSymbol*> varsWithoutDecl;
    map<string, SgSymbol*> vars;
    for (auto& var : allVars)
    {
        if (vars.find(var->identifier()) != vars.end())
            printInternalError(convertFileName(__FILE__).c_str(), __LINE__);
        vars[var->identifier()] = var;
    }

    for (auto& var : allVars)
    {
        if (functionSymbs.count(var->identifier()))
            continue;

        vector<SgStatement*> allDecls;
        declaratedInStmt(var, &allDecls, false, (var->variant() == FUNCTION_NAME) ? scope : NULL);

        bool hasTypeDecls = false;
        for (auto& decl : allDecls)
        {
            const int var = decl->variant();
            if (var == VAR_DECL || var == VAR_DECL_90 || var == STRUCT_DECL)
                hasTypeDecls = true;
        }

        if (!hasTypeDecls)
        {
            const char c = var->identifier()[0];
            auto it = types.find(c);
            if (it != types.end())
            {
                var->setType(it->second);

                auto s = it->second->symbol();
                if (s)
                {
                    auto itS = vars.find(s->identifier());
                    if (itS != vars.end())
                        toRename.insert(itS->second);
                }
            }

            varsWithoutDecl.push_back(var);
        }
    }
    return varsWithoutDecl;
}

static map<char, SgType*> FunctionImplicitCheck(SgStatement* function, const map<SgStatement*, map<char, SgType*>>& typesByFunctions, 
                                                vector<int>& dvmDirErrorsLines)
{
    set<SgSymbol*> allVars, allVarsConst;
    map<char, SgType*> types;
    vector<SgSymbol*> varsWithoutDecl, varsWithoutDeclConst;
    set<SgSymbol*> toRename;
     
    InitTypes(types);
    FillCommonTypes(types);

    auto cp = function->controlParent();
    if (isSgProgHedrStmt(cp))
        if (typesByFunctions.find(cp) != typesByFunctions.end())
            for (auto& parentType : typesByFunctions.at(cp))
                types[parentType.first] = parentType.second;

    auto hasImplicitNone = false;
    auto endOfFunc = function->lastNodeOfStmt();
    for (auto st = function; st != endOfFunc; st = st->lexNext())
    {
        if (st->variant() == IMPL_DECL)
        {
            SgImplicitStmt* implicitStatement = isSgImplicitStmt(st);
            if (implicitStatement != NULL)
            {
                const int numberOfTypes = implicitStatement->numberOfImplicitTypes();

                if (numberOfTypes > 0)
                {
                    for (int j = 0; j < numberOfTypes; ++j)
                    {
                        SgType* type = implicitStatement->implicitType(j);
                        SgExpression* lettersExpression = implicitStatement->implicitRangeList(j);

                        AddLettersToMap(lettersExpression, type, types);
                    }
                }
                else
                    hasImplicitNone = true;
            }
        }
        else if (st->variant() == CONTAINS_STMT || isSgExecutableStatement(st) != NULL)
            break;
    }
    
    set<int> skip = { EXTERN_STAT, PRIVATE_STMT, PUBLIC_STMT };

    set<SgSymbol*> allDataSymbols;
    for (auto s = function->symbol()->next(); s; s = s->next())
        if ((s->attributes() & DATA_BIT) && s->scope() == function)
            allDataSymbols.insert(s);

    set<string> functionSymbs = { function->symbol()->identifier() };
    if (isSgFuncHedrStmt(function))
    {
        SgFuncHedrStmt* hedr = isSgFuncHedrStmt(function);
        if (hedr->resultName())
            functionSymbs.insert(hedr->resultName()->identifier());
    }

    auto prog = isSgProgHedrStmt(function);
    if (prog)
    {
        for (int z = 0; z < prog->numberOfInternalSubroutinesDefined(); ++z)
            functionSymbs.insert(prog->internalSubroutine(z)->symbol()->identifier());
        for (int z = 0; z < prog->numberOfInternalFunctionsDefined(); ++z)
            functionSymbs.insert(prog->internalFunction(z)->symbol()->identifier());
    }

    for (auto st = function->lexNext(); st != endOfFunc && st->variant() != CONTAINS_STMT; st = st->lexNext())
    {
        if (skip.count(st->variant()))
            continue;

        if (isDVM_stat(st))
            dvmDirErrorsLines.push_back(st->lineNumber());

        if (isDVM_stat(st) || isSPF_stat(st))
            continue;

        if (st->variant() == INTERFACE_STMT)
        {
            st = st->lastNodeOfStmt();
            continue;
        }

        if (st->variant() == DATA_DECL)
        {
            const string str = st->expr(0)->thellnd->entry.string_val;
            
            for (auto& data : allDataSymbols)
            {
                if (str.find(data->identifier()) != string::npos)
                    allVars.insert(data);
            }
        }

        for (int i = 0; i < 3; ++i)
            FindAllVars(st->expr(i), allVars, allVarsConst, function);

        if (st->variant() == FOR_NODE)
            if (!isByUse(isSgForStmt(st)->doName(), function))
                allVars.insert(isSgForStmt(st)->doName());
    }
    
    //add parameters
    if (prog)
    {
        for (int z = 0; z < prog->numberOfParameters(); ++z)
        {
            auto s = prog->parameter(z);
            if ((s->attributes() & EXTERNAL_BIT) == 0)
                allVars.insert(s);
        }
    }

    varsWithoutDecl = getVars(functionSymbs, toRename, allVars, types, function);
    varsWithoutDeclConst = getVars(functionSymbs, toRename, allVarsConst, types, NULL);
    
    if (!hasImplicitNone)
    {
        vector<SgStatement*> macro;

        for (auto st = function->lexNext(); 
             st != endOfFunc && st->variant() != CONTAINS_STMT && isSgExecutableStatement(st) == NULL; 
             )
        {
            if (st->variant() == IMPL_DECL)
            {
                auto tmpStatement = st;
                st = st->lexNext();
                tmpStatement->deleteStmt();
            }
            else if (st->variant() == STMTFN_STAT)
            {
                auto stat = st;
                st = st->lexNext();
                macro.push_back(stat->extractStmt());
            }
            else
                st = st->lexNext();
        }

        auto implNone = new SgStatement(IMPL_DECL);
        implNone->setlineNumber(function->lineNumber());
        implNone->setFileName(function->fileName());

        SgStatement* insertPlace = function;
        while (insertPlace->lexNext()->variant() == USE_STMT)
            insertPlace = insertPlace->lexNext();

        insertPlace->insertStmtAfter(*implNone, *function);
        insertPlace = insertPlace->lexNext();

        if (function->variant() == FUNC_HEDR)
        {
            SgFuncHedrStmt* hedr = isSgFuncHedrStmt(function);

            auto type_op = function->expr(1);
            if (type_op == NULL)
            {
                if (hedr->resultName())
                    varsWithoutDecl.push_back(hedr->resultName());
                else
                    varsWithoutDecl.push_back(function->symbol());
            }
        }

        makeDeclaration(varsWithoutDecl, function);

        auto declList = makeDeclaration(varsWithoutDeclConst, NULL);
        for (auto& decl : declList)
            insertPlace->insertStmtAfter(*decl, *function);

        if (macro.size())
        {
            while (!isSgExecutableStatement(insertPlace) && insertPlace != NULL)
                insertPlace = insertPlace->lexNext();

            if (insertPlace == NULL)
                printInternalError(convertFileName(__FILE__).c_str(), __LINE__);

            for (auto& elem : macro)
                insertPlace->insertStmtBefore(*elem, *function);
        }
    }

    for (auto& s : toRename)
        s->changeName(TestAndCorrectName((string(s->identifier()) + "_").c_str()));
    
    return types;
}

void implicitCheck(SgFile* file, vector<int>& dvmDirErrorsLines)
{
    map<SgStatement*, map<char, SgType*>> typesByFunctions;

    vector<SgStatement*> modulesAndFunctions;
    getModulesAndFunctions(file, modulesAndFunctions);
    const string currFile = file->filename();

    for (int func = 0; func < modulesAndFunctions.size(); ++func)
    {
        SgStatement* function = modulesAndFunctions[func];
        if (function->fileName() != currFile)
            continue;
        typesByFunctions[function] = FunctionImplicitCheck(function, typesByFunctions, dvmDirErrorsLines);
    }

    typesByFunctions.clear();
}
