#include "spl_ast.h"

const char *spl_ast_kind_name(spl_ast_node_kind_t kind) {
    static const char *const names[] = {
#define X(a, b, c) #c,
        SPL_AST_KIND_TABLE
#undef X
    };
    if ((usize)kind < sizeof(names) / sizeof(names[0]))
        return names[kind];
    return "?";
}

const char *spl_binary_op_name(spl_binary_op_t op) {
    static const char *const names[] = {
#define X(a, b, c, name) #name,
        SPL_BINARY_OP_TABLE
#undef X
    };
    if ((usize)op < sizeof(names) / sizeof(names[0]))
        return names[op];
    return "?";
}

const char *spl_unary_op_name(spl_unary_op_t op) {
    switch (op) {
    case SPL_UNARY_OP_MINUS:
        return "neg";
    case SPL_UNARY_OP_NOT:
        return "not";
    case SPL_UNARY_OP_BIT_NOT:
        return "bitnot";
    case SPL_UNARY_OP_ADDRESS:
        return "addr";
    case SPL_UNARY_OP_DEREF:
        return "deref";
    default:
        return "?";
    }
}

void spl_ast_init(spl_ast_t *ast, const spl_tok_vec_t *tok_vec) {
    ast->input = *tok_vec;
    vec_init(ast->node_buckets);
    vec_init(ast->str_buckets);
    ast->root = 0;

    spl_ast_node_t n = {0};
    vec_push(ast->node_buckets, n);
}

void spl_ast_drop(spl_ast_t *ast) {
    if (!ast)
        return;
    vec_for(ast->str_buckets, idx) { free(vec_at(ast->str_buckets, idx)); }
    vec_free(ast->str_buckets);
    vec_for(ast->node_buckets, idx) { vec_free(vec_at(ast->node_buckets, idx).attrs); }
    vec_free(ast->node_buckets);
    vec_free(ast->input);
    ast->root = 0;
}

spl_ast_node_t *spl_ast_node(spl_ast_t *ast, spl_ast_node_ref_t node) {
    if (!ast || node == 0 || node >= ast->node_buckets.size)
        return NULL;
    return &ast->node_buckets.data[node];
}

typedef VEC(char) dump_stack_t;

static void dump_node(spl_ast_t *ast, spl_ast_node_ref_t node_ref, dump_stack_t *stack, bool last);

void spl_ast_dump(spl_ast_t *ast, spl_ast_node_ref_t node) {
    if (!ast || !node || node >= ast->node_buckets.size)
        return;
    dump_stack_t stack;
    vec_init(stack);
    printf("AST:\n");
    dump_node(ast, node, &stack, true);
    vec_free(stack);
}

#define dump_vec(_ast, _vec, _stack)                                                               \
    do {                                                                                           \
        vec_for(_vec, _i) { dump_node(_ast, vec_at(_vec, _i), _stack, _i + 1 == vec_size(_vec)); } \
    } while (0)

