/*****************************************************************************
 *
 * Module: check_types
 *
 * Prefix: CT
 *
 * Description:
 *
 * This module does type checking.
 *
 *****************************************************************************/

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

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

/*
 * INFO structure
 */

struct INFO
{
    struct symbol_table* st;
    struct st_entry_fn*  cur_fun;
    type_t               inferred_type;

    /* returned from var for call expressions, so that param dims and var dims can be compared. */
    int                  array_dimensions;
    bool                 top_level;
    bool                 top_level_has_return;
};

/*
 * INFO functions
 */

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

    DBUG_ENTER("MakeInfo");

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

    DBUG_RETURN(result);
}

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

    inf = MEMfree(inf);

    DBUG_RETURN(inf);
}

/*
 * Helper functions.
 */

void check_index_types(struct st_entry_arr* arr_ent, node* arg_node, info* arg_info)
{
    int n_dims = arr_ent->n_dims;
    int n = 0;
    for (node* expr = arg_node; expr != NULL; expr = EXPRS_NEXT(expr), n++)
    {
        type_t ty = GET_INFERRED_TYPE_NODE(EXPRS_EXPR(expr));
        if (ty != TY_int)
            ERR("Invalid type in array indexing expression: '%s' (expected int).", type_to_str(ty));
    }
    if (n != n_dims)
        ERR("Invalid number of dimensions in array indexing (array '%s' expected %d indices, got %d).", arr_ent->name, n_dims, n);
}

/*
 * Traversal functions
 */

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

    arg_info->inferred_type = TY_int;

    DBUG_RETURN(arg_node);
}

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

    arg_info->inferred_type = TY_float;

    DBUG_RETURN(arg_node);
}

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

    arg_info->inferred_type = TY_bool;

    DBUG_RETURN(arg_node);
}

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

    binop  op_type = BINOP_OP(arg_node);
    type_t l_type  = GET_INFERRED_TYPE(BINOP_LEFT);
    type_t r_type  = GET_INFERRED_TYPE(BINOP_RIGHT);

    if (TY_IS_ARRAY(l_type) || TY_IS_ARRAY(r_type))
    {
        ERR("Binop applied to array variables");
    }

    if (l_type != r_type)
        CTIerrorLine(NODE_LINE(arg_node), "Binop applied to variables of different types");
    else if (IS_ARITH_OP(op_type))
    {
        if (op_type == BO_mod && l_type != TY_int)
            CTIerrorLine(NODE_LINE(arg_node),
                         "Modulo operator only compatible with operands of type int");
        else if (op_type & ~(BO_add | BO_mul) && l_type == TY_bool)
            CTIerrorLine(NODE_LINE(arg_node),
                         "Arithmetic operator %s not compatible with operands of type bool",
                         binop_to_str(op_type));
    }
    else if (IS_LOG_OP(op_type))
    {
        if (l_type != TY_bool)
            CTIerrorLine(NODE_LINE(arg_node),
                         "Logical operator only compatible with operands of type bool");
    }
    else if (IS_REL_OP(op_type))
    {
        if (op_type & ~(BO_eq | BO_ne) && l_type == TY_bool)
            CTIerrorLine(NODE_LINE(arg_node),
                         "Relational operator %s not compatible with operands of type bool",
                         binop_to_str(op_type));

        arg_info->inferred_type = TY_bool;
    }
    else
    {
        // GEEN IDEE
    }


    /* Save the type for in later traversals. */
    BINOP_TYPE(arg_node)     = arg_info->inferred_type;
    BINOP_EXPRTYPE(arg_node) = l_type;

    DBUG_RETURN(arg_node);
}

