#include "leak_detector.h"

#include <cstdio>
#include <cstring>
#include <fstream>
#include <iostream>
#include <cstdlib>
#include <set>
#include <string>
#include <vector>
#include <algorithm>
#include <tuple>

#include "CFGraph/CFGraph.h"
#include "CFGraph/IR.h"
#include "Distribution/Array.h"
#include "dvm.h"

#include "errors.h"
#include "SgUtils.h"
#include "graph_calls.h"
#include "graph_calls_func.h"

#include "projectParameters.h"

using std::set;
using std::map;
using std::string;
using std::vector;
using std::tuple;
using std::pair;
using std::make_tuple;
using std::find_if;

static map<FuncInfo*, vector<SAPFOR::Instruction*>> call_sites;

enum class MODE
{
    BEFORE,
    AFTER
};

static tuple<FuncInfo*, SAPFOR::Instruction*, SAPFOR::BasicBlock*>
stmtToIR(const map<FuncInfo*, vector<SAPFOR::BasicBlock*>>& CFGraph, SgStatement* stmt)
{
    SgStatement* cur = stmt;
    cur->switchToFile();
    while (cur->variant() != PROC_HEDR && cur->variant() != PROG_HEDR && cur->variant() != FUNC_HEDR)
        cur = cur->controlParent();

    string funcName = ((SgProcHedrStmt*)cur)->nameWithContains();

    int stmtID = stmt->id();
    for (const auto& [func, bblocks] : CFGraph)
    {
        if (func->funcName != funcName)
            continue;

        for (auto basicBlock : bblocks)
            for (auto ins : basicBlock->getInstructions())
                if (stmtID == ins->getInstruction()->getOperator()->id())
                    return make_tuple(func, ins->getInstruction(), basicBlock);
    }

    printInternalError(convertFileName(__FILE__).c_str(), __LINE__);
    return { NULL, NULL, NULL };
}

static tuple<FuncInfo*, SAPFOR::Instruction*, SAPFOR::BasicBlock*>
IRByNumber(const map<FuncInfo*, vector<SAPFOR::BasicBlock*>>& CFGraph, int num)
{
    if (num < 0)
        return { NULL, NULL, NULL };

    for (const auto& [func, bblocks] : CFGraph)
        for (auto byBB : bblocks)
            if (byBB->getInstructions().front()->getNumber() <= num && byBB->getInstructions().back()->getNumber() >= num)
                return make_tuple(func, getInstructionByNumber(byBB->getInstructions(), num), byBB);

    printInternalError(convertFileName(__FILE__).c_str(), __LINE__);
    return { NULL, NULL, NULL };
}

template<typename Iterator>
static void processArgument(set<SAPFOR::Argument*>& worklist,
                            SAPFOR::Argument* arg,
                            Iterator instr, Iterator first_instr)
{
    if (arg == NULL)
        return;

    if (arg->getType() == SAPFOR::CFG_ARG_TYPE::REG)
        extract_vars_from_reg(worklist, arg, instr, first_instr);
    else if (arg->getType() == SAPFOR::CFG_ARG_TYPE::VAR)
        worklist.insert(arg);
}

template<typename Iterator>
static void extract_vars_from_reg(set<SAPFOR::Argument*>& worklist,
                                  SAPFOR::Argument* reg,
                                  Iterator instr, Iterator first_instr)
{
    for (; instr >= first_instr; instr--)
    {
        if ((*instr)->getInstruction()->getResult() == reg)
        {
            processArgument(worklist, (*instr)->getInstruction()->getArg1(), instr, first_instr);
            processArgument(worklist, (*instr)->getInstruction()->getArg2(), instr, first_instr);
            return;
        }
    }
}

