#include <libasr/asr.h>
#include <libasr/containers.h>
#include <libasr/exception.h>
#include <libasr/asr_utils.h>
#include <libasr/asr_verify.h>
#include <libasr/pass/replace_intrinsic_function.h>
#include <libasr/pass/intrinsic_function_registry.h>
#include <libasr/pass/intrinsic_array_function_registry.h>
#include <libasr/pass/pass_utils.h>

#include <vector>


namespace LCompilers {

/*

This ASR pass replaces the IntrinsicFunction node with a call to an
implementation in ASR that we construct (and cache) on the fly for the actual
arguments.

Call this pass if you do not want to implement intrinsic functions directly
in the backend.

*/

class ReplaceIntrinsicFunctions: public ASR::BaseExprReplacer<ReplaceIntrinsicFunctions> {

    private:

    Allocator& al;
    SymbolTable* global_scope;
    std::map<ASR::symbol_t*, ASRUtils::IntrinsicArrayFunctions>& func2intrinsicid;
    bool& in_debugcheck;
    bool& in_ttype;
    int index_kind;

    public:

    ReplaceIntrinsicFunctions(Allocator& al_, SymbolTable* global_scope_,
    std::map<ASR::symbol_t*, ASRUtils::IntrinsicArrayFunctions>& func2intrinsicid_, bool& in_debugcheck_, bool &in_ttype_,
    int index_kind_) :
        al(al_), global_scope(global_scope_), func2intrinsicid(func2intrinsicid_), in_debugcheck(in_debugcheck_), in_ttype(in_ttype_),
        index_kind(index_kind_) {}


    void replace_IntrinsicElementalFunction(ASR::IntrinsicElementalFunction_t* x) {
        if (x->m_value) {
            *current_expr = x->m_value;
            return;
        }
        // If inside DebugCheckArrayBounds and not inside a type
        // Then keep the IntrinsicElementalFunction because its arguments are used to find ArraySize
        if (in_debugcheck && !in_ttype) return;

        Vec<ASR::call_arg_t> new_args; new_args.reserve(al, x->n_args);
        // Replace any IntrinsicElementalFunctions in the argument first:
        for( size_t i = 0; i < x->n_args; i++ ) {
            ASR::call_arg_t arg0;
            arg0.loc = (*current_expr)->base.loc;
            arg0.m_value = x->m_args[i];
            new_args.push_back(al, arg0);
        }
        // TODO: currently we always instantiate a new function.
        // Rather we should reuse the old instantiation if it has
        // exactly the same arguments. For that we could use the
        // overload_id, and uniquely encode the argument types.
        // We could maintain a mapping of type -> id and look it up.

        ASRUtils::impl_function instantiate_function =
            ASRUtils::IntrinsicElementalFunctionRegistry::get_instantiate_function(x->m_intrinsic_id);
        if( instantiate_function == nullptr ) {
            return ;
        }
        Vec<ASR::ttype_t*> arg_types;
        arg_types.reserve(al, x->n_args);
        for( size_t i = 0; i < x->n_args; i++ ) {
            arg_types.push_back(al, ASRUtils::expr_type(x->m_args[i]));
        }
        ASR::ttype_t* type = nullptr;
        type = ASRUtils::extract_type(x->m_type);
        ASR::expr_t* current_expr_ = instantiate_function(al, x->base.base.loc,
            global_scope, arg_types, type, new_args, x->m_overload_id, index_kind);
        if (current_expr_) {
            *current_expr = current_expr_;
        }
    }