node* CTvardecl(node* arg_node, info* arg_info)
{
    type_t var_type = VARDECL_TYPE(arg_node);

    if (VARDECL_INIT(arg_node) != NULL)
    {
        if (NODE_TYPE(VARDECL_INIT(arg_node)) == N_arrexpr)
            arg_info->array_dimensions = 0;

        type_t init_ty = GET_INFERRED_TYPE(VARDECL_INIT);

        if (NODE_TYPE(VARDECL_INIT(arg_node)) == N_arrexpr && arg_info->array_dimensions != AS_ARRAY(VARDECL_STLINK(arg_node))->n_dims)
        {
            ERR("Initializing array `%s' with wrong-dimensioned initializer (initializer has dims `%d', array has dims `%d').",
                VARDECL_NAME(arg_node), arg_info->array_dimensions, AS_ARRAY(VARDECL_STLINK(arg_node))->n_dims);
        }

        if (init_ty != TY_FUNDAMENTAL_TYPE(var_type))
        {
            ERR("Variable of type `%s' cannot be initialized with a value of type `%s'.",
                type_to_str(var_type), type_to_str(init_ty));
        }
    }

    TO(VARDECL_NEXT);

    return arg_node;
}

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

    type_t expr_type = GET_INFERRED_TYPE(MONOP_EXPR);

    if (MONOP_OP(arg_node) == MO_not && expr_type != TY_bool)
        CTIerrorLine(NODE_LINE(arg_node), "Logical negation applied to non-boolean expression");
    if (MONOP_OP(arg_node) == MO_neg && expr_type == TY_bool)
        CTIerrorLine(NODE_LINE(arg_node), "Negation applied to non-numerical expression");

    MONOP_TYPE(arg_node) = arg_info->inferred_type;

    DBUG_RETURN(arg_node);
}

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

    type_t let_type  = GET_INFERRED_TYPE(ASSIGN_LET);
    type_t expr_type = GET_INFERRED_TYPE(ASSIGN_EXPR);

    if (let_type != expr_type)
        CTIerrorLine(NODE_LINE(arg_node), "Expression of type %s assigned to variable of type %s",
                     type_to_str(expr_type), type_to_str(let_type));

    DBUG_RETURN(arg_node);
}

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

    struct st_entry_fn* fun_entry = CALL_STLINK(arg_node);

    node*            param       = CALL_PARAMS(arg_node);
    struct st_entry* param_entry = fun_entry->child_syms->head;

    int old_array_dims = arg_info->array_dimensions;

    for (int i = 0; i < fun_entry->num_params && param != NULL; i++)
    {
        arg_info->array_dimensions = 0;
        type_t param_type       = GET_INFERRED_TYPE_NODE(EXPRS_EXPR(param));

        /* Passing a T to a function expecting a const T is totally allowed. */
        type_t param_entry_type = TY_CONSTLESS_TYPE(param_entry->type);

        if (param_type != param_entry_type)
        {
            ERR("Parameter %d has type %s, but function '%s' expects a parameter of type %s.",
                i + 1, type_to_str(param_type), fun_entry->name, type_to_str(param_entry_type));
        }

        if (TY_IS_ARRAY(param_entry->type) && AS_ARRAY(param_entry)->n_dims != arg_info->array_dimensions)
        {
            ERR("Invalid dimensions for array parameter %d: function expects %d, but %d given.",
                i + 1, AS_ARRAY(param_entry)->n_dims, arg_info->array_dimensions);
        }

        arg_info->array_dimensions = 0;

        param       = EXPRS_NEXT(param);
        param_entry = param_entry->next;
    }

    arg_info->array_dimensions = old_array_dims;

    arg_info->inferred_type = fun_entry->type;

    DBUG_RETURN(arg_node);
}

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

    type_t fun_type = arg_info->cur_fun->type;
    type_t ret_type;

    if (fun_type == TY_void && RETURN_RETVAL(arg_node))
        ERR("Void functions can't have return expressions.");

    if (RETURN_RETVAL(arg_node))
        ret_type = GET_INFERRED_TYPE(RETURN_RETVAL);
    else
        ret_type = TY_void;

    if (fun_type != ret_type)
        CTIerrorLine(NODE_LINE(arg_node),
                     "Return expression type does not correspond to current function return type");

    RETURN_TYPE(arg_node) = ret_type;

    if (arg_info->top_level) arg_info->top_level_has_return = true;

    DBUG_RETURN(arg_node);
}

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

    TO(CAST_EXPR);

    if (TY_IS_ARRAY(arg_info->inferred_type))
    {
        ERR("Can't cast an array to a different type.");
    }
    CAST_EXPRTYPE(arg_node) = arg_info->inferred_type;

    arg_info->inferred_type = CAST_TYPE(arg_node);

    DBUG_RETURN(arg_node);
}

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

    bool was_top_level  = arg_info->top_level;
    arg_info->top_level = false;

    type_t init_type = GET_INFERRED_TYPE(FOR_VARINITVAL);
    type_t stop_type = GET_INFERRED_TYPE(FOR_VARLIMVAL);
    type_t step_type = GET_INFERRED_TYPE(FOR_VARSTEPVAL);

    if (init_type != TY_int)
        CTIerrorLine(NODE_LINE(arg_node), "For initializing expression is not of type int");
    if (stop_type != TY_int)
        CTIerrorLine(NODE_LINE(arg_node), "For limit expression is not of type int");
    if (FOR_VARSTEPVAL(arg_node) && step_type != TY_int)
        CTIerrorLine(NODE_LINE(arg_node), "For step expression is not of type int");

    arg_info->st = FOR_STLINK(arg_node)->child_syms;

    TO(FOR_BLOCK);

    arg_info->st        = arg_info->st->parent;
    arg_info->top_level = was_top_level;

    DBUG_RETURN(arg_node);
}

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

    bool was_top_level  = arg_info->top_level;
    arg_info->top_level = false;

    if (GET_INFERRED_TYPE(WHILE_CONDITION) != TY_bool)
        CTIerrorLine(NODE_LINE(arg_node), "Condition expression value is not of type bool");

    TO(WHILE_BLOCK);

    arg_info->top_level = was_top_level;

    DBUG_RETURN(arg_node);
}

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

    bool was_top_level  = arg_info->top_level;
    arg_info->top_level = false;

    if (GET_INFERRED_TYPE(IF_CONDITION) != TY_bool)
        CTIerrorLine(NODE_LINE(arg_node), "Condition expression value is not of type bool");

    TO(IF_BLOCK);
    TO(IF_ELSEBLOCK);

    arg_info->top_level = was_top_level;

    DBUG_RETURN(arg_node);
}

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

    struct st_entry_fn* cur_fun = FUNDEF_STLINK(arg_node);

    info old_info = *arg_info;

    *arg_info = (info){ .cur_fun              = cur_fun,
                        .st                   = cur_fun->child_syms,
                        .top_level            = true,
                        .top_level_has_return = false };

    if (FUNDEF_BODY(arg_node))
    {
        TO(FUNDEF_BODY);

        if (!arg_info->top_level_has_return && FUNDEF_TYPE(arg_node) != TY_void)
        {
            ERR("Function has no top-level return statement and is therefore not guaranteed to "
                "return a value.");
        }
    }

    *arg_info = old_info;

    DBUG_RETURN(arg_node);
}

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

    /* We're reading, so constness is irrelevant */
    struct st_entry* var_ent = VAR_STLINK(arg_node);
    type_t ty = TY_CONSTLESS_TYPE(var_ent->type);

    if (VAR_IS_ARRAYIDX(arg_node))
    {
        check_index_types(AS_ARRAY(var_ent), VAR_INDICES(arg_node), arg_info);
    }

    arg_info->inferred_type = ty;

    if (TY_IS_ARRAY(ty) && VAR_IS_ARRAYIDX(arg_node))
        arg_info->inferred_type = TY_ARRAY_ELEMENT_TYPE(ty);
    else if (TY_IS_ARRAY(ty))
    {
        arg_info->array_dimensions = AS_ARRAY(var_ent)->n_dims;
    }

    DBUG_RETURN(arg_node);
}

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

    struct st_entry* var_ent = VARLET_STLINK(arg_node);
    type_t ty = var_ent->type;

    if (VARLET_IS_ARRAYIDX(arg_node))
    {
        check_index_types(AS_ARRAY(var_ent), VARLET_INDICES(arg_node), arg_info);
    }

    if (TY_IS_ARRAY(ty) && VARLET_IS_ARRAYIDX(arg_node))
        arg_info->inferred_type = TY_ARRAY_ELEMENT_TYPE(ty);
    else
        arg_info->inferred_type = ty;

    DBUG_RETURN(arg_node);
}

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

    if (!GLOBALDEF_INIT(arg_node)) DBUG_RETURN(arg_node);

    int old_array_dims = arg_info->array_dimensions;
    arg_info->array_dimensions = 0;

    type_t inferred_ty = GET_INFERRED_TYPE(GLOBALDEF_INIT);

    type_t declared_ty = GLOBALDEF_STLINK(arg_node)->type;
    if (inferred_ty != TY_FUNDAMENTAL_TYPE(declared_ty))
    {
        ERR("Global definition of type `%s' cannot be assigned a value of type `%s'.",
            type_to_str(declared_ty), type_to_str(inferred_ty));
    }

    arg_info->array_dimensions = old_array_dims;

    DBUG_RETURN(arg_node);
}

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

    if (GET_INFERRED_TYPE(CONDEXPR_PREDEXPR) != TY_bool)
        ERR("Predicate of conditional expression is not of type bool.");

    type_t true_expr_ty  = GET_INFERRED_TYPE(CONDEXPR_TRUEEXPR);
    type_t false_expr_ty = GET_INFERRED_TYPE(CONDEXPR_FALSEEXPR);

    if (true_expr_ty != false_expr_ty)
        ERR("Result expressions in conditional expression are of different types.");

    CONDEXPR_TYPE(arg_node) = true_expr_ty;

    DBUG_RETURN(arg_node);
}