static void lookup_for_vars(set<tuple<SgStatement*, string, MODE>>& where_to_add,
                            set<SAPFOR::Argument*>& worklist,
                            SAPFOR::Instruction* instr,
                            SAPFOR::BasicBlock* bblock,
                            FuncInfo* cur_func,
                            const map<FuncInfo*, vector<SAPFOR::BasicBlock*>>& fullIR)
{
    while (bblock)
    {
        auto first_instr = bblock->getInstructions().begin();
        auto cur_instr = find_if(first_instr, bblock->getInstructions().end(), [instr](SAPFOR::IR_Block* i) {
            return i->getInstruction() == instr;
            });

        for (; cur_instr >= bblock->getInstructions().begin(); --cur_instr)
        {
            auto instr = (*cur_instr)->getInstruction();
            auto result_arg = instr->getResult();
            auto arg1 = instr->getArg1();
            auto arg2 = instr->getArg2();

            if (worklist.count(result_arg))
            {
                worklist.erase(result_arg);
                processArgument(worklist, arg1, cur_instr, first_instr);
                processArgument(worklist, arg2, cur_instr, first_instr);
            }

            if (instr->getOperation() == SAPFOR::CFG_OP::PARAM && worklist.count(arg1))
            {
                // skip to F_CALL
                auto f_call_instr = cur_instr;
                while ((*f_call_instr)->getInstruction()->getOperation() != SAPFOR::CFG_OP::F_CALL)
                    f_call_instr++;

                if ((*f_call_instr)->getInstruction()->getArg1()->getValue() == "_READ")
                {
                    auto stmt_before = (*f_call_instr)->getInstruction()->getOperator();
                    auto filename = stmt_before->fileName();
                    auto line = stmt_before->lineNumber();
                    auto var_name = arg1->getValue().substr(arg1->getValue().find('%') + 1);
                    __spf_print(1, "Please specify value of variable %s on line %d of file %s\n", arg1->getValue().c_str(), line, filename);
                    auto toAdd = make_tuple(stmt_before, var_name, MODE::AFTER);
                    where_to_add.insert(toAdd);
                    worklist.erase(arg1);
                }
            }
        }

        const auto& RD = bblock->getRD_In();
        map<SAPFOR::BasicBlock*, SAPFOR::Instruction*> group_by_block;

        set<SAPFOR::Argument*> to_erase;
        for (auto& arg : worklist)
        {
            if (RD.count(arg))
            {
                if (RD.at(arg).size() == 1 && *RD.at(arg).begin() == SAPFOR::CFG_VAL::UNINIT)
                    __spf_print(1, "variable %s has no definition\n", arg->getValue().c_str());
                else if (RD.at(arg).size() > 1)
                {
                    auto stmt_after = (*first_instr)->getInstruction()->getOperator();
                    auto filename = stmt_after->fileName();
                    auto line = stmt_after->lineNumber();
                    auto var_name = arg->getValue().substr(arg->getValue().find('%') + 1);
                    __spf_print(1, "variable %s has multiple reaching definitions, further analysis is impossible\n", arg->getValue().c_str());
                    __spf_print(1, "Please specify value of variable %s on line %d of file %s\n", arg->getValue().c_str(), line, filename);
                    auto toAdd = make_tuple(stmt_after, var_name, MODE::BEFORE);
                    where_to_add.insert(toAdd);
                    to_erase.insert(arg);
                }
                else
                {
                    auto instr_num = *RD.at(arg).begin();
                    auto [func, instr, bblock] = IRByNumber(fullIR, instr_num);
                    if (cur_func == func && (group_by_block[bblock] == NULL || group_by_block[bblock]->getNumber() < instr_num))
                        group_by_block[bblock] = instr;
                }
            }
        }
        for (const auto& arg : to_erase)
            worklist.erase(arg);

        while (bblock && group_by_block.find(bblock) == group_by_block.end())
            bblock = bblock->getDom();

        if (bblock)
            instr = group_by_block[bblock];
    }

    // other variables are from global scope
    const auto& RD = fullIR.at(cur_func).front()->getRD_In();
    for (auto& arg : worklist)
    {
        if (arg->isMemGlobal())
        {
            set<int> found_rd;
            if (RD.count(arg))
                found_rd = RD.at(arg);

            if (found_rd.size() == 0)
            {
                auto call_instr = call_sites[cur_func].size() ? call_sites[cur_func].front() : NULL;
                while (call_instr && found_rd.size() == 0)
                {
                    auto [call_func, _, call_bblock] = IRByNumber(fullIR, call_instr->getNumber());
                    if (call_bblock->getRD_Out().count(arg))
                        found_rd = call_bblock->getRD_Out().at(arg);


                    call_instr = call_sites[call_func].size() ? call_sites[call_func].front() : NULL;
                }
            }

            if (found_rd.size() == 1 && *found_rd.begin() == SAPFOR::CFG_VAL::UNINIT)
                __spf_print(1, "variable %s has no definition\n", arg->getValue().c_str());
            else if (found_rd.size() > 1)
            {
                auto first_instr = fullIR.at(cur_func).front()->getInstructions().begin();
                auto stmt_after = (*first_instr)->getInstruction()->getOperator();
                auto filename = stmt_after->fileName();
                auto line = stmt_after->lineNumber();
                auto var_name = arg->getValue().substr(arg->getValue().find('%') + 1);
                __spf_print(1, "variable %s has multiple reaching definitions, further analysis is impossible\n", arg->getValue().c_str());
                __spf_print(1, "Please specify value of variable %s on line %d of file %s\n", arg->getValue().c_str(), line, filename);
                auto toAdd = make_tuple(stmt_after, var_name, MODE::BEFORE);
                where_to_add.insert(toAdd);
            }
            else
            {
                auto instr_num = *found_rd.begin();
                auto [func, instr, bblock] = IRByNumber(fullIR, instr_num);
                set<SAPFOR::Argument*> new_worklist = { arg };

                lookup_for_vars(where_to_add, new_worklist, instr, bblock, func, fullIR);
            }
        }
    }


    for (const auto& call_instr : call_sites[cur_func])
    {
        set<SAPFOR::Argument*> new_worklist;
        auto params_num = cur_func->funcParams.countOfPars;

        auto [call_func, _, call_bblock] = IRByNumber(fullIR, call_instr->getNumber());
        auto first_instr = call_bblock->getInstructions().begin();
        auto cur_instr = find_if(first_instr, call_bblock->getInstructions().end(), [call_instr](SAPFOR::IR_Block* i) {
            return i->getInstruction() == call_instr;
            });

        for (auto& arg : worklist)
        {
            if (arg->getMemType() == SAPFOR::CFG_MEM_TYPE::FUNC_PARAM_)
            {
                auto param_num = stoi(arg->getValue().substr(arg->getValue().find('%', arg->getValue().find('%') + 1) + 1));
                auto param_instr = (cur_instr - (params_num - param_num));
                auto param_arg = (*param_instr)->getInstruction()->getArg1();
                processArgument(new_worklist, param_arg, param_instr, first_instr);
            }
        }
        lookup_for_vars(where_to_add, new_worklist, call_instr, call_bblock, call_func, fullIR);
    }

}

