#include "vardecl_rewrite.h"
#include "assert.h"
#include "copy.h"
#include "dbug.h"
#include "free.h"
#include "memory.h"
#include "mymacros.h"
#include "symbol_table.h"
#include "traverse.h"
#include "tree_basic.h"
#include "utils.h"
#include <string.h>

struct INFO
{
    node* local_stmts_head;
    node* place_to_append_new_vardecl_assign;
    node* module_init_stmts_head;
    node* module_init_stmts_tail;
    bool  past_vardecl_head;
    int   cur_n_dims;
};

static info* info_init(void)
{
    info* result;
    DBUG_ENTER("info_init");

    result = (info*)MEMmalloc(sizeof(struct INFO));
    memset(result, 0, sizeof(struct INFO));

    DBUG_RETURN(result);
}

static void info_free(info* arg_info) { MEMfree(arg_info); }

/* Prepend an assign statement to the stmts head. Returns the new stmt. */
node* assign_stmt_prepend_head(node* stmt, struct st_entry* var, node* expr, bool global)
{
    node* varlet              = TBmakeVarlet(strdup(var->name), NULL);
    VARLET_STLINK(varlet)     = var;
    VARLET_RELFRAMENO(varlet) = 0;
    node* assign              = TBmakeAssign(varlet, expr);

    return TBmakeStmts(assign, stmt);
}

/* Append an assign statement after the given statement. Returns the new stmt.
 */
node* assign_stmt_append(node* stmt, struct st_entry* var, node* expr, bool global)
{
    assert(stmt != NULL);
    node* varlet              = TBmakeVarlet(strdup(var->name), NULL);
    VARLET_STLINK(varlet)     = var;
    VARLET_RELFRAMENO(varlet) = 0;
    node* assign              = TBmakeAssign(varlet, expr);

    node* new_stmt   = TBmakeStmts(assign, STMTS_NEXT(stmt));
    STMTS_NEXT(stmt) = new_stmt;
    return new_stmt;
}

/* Generate assignments to array indices, for array initializers that are fixed-size
 * (so the size of what to initialize is known at compile time)
 * (i.e. int[...] a = [1, 2];)
 */
void generate_array_assignments(node* vardecl_init, struct st_entry* arr_ent, info* arg_info)
{
    bool is_global = IS_GLOBAL_SYMBOL_TABLE(arr_ent->containing_table);

    vardecl_init = flatten_arrexpr(vardecl_init, NULL, NULL, NULL);

    node* init_exprs = ARREXPR_EXPRS(vardecl_init);
    node* new_tail   = NULL;
    node* next       = NULL;

    int index = 0;

    for (; init_exprs != NULL; init_exprs = next)
    {
        node* expr = EXPRS_EXPR(init_exprs);
        next       = EXPRS_NEXT(init_exprs);
        if (is_global)
        {
            new_tail = assign_stmt_append(arg_info->module_init_stmts_tail, arr_ent, expr, false);
            arg_info->module_init_stmts_tail = new_tail;
        }
        else
        {
            new_tail = assign_stmt_append(arg_info->place_to_append_new_vardecl_assign, arr_ent,
                                          expr, false);
            arg_info->place_to_append_new_vardecl_assign = new_tail;
        }
        VARLET_INDICES(ASSIGN_LET(STMTS_STMT(new_tail))) = TBmakeExprs(TBmakeInt(index++), NULL);
        EXPRS_EXPR(init_exprs)                           = NULL;
        EXPRS_NEXT(init_exprs)                           = NULL;
        FREEdoFreeTree(init_exprs);
    }

    ARREXPR_EXPRS(vardecl_init) = NULL;
    FREEdoFreeTree(vardecl_init);
}

