/*****************************************************************************
 *
 * Module: st_construct
 *
 * Prefix: ST
 *
 * Description:
 *
 * This module constructs a symbol table from the AST
 *
 *****************************************************************************/

#include "st_construct.h"
#include "dbug.h"
#include "mytypes.h"
#include "queue.h"
#include "symbol_table.h"
#include "traverse.h"
#include "tree_basic.h"
#include "types.h"
#include "utils.h"

#include "ctinfo.h"
#include "memory.h"
#include "mymacros.h"
#include "free.h"

/*
 * INFO structure
 */

typedef struct q_item
{
    struct q_item* next;
    node*          node;
    info*          info;
} q_item;

struct INFO
{
    struct symbol_table* st;
    struct queue*        late_traversal_nodes;
    /* For loops will influence the relative frame number, since they count as another scope (for
     * the induction variable) but are in reality in the same frame as the rest of the function.
     * This variable accounts for the number of for loops we're nested in, and should be subtracted
     * from any loads/stores (var/varlet). */
    int for_depth;
};

/*
 * INFO functions
 */

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

    DBUG_ENTER("info_init");

    result = (info*)MEMmalloc(sizeof(struct INFO));

    result->st                   = st_init(NULL);
    result->late_traversal_nodes = q_init();
    result->for_depth            = 0;
    DBUG_RETURN(result);
}

static void info_free(info* inf)
{
    DBUG_ENTER("info_free");

    MEMfree(inf);
}

static info* info_copy(info* orig)
{
    if (!orig) return NULL;

    info* copy = malloc(sizeof(info));
    *copy      = *orig;

    return copy;
}

/*
 * Helper functions
 */

struct st_entry_arr* def_register_array(struct symbol_table* st, char* name, type_t type, int n_dims,
                                bool isexported, bool isextern)
{
    struct st_entry_arr* arr_ent = st_insert_arr(st, name, type, n_dims, isexported, isextern);

    for (int i = 0; i < n_dims; i++)
    {
        char*            dim_name = get_dim_name(name, i);
        struct st_entry* dim_var  = st_lookup(st, dim_name, LOOKUP_LOCAL | LOOKUP_VAR, NULL);
        if (!dim_var)
        {
            CTIabort("Unknown dimension variable (or not in local scope even though this is a "
                     "vardecl/globaldef?) '%s'\n",
                     dim_name);
        }
        arr_ent->dim_vars[i] = dim_var;
        free(dim_name);
    }

    return arr_ent;
}


/* Register an array parameter. This is similar but different from register_array for vardecl or globaldef,
 * so, sadly, here's a mostly-carbon-copy function.
 */
struct st_entry_arr* param_register_array(struct symbol_table* st, node* ids, char* name, type_t type, int n_dims,
                                bool isexported, bool isextern)
{
    struct st_entry_arr* arr_ent = st_insert_arr(st, name, type, n_dims, isexported, isextern);

    node* id = ids;

    for (int i = 0; i < n_dims; id = IDS_NEXT(id), i++)
    {
        char*            dim_name = IDS_NAME(id);
        struct st_entry* dim_var  = st_lookup(st, dim_name, LOOKUP_LOCAL | LOOKUP_VAR, NULL);
        if (!dim_var)
        {
            CTIabort("Unknown dimension variable (or not in local scope even though this is a "
                     "vardecl/globaldef?) '%s'\n",
                     dim_name);
        }
        arr_ent->dim_vars[i] = dim_var;
    }

    return arr_ent;
}