    void replace_IntrinsicArrayFunction(ASR::IntrinsicArrayFunction_t* x) {
        std::string intrinsic_name_ = std::string(ASRUtils::get_array_intrinsic_name(x->m_arr_intrinsic_id));
        if (x->m_value) {
            *current_expr = x->m_value;
            return;
        }
        replace_ttype(x->m_type);
        Vec<ASR::call_arg_t> new_args; new_args.reserve(al, x->n_args);
        // Replace any IntrinsicArrayFunctions in the argument first:
        for( size_t i = 0; i < x->n_args; i++ ) {
            ASR::call_arg_t arg0;
            arg0.loc = (*current_expr)->base.loc;
            arg0.m_value = x->m_args[i];
            new_args.push_back(al, arg0);
        }

        // TODO: currently we always instantiate a new function.
        // Rather we should reuse the old instantiation if it has
        // exactly the same arguments. For that we could use the
        // overload_id, and uniquely encode the argument types.
        // We could maintain a mapping of type -> id and look it up.

        ASRUtils::impl_function instantiate_function =
            ASRUtils::IntrinsicArrayFunctionRegistry::get_instantiate_function(x->m_arr_intrinsic_id);
        if( instantiate_function == nullptr ) {
            return ;
        }
        Vec<ASR::ttype_t*> arg_types;
        arg_types.reserve(al, x->n_args);
        for( size_t i = 0; i < x->n_args; i++ ) {
            arg_types.push_back(al, ASRUtils::expr_type(x->m_args[i]));
        }
        ASR::expr_t* current_expr_ = instantiate_function(al, x->base.base.loc,
            global_scope, arg_types, x->m_type, new_args, x->m_overload_id, index_kind);
        ASR::expr_t* func_call = current_expr_;
        *current_expr = current_expr_;
        bool condition = ASR::is_a<ASR::FunctionCall_t>(*func_call);
        condition = condition && ASRUtils::is_array(x->m_type);
        if (condition) {
            ASR::symbol_t *call_sym = ASRUtils::symbol_get_past_external(
                ASR::down_cast<ASR::FunctionCall_t>(func_call)->m_name);
            func2intrinsicid[call_sym] = (ASRUtils::IntrinsicArrayFunctions) x->m_arr_intrinsic_id;
        }
    }
};

/*
The following visitor calls the above replacer i.e., ReplaceFunctionCalls
on expressions present in ASR so that FunctionCall get replaced everywhere
and we don't end up with false positives.
*/
class ReplaceIntrinsicFunctionsVisitor : public ASR::CallReplacerOnExpressionsVisitor<ReplaceIntrinsicFunctionsVisitor>
{
    private:

        ReplaceIntrinsicFunctions replacer;

    public:
        bool in_debugcheck = false;
        bool in_ttype = false;

        ReplaceIntrinsicFunctionsVisitor(Allocator& al_, SymbolTable* global_scope_,
            std::map<ASR::symbol_t*, ASRUtils::IntrinsicArrayFunctions>& func2intrinsicid_,
            int index_kind_) :
            replacer(al_, global_scope_, func2intrinsicid_, in_debugcheck, in_ttype, index_kind_) {}

        // Don't replace inside DebugCheckArrayBounds, the arguments for elemental functions might be arrays
        void visit_DebugCheckArrayBounds(const ASR::DebugCheckArrayBounds_t& x) {
            bool in_debugcheck_copy = in_debugcheck;
            in_debugcheck = true;
            ASR::CallReplacerOnExpressionsVisitor<ReplaceIntrinsicFunctionsVisitor>::visit_DebugCheckArrayBounds(x);
            in_debugcheck = in_debugcheck_copy;
        }

        void visit_ttype(const ASR::ttype_t& x) {
            bool in_ttype_copy = in_ttype;
            in_ttype = true;
            ASR::CallReplacerOnExpressionsVisitor<ReplaceIntrinsicFunctionsVisitor>::visit_ttype(x);
            in_ttype = in_ttype_copy;
        }

        void call_replacer() {
            replacer.current_expr = current_expr;
            replacer.replace_expr(*current_expr);
        }

};

class ReplaceFunctionCallReturningArray: public ASR::BaseExprReplacer<ReplaceFunctionCallReturningArray> {

    private:

    Allocator& al;
    Vec<ASR::stmt_t*>& pass_result;
    size_t result_counter;
    std::map<ASR::symbol_t*, ASRUtils::IntrinsicArrayFunctions>& func2intrinsicid;

    public:
    ASR::expr_t* result_var_; // Declared in array_struct_temporary
    SymbolTable* current_scope;

    ReplaceFunctionCallReturningArray(Allocator& al_, Vec<ASR::stmt_t*>& pass_result_,
    std::map<ASR::symbol_t*, ASRUtils::IntrinsicArrayFunctions>& func2intrinsicid_) :
    al(al_), pass_result(pass_result_), result_counter(0),
    func2intrinsicid(func2intrinsicid_),
    result_var_(nullptr),
    current_scope(nullptr) {}