static void handle_single_allocate(set<tuple<SgStatement*, string, MODE>>& where_to_add,
                                   SgStatement* alloc_statement,
                                   const map<FuncInfo*, vector<SAPFOR::BasicBlock*>>& fullIR)
{
    auto [func, instr, bblock] = stmtToIR(fullIR, alloc_statement);

    auto first_instr = bblock->getInstructions().begin();
    auto cur_instr = find_if(first_instr, bblock->getInstructions().end(), [instr](SAPFOR::IR_Block* i) {
        return i->getInstruction() == instr;
        });
    auto alloc_instr = cur_instr;

    // skip to F_CALL _ALLOC n
    while ((*alloc_instr)->getInstruction()->getOperation() != SAPFOR::CFG_OP::F_CALL ||
        (*alloc_instr)->getInstruction()->getArg1()->getValue() != "_ALLOC")
        alloc_instr++;

    auto arrays_num = stoi((*alloc_instr)->getInstruction()->getArg2()->getValue());

    set<SAPFOR::Argument*> worklist;
    for (int i = 0; i < arrays_num; i++)
    {
        auto param_instr = --alloc_instr;
        auto param_reg = (*param_instr)->getInstruction()->getArg1();

        while ((*param_instr)->getInstruction()->getOperation() != SAPFOR::CFG_OP::LOAD ||
            (*param_instr)->getInstruction()->getResult() != param_reg)
            param_instr--;


        auto dimensions_num = stoi((*param_instr)->getInstruction()->getArg2()->getValue());

        for (int j = 0; j < dimensions_num; j++)
        {
            auto ref_instr = --param_instr;
            if ((*ref_instr)->getInstruction()->getOperation() == SAPFOR::CFG_OP::RANGE)
            {
                vector<SAPFOR::Argument*> range_args = { (*ref_instr)->getInstruction()->getArg1(),
                                                        (*ref_instr)->getInstruction()->getArg2(),
                                                        (*ref_instr)->getInstruction()->getResult() };
                for (auto& arg : range_args)
                    processArgument(worklist, arg, ref_instr, first_instr);

            }
            else
            {
                auto arg = (*ref_instr)->getInstruction()->getArg1();
                processArgument(worklist, arg, ref_instr, first_instr);
            }
        }
    }
    lookup_for_vars(where_to_add, worklist, instr, bblock, func, fullIR);
}