static void dump_node(spl_ast_t *ast, spl_ast_node_ref_t node_ref, dump_stack_t *stack, bool last) {
    spl_ast_node_t *node = &ast->node_buckets.data[node_ref];

    for (usize i = 0; i < vec_size(*stack); i++)
        printf("%s", vec_at(*stack, i) ? "  " : "| ");
    printf("%s%s #%zu", last ? "`-" : "|-", spl_ast_kind_name(node->kind), node_ref);
    if (node->dbg.fname)
        printf(" [%d:%d]", node->dbg.line, node->dbg.col);
    printf("\n");

    vec_push(*stack, last);

    switch (node->kind) {
    case SPL_AST_ROOT:
        if (node->root.members)
            dump_node(ast, node->root.members, stack, true);
        break;

    case SPL_AST_CONTAINER_MEMBERS:
        dump_vec(ast, node->container_members.external_tags, stack);
        dump_vec(ast, node->container_members.declarations, stack);
        break;

    case SPL_AST_TAG_ITEM:
    case SPL_AST_EXTERNAL_TAG_ITEM:
        dump_vec(ast, node->tag_item.args, stack);
        break;

    case SPL_AST_FN_DECL:
        dump_vec(ast, node->fn_decl.tag_list, stack);
        dump_vec(ast, node->fn_decl.params, stack);
        if (node->fn_decl.ret_type)
            dump_node(ast, node->fn_decl.ret_type, stack, false);
        if (node->fn_decl.body)
            dump_node(ast, node->fn_decl.body, stack, true);
        break;

    case SPL_AST_VAR_DECL:
    case SPL_AST_CONST_DECL:
        dump_vec(ast, node->var_const_decl.attr_list, stack);
        if (node->var_const_decl.pattern)
            dump_node(ast, node->var_const_decl.pattern, stack, false);
        if (node->var_const_decl.type)
            dump_node(ast, node->var_const_decl.type, stack, false);
        if (node->var_const_decl.init_expr)
            dump_node(ast, node->var_const_decl.init_expr, stack, true);
        break;

    case SPL_AST_TYPE_DECL:
        dump_vec(ast, node->type_decl.attr_list, stack);
        if (node->type_decl.pattern)
            dump_node(ast, node->type_decl.pattern, stack, false);
        if (node->type_decl.type_expr)
            dump_node(ast, node->type_decl.type_expr, stack, true);
        break;

    case SPL_AST_MEMBER_DECL:
        dump_vec(ast, node->member_decl.attr_list, stack);
        if (node->member_decl.type_expr)
            dump_node(ast, node->member_decl.type_expr, stack, false);
        if (node->member_decl.init_expr)
            dump_node(ast, node->member_decl.init_expr, stack, true);
        break;

    case SPL_AST_PARAM_DECL:
        dump_vec(ast, node->param_decl.attr_list, stack);
        if (node->param_decl.type_expr)
            dump_node(ast, node->param_decl.type_expr, stack, true);
        break;

    case SPL_AST_COMPTIME_DECL:
        if (node->comptime_decl.body)
            dump_node(ast, node->comptime_decl.body, stack, true);
        break;

    case SPL_AST_TOP_LEVEL_BLOCK:
        dump_vec(ast, node->top_level_block.tag_list, stack);
        if (node->top_level_block.block)
            dump_node(ast, node->top_level_block.block, stack, true);
        break;

    case SPL_AST_STATEMENT_RET:
        if (node->statement_ret.expr)
            dump_node(ast, node->statement_ret.expr, stack, true);
        break;

    case SPL_AST_STATEMENT_BREAK:
        if (node->statement_break.label)
            dump_node(ast, node->statement_break.label, stack, false);
        if (node->statement_break.expr)
            dump_node(ast, node->statement_break.expr, stack, true);
        break;

    case SPL_AST_STATEMENT_CONTINUE:
        if (node->statement_continue.label)
            dump_node(ast, node->statement_continue.label, stack, true);
        break;

    case SPL_AST_STATEMENT_DEFER:
    case SPL_AST_STATEMENT_COMPTIME:
        if (node->statement_defer.body)
            dump_node(ast, node->statement_defer.body, stack, true);
        break;

    case SPL_AST_STATEMENT_DECL:
        if (node->statement_decl.decl)
            dump_node(ast, node->statement_decl.decl, stack, true);
        break;

    case SPL_AST_STATEMENT_EXPR:
        if (node->statement_expr.expr)
            dump_node(ast, node->statement_expr.expr, stack, true);
        break;

    case SPL_AST_AGGREGATE_MEMBER:
        dump_vec(ast, node->aggregate_member.attr_list, stack);
        if (node->aggregate_member.type_expr)
            dump_node(ast, node->aggregate_member.type_expr, stack, false);
        if (node->aggregate_member.init_expr)
            dump_node(ast, node->aggregate_member.init_expr, stack, true);
        break;

    case SPL_AST_STRUCT_INIT_ITEM:
        if (node->struct_init_item.expr)
            dump_node(ast, node->struct_init_item.expr, stack, true);
        break;

    case SPL_AST_MATCH_ARM:
        if (node->match_arm.pattern)
            dump_node(ast, node->match_arm.pattern, stack, false);
        if (node->match_arm.body)
            dump_node(ast, node->match_arm.body, stack, true);
        break;

    case SPL_AST_LABEL:
        break;

    case SPL_AST_EXPR_INT_LIT:
    case SPL_AST_EXPR_FLOAT_LIT:
    case SPL_AST_EXPR_CHAR_LIT:
    case SPL_AST_EXPR_STRING_LIT:
    case SPL_AST_EXPR_BOOL_LIT:
    case SPL_AST_EXPR_NULL_LIT:
    case SPL_AST_EXPR_UNDEFINED_LIT:
    case SPL_AST_EXPR_IDENT:
        break;

    case SPL_AST_EXPR_UNARY:
        if (node->unary.operand)
            dump_node(ast, node->unary.operand, stack, true);
        break;

    case SPL_AST_EXPR_BINARY:
        if (node->binary.left)
            dump_node(ast, node->binary.left, stack, false);
        if (node->binary.right)
            dump_node(ast, node->binary.right, stack, true);
        break;

    case SPL_AST_EXPR_CALL:
        if (node->call.callee)
            dump_node(ast, node->call.callee, stack, false);
        dump_vec(ast, node->call.args, stack);
        break;

    case SPL_AST_EXPR_FIELD:
        if (node->field.obj)
            dump_node(ast, node->field.obj, stack, true);
        break;

    case SPL_AST_EXPR_INDEX:
        if (node->index.obj)
            dump_node(ast, node->index.obj, stack, false);
        if (node->index.index)
            dump_node(ast, node->index.index, stack, true);
        break;

    case SPL_AST_EXPR_SLICE:
        if (node->slice.obj)
            dump_node(ast, node->slice.obj, stack, false);
        if (node->slice.begin)
            dump_node(ast, node->slice.begin, stack, false);
        if (node->slice.end)
            dump_node(ast, node->slice.end, stack, true);
        break;

    case SPL_AST_EXPR_CAST:
        if (node->cast.expr)
            dump_node(ast, node->cast.expr, stack, false);
        if (node->cast.type)
            dump_node(ast, node->cast.type, stack, true);
        break;

    case SPL_AST_EXPR_TUPLE:
        dump_vec(ast, node->tuple.elems, stack);
        break;

    case SPL_AST_EXPR_ARRAY_LIT:
        if (node->array_literal.elem_type)
            dump_node(ast, node->array_literal.elem_type, stack, false);
        dump_vec(ast, node->array_literal.elems, stack);
        break;

    case SPL_AST_EXPR_AGGREGATE_LIT:
        dump_vec(ast, node->aggregate_literal.fields, stack);
        break;

    case SPL_AST_EXPR_COMPTIME_CALL:
        dump_vec(ast, node->comptime_call.args, stack);
        break;

    case SPL_AST_EXPR_BLOCK:
        dump_vec(ast, node->block.stmts, stack);
        if (node->block.tail_expr)
            dump_node(ast, node->block.tail_expr, stack, true);
        break;

    case SPL_AST_EXPR_IF:
        if (node->if_expr.cond)
            dump_node(ast, node->if_expr.cond, stack, false);
        if (node->if_expr.then_block)
            dump_node(ast, node->if_expr.then_block, stack, false);
        if (node->if_expr.else_branch)
            dump_node(ast, node->if_expr.else_branch, stack, true);
        break;

    case SPL_AST_EXPR_IFVAR:
        if (node->ifvar_expr.var_decl)
            dump_node(ast, node->ifvar_expr.var_decl, stack, false);
        if (node->ifvar_expr.then_block)
            dump_node(ast, node->ifvar_expr.then_block, stack, false);
        if (node->ifvar_expr.else_branch)
            dump_node(ast, node->ifvar_expr.else_branch, stack, true);
        break;

    case SPL_AST_EXPR_WHILE:
        if (node->while_expr.label)
            dump_node(ast, node->while_expr.label, stack, false);
        if (node->while_expr.cond)
            dump_node(ast, node->while_expr.cond, stack, false);
        if (node->while_expr.body)
            dump_node(ast, node->while_expr.body, stack, false);
        if (node->while_expr.else_branch)
            dump_node(ast, node->while_expr.else_branch, stack, true);
        break;

    case SPL_AST_EXPR_LOOP:
        if (node->loop_expr.label)
            dump_node(ast, node->loop_expr.label, stack, false);
        if (node->loop_expr.body)
            dump_node(ast, node->loop_expr.body, stack, true);
        break;

    case SPL_AST_EXPR_FOR:
        if (node->for_expr.label)
            dump_node(ast, node->for_expr.label, stack, false);
        dump_vec(ast, node->for_expr.iter_exprs, stack);
        if (node->for_expr.pattern)
            dump_node(ast, node->for_expr.pattern, stack, false);
        if (node->for_expr.body)
            dump_node(ast, node->for_expr.body, stack, false);
        if (node->for_expr.else_branch)
            dump_node(ast, node->for_expr.else_branch, stack, true);
        break;

    case SPL_AST_EXPR_MATCH:
        if (node->match_expr.matched_expr)
            dump_node(ast, node->match_expr.matched_expr, stack, false);
        dump_vec(ast, node->match_expr.arms, stack);
        break;

    case SPL_AST_EXPR_FN_TYPE:
        dump_vec(ast, node->fn_type.params, stack);
        if (node->fn_type.ret_type)
            dump_node(ast, node->fn_type.ret_type, stack, true);
        break;

    case SPL_AST_TYPE_PATH:
        break;

    case SPL_AST_TYPE_POINTER:
        if (node->type_pointer.pointee)
            dump_node(ast, node->type_pointer.pointee, stack, true);
        break;

    case SPL_AST_TYPE_ARRAY:
        if (node->type_array.elem_type)
            dump_node(ast, node->type_array.elem_type, stack, true);
        break;

    case SPL_AST_TYPE_SLICE:
        if (node->type_slice.elem_type)
            dump_node(ast, node->type_slice.elem_type, stack, true);
        break;

    case SPL_AST_TYPE_FN:
        dump_vec(ast, node->type_fn.params, stack);
        if (node->type_fn.ret_type)
            dump_node(ast, node->type_fn.ret_type, stack, true);
        break;

    case SPL_AST_TYPE_PAREN:
        dump_vec(ast, node->type_paren.types, stack);
        break;

    case SPL_AST_TYPE_AGGREGATE:
        dump_vec(ast, node->type_aggregate.members, stack);
        break;

    case SPL_AST_TYPE_COMPTIME_CALL:
        dump_vec(ast, node->type_comptime_call.args, stack);
        break;

    case SPL_AST_PATTERN_IDENT:
    case SPL_AST_PATTERN_UNDERSCORE:
        break;

    case SPL_AST_PATTERN_DOT:
        if (node->pattern_dot.sub_pattern)
            dump_node(ast, node->pattern_dot.sub_pattern, stack, true);
        break;

    case SPL_AST_PATTERN_TUPLE:
        dump_vec(ast, node->pattern_tuple.patterns, stack);
        break;

    case SPL_AST_PATTERN_LIST:
        dump_vec(ast, node->pattern_list.patterns, stack);
        break;

    case SPL_AST_PATTERN_AGGREGATE:
        dump_vec(ast, node->pattern_aggregate.fields, stack);
        break;

    case SPL_AST_PATTERN_FIELD:
        if (node->pattern_field.pattern)
            dump_node(ast, node->pattern_field.pattern, stack, true);
        break;

    default:
        break;
    }

    (void)vec_pop(*stack);
}