    // Not called from anywhere but kept for future use.
    // Especially if we don't find alternatives to allocatables
    ASR::expr_t* get_result_var_for_runtime_dim(ASR::expr_t* dim, int n_dims,
        std::string m_name, const Location& loc, ASR::ttype_t* m_type,
        ASR::expr_t* input_array) {
        m_type = ASRUtils::TYPE(ASR::make_Allocatable_t(al, loc,
            ASRUtils::type_get_past_allocatable(m_type)));
        ASR::expr_t* result_var_ = PassUtils::create_var(result_counter,
                                    m_name + "_res",
                                    loc, m_type, al, current_scope);
        ASR::stmt_t** else_ = nullptr;
        size_t else_n = 0;
        const Location& loc_ = dim->base.loc;
        for( int i = 1; i <= n_dims + 1; i++ ) {
            ASR::expr_t* current_dim = ASRUtils::EXPR(ASR::make_IntegerConstant_t(
                                        al, loc_, i, ASRUtils::expr_type(dim)));
            ASR::expr_t* test_expr = ASRUtils::EXPR(ASR::make_IntegerCompare_t(
                                        al, loc_, dim, ASR::cmpopType::Eq,
                                        current_dim, ASRUtils::TYPE(ASR::make_Logical_t(
                                            al, loc_, 4)), nullptr));

            ASR::alloc_arg_t alloc_arg;
            alloc_arg.m_len_expr = nullptr;
            alloc_arg.m_type = nullptr;
            alloc_arg.loc = loc_;
            alloc_arg.m_a = result_var_;
            alloc_arg.m_sym_subclass = nullptr;
            Vec<ASR::dimension_t> alloc_dims;
            alloc_dims.reserve(al, n_dims);
            for( int j = 1; j <= n_dims + 1; j++ ) {
                if( j == i ) {
                    continue ;
                }
                ASR::dimension_t m_dim;
                m_dim.loc = loc_;
                m_dim.m_start = ASRUtils::EXPR(ASR::make_IntegerConstant_t(
                                al, loc_, 1, ASRUtils::expr_type(dim)));
                // Assuming that first argument is the array
                m_dim.m_length = ASRUtils::get_size(input_array, j, al);
                alloc_dims.push_back(al, m_dim);
            }
            alloc_arg.m_dims = alloc_dims.p;
            alloc_arg.n_dims = alloc_dims.size();
            Vec<ASR::alloc_arg_t> alloc_args;
            alloc_args.reserve(al, 1);
            alloc_args.push_back(al, alloc_arg);
            Vec<ASR::expr_t*> to_be_deallocated;
            to_be_deallocated.reserve(al, alloc_args.size());
            for( size_t i = 0; i < alloc_args.size(); i++ ) {
                to_be_deallocated.push_back(al, alloc_args.p[i].m_a);
            }
            ASR::stmt_t* allocate_stmt = ASRUtils::STMT(ASR::make_Allocate_t(
                al, loc_, alloc_args.p, alloc_args.size(), nullptr, nullptr, nullptr));
            Vec<ASR::stmt_t*> if_body;
            if_body.reserve(al, 2);
            if_body.push_back(al, ASRUtils::STMT(ASR::make_ExplicitDeallocate_t(
                al, loc, to_be_deallocated.p, to_be_deallocated.size())));
            if_body.push_back(al, allocate_stmt);
            ASR::stmt_t* if_ = ASRUtils::STMT(ASR::make_If_t(al, loc_, nullptr, test_expr,
                                if_body.p, if_body.size(), else_, else_n));
            Vec<ASR::stmt_t*> if_else_if;
            if_else_if.reserve(al, 1);
            if_else_if.push_back(al, if_);
            else_ = if_else_if.p;
            else_n = if_else_if.size();
        }
        pass_result.push_back(al, else_[0]);
        return result_var_;
    }

    ASR::expr_t* get_result_var_for_constant_dim(int dim, int n_dims,
        std::string m_name, const Location& loc, ASR::ttype_t* m_type,
        ASR::expr_t* input_array) {
        Vec<ASR::dimension_t> result_dims;
        result_dims.reserve(al, n_dims);
        for( int j = 1; j <= n_dims + 1; j++ ) {
            if( j == dim ) {
                continue ;
            }
            ASR::dimension_t m_dim;
            m_dim.loc = loc;
            m_dim.m_start = ASRUtils::EXPR(ASR::make_IntegerConstant_t(
                            al, loc, 1, ASRUtils::TYPE(ASR::make_Integer_t(al, loc, 4)) ));
            // Assuming that first argument is the array
            m_dim.m_length = ASRUtils::get_size(input_array, j, al);
            result_dims.push_back(al, m_dim);
        }
        ASR::ttype_t* result_type = ASRUtils::duplicate_type(al, m_type, &result_dims);
        return PassUtils::create_var(result_counter,
                                    m_name + "_res",
                                    loc, result_type, al, current_scope);
    }