/* Generates a while loop to initialize arrays to a scalar value. */
void generate_array_while(node* vardecl_init, struct st_entry_arr* arr_ent, info* arg_info)
{
    bool  is_global = IS_GLOBAL_SYMBOL_TABLE(arr_ent->containing_table);
    node* dims_expr = get_dims_expr(arr_ent);

    // Fix while induction variable
    node* new_var = TBmakeVar(get_array_while_name(arr_ent->name), NULL);
    VAR_STLINK(new_var) =
        st_insert(arr_ent->containing_table, VAR_NAME(new_var), TY_int, false, false);

    /* This is the variable that each element is initialized to. Since the initializer is only
     * evaluated once, this needs to be saved in a separate variable. */
    node* new_var_while         = TBmakeVar(get_array_while_var(arr_ent->name), NULL);
    VAR_STLINK(new_var_while)   = st_insert(arr_ent->containing_table, VAR_NAME(new_var_while),
                                          TY_ARRAY_ELEMENT_TYPE(arr_ent->type), false, false);

    node* varlet_while          = TBmakeVarlet(strdup(VAR_NAME(new_var_while)), NULL);
    VARLET_STLINK(varlet_while) = VAR_STLINK(new_var_while);
    node* while_var_ass         = TBmakeAssign(varlet_while, vardecl_init);

    // Gen assign
    node* varlet          = TBmakeVarlet(strdup(arr_ent->name), TBmakeExprs(new_var, NULL));
    VARLET_STLINK(varlet) = AS_VAR(arr_ent);
    node* assign          = TBmakeAssign(varlet, new_var_while);

    // Gen condition
    node* condition           = TBmakeBinop(BO_lt, COPYdoCopy(new_var), dims_expr);
    BINOP_TYPE(condition)     = TY_bool;
    BINOP_EXPRTYPE(condition) = TY_int;

    // Gen while
    node* varlet_incr          = TBmakeVarlet(strdup(VAR_NAME(new_var)), NULL);
    VARLET_STLINK(varlet_incr) = VAR_STLINK(new_var);
    node* binop_incr           = TBmakeBinop(BO_add, COPYdoCopy(new_var), TBmakeInt(1));
    BINOP_TYPE(binop_incr)     = TY_int;
    BINOP_EXPRTYPE(binop_incr) = TY_int;
    node* incr                 = TBmakeAssign(varlet_incr, binop_incr);
    node* whl =
        TBmakeStmts(TBmakeWhile(condition, TBmakeStmts(assign, TBmakeStmts(incr, NULL))), NULL);
    node* stmt = TBmakeStmts(while_var_ass,
                             TBmakeStmts(TBmakeAssign(COPYdoCopy(varlet_incr), TBmakeInt(0)), whl));

    node* tail = (is_global ? arg_info->module_init_stmts_tail
                            : arg_info->place_to_append_new_vardecl_assign);

    STMTS_NEXT(whl)  = STMTS_NEXT(tail);
    STMTS_NEXT(tail) = stmt;

    arg_info->place_to_append_new_vardecl_assign = whl;
}

node* VDRvardecl(node* arg_node, info* arg_info)
{
    DBUG_ENTER("VDRvardecl");

    if (VARDECL_INIT(arg_node))
    {
        node* new_stmt;
        if (arg_info->past_vardecl_head)
        {
            new_stmt = assign_stmt_append(arg_info->place_to_append_new_vardecl_assign,
                                          VARDECL_STLINK(arg_node), VARDECL_INIT(arg_node), false);
        }
        else
        {
            new_stmt =
                assign_stmt_prepend_head(arg_info->local_stmts_head, VARDECL_STLINK(arg_node),
                                         VARDECL_INIT(arg_node), false);
            // Since we're prepending the statement we need to override
            // arg_info's local_stmts head.
            arg_info->local_stmts_head  = new_stmt;
            arg_info->past_vardecl_head = true;
        }
        arg_info->place_to_append_new_vardecl_assign = new_stmt;

        // Initialized with an array expression
        if (NODE_TYPE(VARDECL_INIT(arg_node)) == N_arrexpr)
        {
            generate_array_assignments(VARDECL_INIT(arg_node), VARDECL_STLINK(arg_node), arg_info);
            node* dims_expr                   = get_dims_expr(AS_ARRAY(VARDECL_STLINK(arg_node)));
            ASSIGN_EXPR(STMTS_STMT(new_stmt)) = dims_expr;
        }
        // Initialized with a scalar value
        else if (VARDECL_IS_ARRAY(arg_node))
        {
            generate_array_while(VARDECL_INIT(arg_node), AS_ARRAY(VARDECL_STLINK(arg_node)),
                                 arg_info);
            node* dims_expr                   = get_dims_expr(AS_ARRAY(VARDECL_STLINK(arg_node)));
            ASSIGN_EXPR(STMTS_STMT(new_stmt)) = dims_expr;
        }
    }

    /* We've just moved any initialisation to an assign statement,
     * so nullify the init expression */
    VARDECL_INIT(arg_node) = NULL;
    VARDECL_NEXT(arg_node) = TRAVopt(VARDECL_NEXT(arg_node), arg_info);

    DBUG_RETURN(arg_node);
}