static void handle_single_loop(set<tuple<SgStatement*, string, MODE>>& where_to_add,
                               SgStatement* loop_stmt,
                               const map<FuncInfo*, vector<SAPFOR::BasicBlock*>>& fullIR)
{
    auto [func, instr, bblock] = stmtToIR(fullIR, loop_stmt);
    auto cur_instr = bblock->getInstructions().end() - 1;

    set<SAPFOR::Argument*> worklist;
    extract_vars_from_reg(worklist, (*cur_instr)->getInstruction()->getResult(), cur_instr, bblock->getInstructions().begin());

    lookup_for_vars(where_to_add, worklist, (*cur_instr)->getInstruction(), bblock, func, fullIR);
}

void findParameters(ResultSet& foundParameters,
                    map<FuncInfo*, vector<SAPFOR::BasicBlock*>>& fullIR,
                    const map<tuple<int, string, string>, pair<DIST::Array*, DIST::ArrayAccessInfo*>>& declaredArrays)
{
    set<tuple<SgStatement*, string, MODE>> where_to_add;

    map<string, FuncInfo*> name_to_func;
    for (const auto& [func, _] : fullIR)
        name_to_func[func->funcName] = func;

    for (auto& [func, bblocks] : fullIR)
    {
        for (const auto& block : bblocks)
        {
            for (const auto& ir_block : block->getInstructions())
            {
                auto instr = ir_block->getInstruction();
                if (instr->getOperation() == SAPFOR::CFG_OP::F_CALL)
                {
                    auto func_name = instr->getArg1()->getValue();
                    auto func_info = name_to_func.find(func_name);

                    if (func_info != name_to_func.end())
                        call_sites[func_info->second].push_back(instr);
                }
            }
        }
    }

    set<SgStatement*> alloc_statements;
    for (const auto& [func, bblocks] : fullIR)
    {
        for (const auto& block : bblocks)
        {
            for (auto instr = block->getInstructions().begin(); instr != block->getInstructions().end(); ++instr)
            {
                auto op = (*instr)->getInstruction()->getOperator();
                if (op && op->variant() == ALLOCATE_STMT)
                    alloc_statements.insert(op);
            }
        }
    }

    set<SgStatement*> for_statements;
    // Find all FOR statements in the program
    for (const auto& [func, bblocks] : fullIR)
    {
        for (const auto& block : bblocks)
        {
            for (auto instr = block->getInstructions().begin(); instr != block->getInstructions().end(); ++instr)
            {
                auto op = (*instr)->getInstruction()->getOperator();
                if (op && op->variant() == FOR_NODE)
                    for_statements.insert(op);
            }
        }
    }

    for (const auto& alloc_statement : alloc_statements)
        handle_single_allocate(where_to_add, alloc_statement, fullIR);

    for (const auto& stmt : for_statements)
        handle_single_loop(where_to_add, stmt, fullIR);

    for (const auto& [stmt_before, var_name, mode] : where_to_add)
    {
        stmt_before->switchToFile();

        SgVariableSymb* var_symb = new SgVariableSymb(var_name.c_str());
        SgVarRefExp* var = new SgVarRefExp(var_symb);
        SgExprListExp* ex = new SgExprListExp();
        auto assgn_op = new SgExpression(ASSGN_OP, var, NULL);
        ex->setLhs(assgn_op);

        SgExpression* parameter_op = new SgExpression(SPF_PARAMETER_OP, ex);
        auto dir_list = new SgExprListExp();
        dir_list->setLhs(parameter_op);
        SgStatement* toAdd = new SgStatement(SPF_ANALYSIS_DIR, NULL, NULL, dir_list, NULL, NULL);

        toAdd->setlineNumber(stmt_before->lineNumber());
        toAdd->setLocalLineNumber(stmt_before->lineNumber());
        toAdd->setFileId(stmt_before->getFileId());
        toAdd->setProject(stmt_before->getProject());

        //NOTE: only for debbuging, results will be transferred to the visualizer
        /*if (mode == MODE::AFTER)
            stmt_before->insertStmtAfter(*toAdd, *stmt_before->controlParent());
        else
            stmt_before->insertStmtBefore(*toAdd, *stmt_before->controlParent());*/

        foundParameters.insert(make_tuple(stmt_before->fileName(), stmt_before->lineNumber(), var_name));
    }
}