    void replace_FunctionCall(ASR::FunctionCall_t* x) {
        ASR::symbol_t* x_m_name = ASRUtils::symbol_get_past_external(x->m_name);
        if( func2intrinsicid.find(x_m_name) == func2intrinsicid.end() ) {
            return;
        }

        if( !this->result_var_ ) {
            // No result variable (e.g., parameter initialization);
            // the value is compile-time constant, skip replacement.
            return;
        }

        ASR::expr_t* target_var = this->result_var_;
        bool is_param = false;
        if (ASR::is_a<ASR::Var_t>(*target_var)) {
            ASR::Variable_t* v = ASRUtils::EXPR2VAR(target_var);
            if (v->m_storage == ASR::storage_typeType::Parameter) {
                is_param = true;
            }
        }

        if (is_param) {
            ASR::ttype_t* target_type = ASRUtils::expr_type(target_var);
            target_var = PassUtils::create_var(++result_counter,
                                        "_param_init_temp",
                                        x->base.base.loc, target_type, al, current_scope);
        }

        Vec<ASR::call_arg_t> new_args;
        new_args.reserve(al, x->n_args + 1);
        for( size_t i = 0; i < x->n_args; i++ ) {
            ASR::call_arg_t new_arg;
            new_arg.loc = x->m_args[i].loc;
            new_arg.m_value = x->m_args[i].m_value;
            new_args.push_back(al, new_arg);
        }
        ASR::call_arg_t new_arg;
        LCOMPILERS_ASSERT(target_var)
        new_arg.loc = target_var->base.loc;
        new_arg.m_value = target_var;
        new_args.push_back(al, new_arg);


        ASR::stmt_t* subrout_call = ASRUtils::STMT(ASRUtils::make_SubroutineCall_t_util(
            al, x->base.base.loc, x->m_name, x->m_original_name, new_args.p,
            new_args.size(), x->m_dt, nullptr, false));
        pass_result.push_back(al, subrout_call);

        if (is_param) {
            ASR::stmt_t* param_assign = ASRUtils::STMT(ASR::make_Assignment_t(
                al, x->base.base.loc, this->result_var_, target_var, nullptr, false, false));
            pass_result.push_back(al, param_assign);
        }
    }
};

class ReplaceFunctionCallReturningArrayVisitor : public ASR::CallReplacerOnExpressionsVisitor<ReplaceFunctionCallReturningArrayVisitor>
{
    private:

        Allocator& al;
        ReplaceFunctionCallReturningArray replacer;
        Vec<ASR::stmt_t*> pass_result;
        Vec<ASR::stmt_t*>* parent_body;

    public:

        ReplaceFunctionCallReturningArrayVisitor(Allocator& al_,
            std::map<ASR::symbol_t*, ASRUtils::IntrinsicArrayFunctions>& func2intrinsicid_) :
        al(al_),
        replacer(al_, pass_result, func2intrinsicid_),
        parent_body(nullptr) {
            pass_result.n = 0;
        }


        void call_replacer() {
            replacer.current_expr = current_expr;
            replacer.current_scope = current_scope;
            replacer.replace_expr(*current_expr);
        }

        void transform_stmts(ASR::stmt_t **&m_body, size_t &n_body) {
            Vec<ASR::stmt_t*> body;
            body.reserve(al, n_body);
            if( parent_body ) {
                for (size_t j=0; j < pass_result.size(); j++) {
                    parent_body->push_back(al, pass_result[j]);
                }
            }

            for (size_t i=0; i<n_body; i++) {
                pass_result.n = 0;
                pass_result.reserve(al, 1);
                Vec<ASR::stmt_t*>* parent_body_copy = parent_body;
                parent_body = &body;
                visit_stmt(*m_body[i]);
                parent_body = parent_body_copy;
                if( pass_result.size() > 0 ) {
                    for (size_t j=0; j < pass_result.size(); j++) {
                        body.push_back(al, pass_result[j]);
                    }
                } else {
                    body.push_back(al, m_body[i]);
                }
            }
            m_body = body.p;
            n_body = body.size();
            pass_result.n = 0;
        }

        void visit_Variable(const ASR::Variable_t& x) {
            if (x.m_storage == ASR::storage_typeType::Parameter) {
                return;
            }
            ASR::CallReplacerOnExpressionsVisitor<ReplaceFunctionCallReturningArrayVisitor>::visit_Variable(x);
        }

        void visit_Assignment(const ASR::Assignment_t& x) {
            replacer.result_var_ = x.m_target;
            ASR::CallReplacerOnExpressionsVisitor<ReplaceFunctionCallReturningArrayVisitor>::visit_Assignment(x);
            replacer.result_var_ = nullptr;
        }

};

void pass_replace_intrinsic_function(Allocator &al, ASR::TranslationUnit_t &unit,
                            const LCompilers::PassOptions& pass_options) {
    int index_kind = pass_options.descriptor_index_64 ? 8 : 4;
    std::map<ASR::symbol_t*, ASRUtils::IntrinsicArrayFunctions> func2intrinsicid;
    ReplaceIntrinsicFunctionsVisitor v(al, unit.m_symtab, func2intrinsicid, index_kind);
    v.visit_TranslationUnit(unit);
    ReplaceFunctionCallReturningArrayVisitor u(al, func2intrinsicid);
    u.visit_TranslationUnit(unit);
    PassUtils::UpdateDependenciesVisitor w(al);
    w.visit_TranslationUnit(unit);
}


} // namespace LCompilers