node* CTarrexpr(node* arg_node, info* arg_info)
{
    int old_dims = arg_info->array_dimensions;
    int further_dims = -1;

    type_t first_ty = TY_unknown;

    for (node* exprs = ARREXPR_EXPRS(arg_node); exprs != NULL; exprs = EXPRS_NEXT(exprs))
    {
        arg_info->array_dimensions = 0;
        type_t expr_type = GET_INFERRED_TYPE_NODE(EXPRS_EXPR(exprs));
        if (further_dims == -1) further_dims = arg_info->array_dimensions;
        else if (further_dims != arg_info->array_dimensions)
        {
            ERR("Array initializer has no homogenous depth.");
        }
        if (first_ty == TY_unknown) first_ty = expr_type;
        if (expr_type != first_ty)
            ERR("Array initializer has no homogenous type.");
        if (TY_IS_ARRAY(expr_type))
            ERR("Array initializer contains array variables (only fundamental types are allowed).");
    }

    arg_info->array_dimensions = old_dims + further_dims + 1;
    arg_info->inferred_type = first_ty;
    return arg_node;
}

/*
 * Traversal start function
 */

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

    DBUG_ENTER("CTdoCheckTypes");

    arg_info     = info_init();
    arg_info->st = PROGRAM_SYMBOLTABLE(syntaxtree);

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

    arg_info = info_free(arg_info);

    DBUG_RETURN(syntaxtree);
}