/*
 * Traversal functions
 */

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

    char* name = VARLET_NAME(arg_node);
    int   frameno;

    struct st_entry* ent = st_lookup(arg_info->st, name, LOOKUP_GLOBAL | LOOKUP_VAR, &frameno);
    if (!ent)
    {
        ERR("Variable `%s' not declared in current scope", name);
        DBUG_RETURN(arg_node);
    }

    TO(VARLET_INDICES);

    VARLET_STLINK(arg_node)     = ent;
    VARLET_RELFRAMENO(arg_node) = MAX(frameno - arg_info->for_depth, 0);

    DBUG_RETURN(arg_node);
}

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

    char* name = VAR_NAME(arg_node);
    int   frameno;

    struct st_entry* ent = st_lookup(arg_info->st, name, LOOKUP_GLOBAL | LOOKUP_VAR, &frameno);
    if (!ent)
    {
        ERR("Variable `%s' not declared in current scope", name);
        DBUG_RETURN(arg_node);
    }

    TO(VAR_INDICES);

    VAR_STLINK(arg_node)     = ent;
    VAR_RELFRAMENO(arg_node) = MAX(frameno - arg_info->for_depth, 0);

    DBUG_RETURN(arg_node);
}

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

    char* name = CALL_NAME(arg_node);

    struct st_entry_fn* ent =
        (struct st_entry_fn*)st_lookup(arg_info->st, name, LOOKUP_GLOBAL | LOOKUP_FUN, NULL);
    if (!ent)
    {
        ERR("Function `%s' not declared in current scope", name);
        DBUG_RETURN(arg_node);
    }

    CALL_STLINK(arg_node) = ent;

    TO(CALL_PARAMS);

    DBUG_RETURN(arg_node);
}

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

    char*                name = VARDECL_NAME(arg_node);
    struct symbol_table* st   = arg_info->st;

    if (st_check(st, name, LOOKUP_LOCAL | LOOKUP_VAR))
    {
        CTIerrorLine(NODE_LINE(arg_node), "Variable %s already declared in local scope", name);
        DBUG_RETURN(arg_node);
    }

    /* Traverse to the initialiser first, since we can refer to an upper scoped var with the same
     * name (i.e. int a = a in a local function)
     */
    TO(VARDECL_INIT);

    struct st_entry* ent;

    if (VARDECL_IS_ARRAY(arg_node))
    {
        int    n_dims = count_exprs(VARDECL_DIMS(arg_node));
        type_t type   = VARDECL_TYPE(arg_node);
        ent           = (struct st_entry*)def_register_array(st, name, type, n_dims, false, false);

        FREEdoFreeTree(VARDECL_DIMS(arg_node));
        VARDECL_DIMS(arg_node) = NULL;
    }
    else
    {
        ent = st_insert(st, name, VARDECL_TYPE(arg_node), false, false);
    }

    VARDECL_NEXT(arg_node)   = TRAVopt(VARDECL_NEXT(arg_node), arg_info);
    VARDECL_STLINK(arg_node) = ent;

    DBUG_RETURN(arg_node);
}

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

    char*                name = FUNDEF_NAME(arg_node);
    struct symbol_table* st   = arg_info->st;

    if (st_check(st, name, LOOKUP_LOCAL | LOOKUP_FUN))
    {
        CTIerrorLine(NODE_LINE(arg_node), "Function %s already declared in local scope", name);
        DBUG_RETURN(arg_node);
    }

    int   num_params = 0;
    node* param      = FUNDEF_PARAMS(arg_node);
    while (param)
    {
        num_params++;
        param = PARAM_NEXT(param);
    }

    struct st_entry_fn* ent = st_insert_fn(arg_info->st, name, FUNDEF_TYPE(arg_node), num_params,
                                           FUNDEF_ISEXPORTED(arg_node), FUNDEF_ISEXTERN(arg_node));

    FUNDEF_STLINK(arg_node) = ent;

    arg_info->st = ent->child_syms;

    printf("Processing function definition: %s\n", name);

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

    q_nq(arg_info->late_traversal_nodes, FUNDEF_BODY(arg_node), info_copy(arg_info));

    arg_info->st = st;

    DBUG_RETURN(arg_node);
}

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

    char*                name = GLOBALDEF_NAME(arg_node);
    struct symbol_table* st   = arg_info->st;

    printf("Processing global definition: %s\n", name);

    if (st_check(st, name, LOOKUP_LOCAL | LOOKUP_VAR))
    {
        CTIerrorLine(NODE_LINE(arg_node), "Variable %s already declared in local scope", name);
    }

    struct st_entry* ent;

    if (GLOBALDEF_IS_ARRAY(arg_node))
    {
        int n_dims;
        node* dims = GLOBALDEF_DIMS(arg_node);
        if (NODE_TYPE(dims) == N_ids)
            n_dims = count_ids(dims);
        else
            n_dims = count_exprs(dims);

        type_t type   = GLOBALDEF_TYPE(arg_node);
        ent           = (struct st_entry*)def_register_array(st, name, type, n_dims, GLOBALDEF_ISEXPORTED(arg_node),
                             GLOBALDEF_ISEXTERN(arg_node));

        FREEdoFreeTree(GLOBALDEF_DIMS(arg_node));
        GLOBALDEF_DIMS(arg_node) = NULL;
    }
    else
    {
        ent = st_insert(st, name, GLOBALDEF_TYPE(arg_node), GLOBALDEF_ISEXPORTED(arg_node),
                        GLOBALDEF_ISEXTERN(arg_node));
    }

    TO(GLOBALDEF_DIMS);
    TO(GLOBALDEF_INIT);

    GLOBALDEF_STLINK(arg_node) = ent;

    DBUG_RETURN(arg_node);
}

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

    char*                name = PARAM_NAME(arg_node);
    struct symbol_table* st   = arg_info->st;

    if (st_check(st, name, LOOKUP_LOCAL | LOOKUP_VAR))
    {
        CTIerrorLine(NODE_LINE(arg_node), "Parameter %s already declared in local scope", name);
    }

    struct st_entry* ent;

    if (PARAM_IS_ARRAY(arg_node))
    {
        int n_dims = count_ids(PARAM_DIMS(arg_node));
        type_t type = PARAM_TYPE(arg_node);
        ent = AS_VAR(param_register_array(st, PARAM_DIMS(arg_node), name, type, n_dims, false, false));
    }
    else
    {
        ent   = st_insert(st, name, PARAM_TYPE(arg_node), false, false);
    }

    PARAM_STLINK(arg_node) = ent;

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

    DBUG_RETURN(arg_node);
}

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

    char*                name = "_for";
    struct symbol_table* st   = arg_info->st;
    arg_info->for_depth++;

    struct st_entry_fn* ent = st_insert_fn(arg_info->st, name, 0, TY_unknown, false, false);

    FOR_STLINK(arg_node) = ent;

    arg_info->st = ent->child_syms;

    st_insert(arg_info->st, FOR_VARID(arg_node), TY_int | TY_const, false, false);

    TO(FOR_VARINITVAL);
    TO(FOR_VARLIMVAL);
    TO(FOR_VARSTEPVAL);
    TO(FOR_BLOCK);

    arg_info->for_depth--;

    arg_info->st = st;

    DBUG_RETURN(arg_node);
}

void traverse_queue(struct queue* queue)
{
    struct q_item* cur_item;
    while ((cur_item = q_dq(queue)))
    {
        node* n   = cur_item->node;
        info* inf = cur_item->info;

        TRAVopt(n, inf);

        free(inf);
        free(cur_item);
    }
}

/*
 * Traversal start function
 */

node* STdoCreateSymTable(node* syntaxtree)
{
    info* arg_info;

    DBUG_ENTER("STdoCreateSymTable");

    arg_info = info_init();

    TRAVpush(TR_st);
    syntaxtree = TRAVdo(syntaxtree, arg_info);
    traverse_queue(arg_info->late_traversal_nodes);
    TRAVpop();

    q_clean(arg_info->late_traversal_nodes);

    st_print(arg_info->st, NULL, 1);

    PROGRAM_SYMBOLTABLE(syntaxtree) = arg_info->st;

    info_free(arg_info);

    DBUG_RETURN(syntaxtree);
}