node* VDRfunbody(node* arg_node, info* arg_info)
{
    DBUG_ENTER("VDRfunbody");

    info saved_info = *arg_info;

    arg_info->local_stmts_head                   = FUNBODY_STMTS(arg_node);
    arg_info->place_to_append_new_vardecl_assign = NULL; // first prepend before local_stmts_head
    arg_info->past_vardecl_head                  = false;

    FUNBODY_VARDECLS(arg_node) = TRAVopt(FUNBODY_VARDECLS(arg_node), arg_info);

    /* If FUNBODY_STMTS was NULL, we need the result. */
    FUNBODY_STMTS(arg_node) = arg_info->local_stmts_head;

    TO(FUNBODY_FUNDEFS);

    *arg_info = saved_info;

    DBUG_RETURN(arg_node);
}

node* VDRglobaldef(node* arg_node, info* arg_info)
{
    DBUG_ENTER("VDRglobaldef");
    if (GLOBALDEF_ISEXTERN(arg_node)) DBUG_RETURN(arg_node);
    if (!GLOBALDEF_INIT(arg_node) && !GLOBALDEF_TY_IS_ARRAY(arg_node)) DBUG_RETURN(arg_node);

    node* new_stmt;

    if (!arg_info->module_init_stmts_head)
    {
        new_stmt =
            assign_stmt_prepend_head(arg_info->module_init_stmts_head, GLOBALDEF_STLINK(arg_node),
                                     GLOBALDEF_INIT(arg_node), true);
        arg_info->module_init_stmts_head = new_stmt;
    }
    else
    {
        new_stmt = assign_stmt_append(arg_info->module_init_stmts_tail, GLOBALDEF_STLINK(arg_node),
                                      GLOBALDEF_INIT(arg_node), true);
    }
    arg_info->module_init_stmts_tail = new_stmt;

    if (GLOBALDEF_INIT(arg_node))
    {
        if (NODE_TYPE(GLOBALDEF_INIT(arg_node)) == N_arrexpr)
        {
            generate_array_assignments(GLOBALDEF_INIT(arg_node), GLOBALDEF_STLINK(arg_node),
                                       arg_info);
        }
        else if (GLOBALDEF_IS_ARRAY(arg_node))
        {
            generate_array_while(GLOBALDEF_INIT(arg_node), AS_ARRAY(GLOBALDEF_STLINK(arg_node)),
                                 arg_info);
        }
    }

    if (GLOBALDEF_TY_IS_ARRAY(arg_node))
    {
        node* dims_expr                   = get_dims_expr(AS_ARRAY(GLOBALDEF_STLINK(arg_node)));
        ASSIGN_EXPR(STMTS_STMT(new_stmt)) = dims_expr;
    }

    GLOBALDEF_INIT(arg_node) = NULL;

    DBUG_RETURN(arg_node);
}

/* Insert the module-wide __init function which initialises
 * global definitions.
 */
void insert_module_init_fn(node* program, node* stmts, node* stmts_tail)
{
    if (stmts == NULL) return;

    node* ret              = TBmakeReturn(NULL);
    RETURN_TYPE(ret)       = TY_void;
    node* stmt             = TBmakeStmts(ret, NULL);
    STMTS_NEXT(stmts_tail) = stmt;

    node* body              = TBmakeFunbody(NULL, NULL, stmts);
    node* fundef            = TBmakeFundef(TY_void, strdup("__init"), NULL, body);
    node* declhead          = TBmakeDecls(fundef, PROGRAM_DECLS(program));
    PROGRAM_DECLS(program)  = declhead;
    struct symbol_table* st = PROGRAM_SYMBOLTABLE(program);
    FUNDEF_STLINK(fundef)   = st_insert_fn(st, "__init", TY_void, 0, true, false);
}

node* VDRdoRewrite(node* syntaxtree)
{
    DBUG_ENTER("VDRdoRewrite");
    info* arg_info = info_init();

    TRAVpush(TR_vdr);
    syntaxtree = TRAVdo(syntaxtree, arg_info);
    TRAVpop();

    insert_module_init_fn(syntaxtree, arg_info->module_init_stmts_head,
                          arg_info->module_init_stmts_tail);
    info_free(arg_info);

    DBUG_RETURN(syntaxtree);
}