#include "spl_ast2ir.h"

static spl_ir_node_ref_t spl_ir_const_comptime_int(spl_ir_builder_t *ir, const spl_sema_t *sema,
                                                   usize i) {
    return spl_ir_builder_type_const_int(ir, spl_type_builder_int(sema->type, 0, 0), i);
}

static spl_type_id_t tid_from_def(const spl_sema_t *sema, spl_def_id_t def) {
    spl_def_node_t *d = spl_def_node(sema->type, def);
    return d ? d->type_id : 0;
}

static spl_ir_node_ref_t transit(spl_ast2ir_t *ctx, spl_ast_node_ref_t ref);

static spl_ir_node_ref_t transit_op(spl_ir_builder_t *ir, spl_ast_node_kind_t kind,
                                    spl_type_id_t tid, spl_ir_node_ref_t lref,
                                    spl_ir_node_ref_t rref) {
    spl_ir_node_ref_t ret = 0;
    switch (kind) {
    case SPL_AST_CMP_EQ_EXPR: {
        ret = spl_ir_builder_cmp_eq(ir, tid, lref, rref);
    } break;
    case SPL_AST_CMP_NE_EXPR: {
        ret = spl_ir_builder_cmp_ne(ir, tid, lref, rref);
    } break;
    case SPL_AST_CMP_LE_EXPR: {
        ret = spl_ir_builder_cmp_le(ir, tid, lref, rref);
    } break;
    case SPL_AST_CMP_GE_EXPR: {
        ret = spl_ir_builder_cmp_ge(ir, tid, lref, rref);
    } break;
    case SPL_AST_CMP_LT_EXPR: {
        ret = spl_ir_builder_cmp_lt(ir, tid, lref, rref);
    } break;
    case SPL_AST_CMP_GT_EXPR: {
        ret = spl_ir_builder_cmp_gt(ir, tid, lref, rref);
    } break;
    case SPL_AST_ASSIGN_EXPR: {
        ret = rref;
    } break;
    case SPL_AST_ASSIGN_LSHIFT_EXPR:
    case SPL_AST_LSHIFT_EXPR: {
        ret = spl_ir_builder_arith_shl(ir, tid, lref, rref);
    } break;
    case SPL_AST_ASSIGN_RSHIFT_EXPR:
    case SPL_AST_RSHIFT_EXPR: {
        ret = spl_ir_builder_arith_shr(ir, tid, lref, rref);
    } break;
    case SPL_AST_ASSIGN_ADD_EXPR:
    case SPL_AST_ADD_EXPR: {
        ret = spl_ir_builder_arith_add(ir, tid, lref, rref);
    } break;
    case SPL_AST_SUB_EXPR:
    case SPL_AST_ASSIGN_SUB_EXPR: {
        ret = spl_ir_builder_arith_sub(ir, tid, lref, rref);
    } break;
    case SPL_AST_ASSIGN_MUL_EXPR:
    case SPL_AST_MUL_EXPR: {
        ret = spl_ir_builder_arith_mul(ir, tid, lref, rref);
    } break;
    case SPL_AST_DIV_EXPR:
    case SPL_AST_ASSIGN_DIV_EXPR: {
        ret = spl_ir_builder_arith_div(ir, tid, lref, rref);
    } break;
    case SPL_AST_ASSIGN_MOD_EXPR:
    case SPL_AST_MOD_EXPR: {
        ret = spl_ir_builder_arith_rem(ir, tid, lref, rref);
    } break;
    case SPL_AST_ASSIGN_AND_EXPR:
    case SPL_AST_BIT_AND_EXPR: {
        ret = spl_ir_builder_arith_and(ir, tid, lref, rref);
    } break;
    case SPL_AST_ASSIGN_OR_EXPR:
    case SPL_AST_BIT_OR_EXPR: {
        ret = spl_ir_builder_arith_or(ir, tid, lref, rref);
    } break;
    case SPL_AST_ASSIGN_XOR_EXPR:
    case SPL_AST_BIT_XOR_EXPR: {
        ret = spl_ir_builder_arith_xor(ir, tid, lref, rref);
    } break;
    default:
        UNREACHABLE();
        break;
    }
    return ret;
}

static spl_ir_node_ref_t transit_expr(spl_ast2ir_t *ctx, spl_ast_node_ref_t ref, bool is_lval) {
    spl_ir_builder_t *ir = ctx->ir;
    const spl_sema_t *sema = ctx->sema;
    if (ref == 0)
        return 0;
    Assert(sema != NULL);
    spl_ast_node_t *n = spl_ast_node(sema->ast, ref);
    spl_ir_node_ref_t ret = 0;
    Assert(n != NULL);

    switch (n->kind) {
    case SPL_AST_PACKED_EXPR: {
        TODO();
    } break;
    case SPL_AST_ASSIGN_EXPR: {
        spl_ir_node_ref_t rval = transit_expr(ctx, n->op_expr.right, false);
        spl_ir_node_ref_t lval = transit_expr(ctx, n->op_expr.left, true);
        ret = spl_ir_builder_mem_store(ir, n->resolved_type_id, lval, rval);
    } break;
    case SPL_AST_ASSIGN_ADD_EXPR:
    case SPL_AST_ASSIGN_SUB_EXPR:
    case SPL_AST_ASSIGN_MUL_EXPR:
    case SPL_AST_ASSIGN_DIV_EXPR:
    case SPL_AST_ASSIGN_MOD_EXPR:
    case SPL_AST_ASSIGN_AND_EXPR:
    case SPL_AST_ASSIGN_OR_EXPR:
    case SPL_AST_ASSIGN_XOR_EXPR:
    case SPL_AST_ASSIGN_LSHIFT_EXPR:
    case SPL_AST_ASSIGN_RSHIFT_EXPR: {
        spl_ir_node_ref_t rval = transit_expr(ctx, n->op_expr.right, false);
        spl_ir_node_ref_t lval = transit_expr(ctx, n->op_expr.left, false);
        rval = transit_op(ir, n->kind, n->resolved_type_id, lval, rval);
        lval = transit_expr(ctx, n->op_expr.left, true);
        ret = spl_ir_builder_mem_store(ir, n->resolved_type_id, lval, rval);
    } break;
    case SPL_AST_CMP_EQ_EXPR:
    case SPL_AST_CMP_NE_EXPR:
    case SPL_AST_CMP_LE_EXPR:
    case SPL_AST_CMP_GE_EXPR:
    case SPL_AST_CMP_LT_EXPR:
    case SPL_AST_CMP_GT_EXPR: {
        spl_ir_node_ref_t rval = transit_expr(ctx, n->op_expr.right, false);
        spl_ir_node_ref_t lval = transit_expr(ctx, n->op_expr.left, false);
        ret = transit_op(ir, n->kind, n->resolved_type_id, lval, rval);
    } break;
    case SPL_AST_RANGE_EXPR: {
        TODO();
    } break;
    case SPL_AST_BOOL_OR_EXPR:
    case SPL_AST_BOOL_AND_EXPR: {
        TODO();
    } break;
    case SPL_AST_BIT_OR_EXPR:
    case SPL_AST_BIT_XOR_EXPR:
    case SPL_AST_BIT_AND_EXPR:
    case SPL_AST_LSHIFT_EXPR:
    case SPL_AST_RSHIFT_EXPR:
    case SPL_AST_ADD_EXPR:
    case SPL_AST_SUB_EXPR:
    case SPL_AST_MUL_EXPR:
    case SPL_AST_DIV_EXPR:
    case SPL_AST_MOD_EXPR: {
        spl_ir_node_ref_t rval = transit_expr(ctx, n->op_expr.right, false);
        spl_ir_node_ref_t lval = transit_expr(ctx, n->op_expr.left, false);
        // TODO assert type is bool
        ret = transit_op(ir, n->kind, n->resolved_type_id, lval, rval);
    } break;
    case SPL_AST_MINUS_EXPR: {
        ret = transit_expr(ctx, n->prefix_expr.postfix_expr, false);
        ret = spl_ir_builder_arith_neg(ir, n->resolved_type_id, ret);
    } break;
    case SPL_AST_NOT_EXPR: {
        ret = transit_expr(ctx, n->prefix_expr.postfix_expr, false);
        ret = spl_ir_builder_arith_not(ir, n->resolved_type_id, ret);
        TODO(); // return bool type
    } break;
    case SPL_AST_BIT_NOT_EXPR: {
        ret = transit_expr(ctx, n->prefix_expr.postfix_expr, false);
        ret = spl_ir_builder_arith_not(ir, n->resolved_type_id, ret);
    } break;
    case SPL_AST_ADDRESS_EXPR: {
        TODO();
    } break;
    case SPL_AST_CALL_EXPR: {
        spl_ast_node_t *fn_sym = spl_ast_node(ctx->sema->ast, n->postfix_expr.primary_expr);
        Assert(fn_sym != NULL);
        spl_def_node_t *dn = spl_def_node(sema->type, fn_sym->resolved_def_id);
        Assert(dn != NULL);
        spl_ir_node_ref_t ir_ref = 0;
        if (!map_get(ctx->def2ir_map, fn_sym->resolved_def_id, &ir_ref)) {
            Panic("can't resolved ir map in def2ir");
        }
        Assert(ir_ref != 0);

        VEC(spl_ir_node_ref_t) params;
        vec_init(params);
        vec_for(n->postfix_expr.call_expr, i) {
            spl_ast_node_ref_t param_node = vec_at(n->postfix_expr.call_expr, i);
            spl_ir_node_ref_t ref = transit_expr(ctx, param_node, false);
            vec_push(params, ref);
        }
        spl_ir_node_ref_t call_ref = spl_ir_builder_control_call(ir, n->resolved_type_id, ir_ref);
        vec_for(params, i) { spl_ir_builder_control_call_param(ir, call_ref, vec_at(params, i)); }
        vec_free(params);
        ret = call_ref;
    } break;
    case SPL_AST_FIELD_EXPR:
    case SPL_AST_DEREF_EXPR:
    case SPL_AST_INDEX_EXPR:
    case SPL_AST_SLICE_EXPR:
    case SPL_AST_AS_EXPR:
        TODO();
        break;
    case SPL_AST_EXPR_INTEGER_LIT: {
        ret = spl_ir_builder_type_const_int(ir, n->resolved_type_id, n->primary_expr.integer_expr);
        break;
    }
    case SPL_AST_EXPR_FLOAT_LIT:
    case SPL_AST_EXPR_CHAR_LIT:
    case SPL_AST_EXPR_STRING_LIT: {
        TODO();
    } break;
    case SPL_AST_EXPR_TRUE: {
        return spl_ir_builder_type_const_bool(
            ir, spl_type_builder_simple(sema->type, SPL_TYPE_BOOL), true);
    } break;
    case SPL_AST_EXPR_FALSE: {
        return spl_ir_builder_type_const_bool(
            ir, spl_type_builder_simple(sema->type, SPL_TYPE_BOOL), false);
    }
    case SPL_AST_EXPR_NULL: {
        return spl_ir_builder_type_const_bool(ir, n->resolved_type_id, 0);
    } break;
    case SPL_AST_EXPR_UNDEFINED:
        return 0;
    case SPL_AST_EXPR_IDENT: {
        spl_def_node_t *dn = spl_def_node(sema->type, n->resolved_def_id);
        Assert(dn != NULL);
        spl_ir_node_ref_t ir_ref = 0;
        if (!map_get(ctx->def2ir_map, n->resolved_def_id, &ir_ref)) {
            Panic("can't resolved ir map in def2ir");
        }
        Assert(ir_ref != 0);
        if (is_lval) {
            return ir_ref;
        } else {
            return spl_ir_builder_mem_load(ir, dn->type_id, ir_ref);
        }
    } break;
    case SPL_AST_ARGGREGATE_INIT: {
        vec_for(n->primary_expr.aggregate_init.expr, i) {
            transit_expr(ctx, vec_at(n->primary_expr.aggregate_init.expr, i), false);
        }
    } break;
    case SPL_AST_EXPR_EXPR: {
        transit_expr(ctx, n->primary_expr.expr, is_lval);
    } break;
    case SPL_AST_ARRAY_LIT: {
        TODO();
    } break;
    case SPL_AST_BUILTIN_EXPR: {
        TODO();
    } break;
    case SPL_AST_BLOCK_EXPR: {
        vec_for(n->primary_expr.block_expr, i) {
            transit(ctx, vec_at(n->primary_expr.block_expr, i));
        }
    } break;
    default:
        UNREACHABLE();
        break;
    }
    return ret;
}

static spl_ir_node_ref_t transit(spl_ast2ir_t *ctx, spl_ast_node_ref_t ref) {
    spl_ir_builder_t *ir = ctx->ir;
    const spl_sema_t *sema = ctx->sema;
    if (ref == 0)
        return 0;
    Assert(sema != NULL);
    spl_ast_node_t *n = spl_ast_node(sema->ast, ref);
    spl_ir_node_ref_t ret = 0;
    Assert(n != NULL);

    switch (n->kind) {
    case SPL_AST_NONE:
        UNREACHABLE();
        break;
    case SPL_AST_CONTAINER_MEMBERS: {
        vec_for(n->container_members, i) { transit(ctx, vec_at(n->container_members, i)); }
    } break;
    case SPL_AST_FN_DECL:
        TODO();
        break;
    case SPL_AST_FN_DEFINE: {
        spl_type_id_t tid = tid_from_def(sema, n->resolved_def_id);
        spl_type_node_t *fn_type = spl_type_node(sema->type, tid);
        Assert(fn_type && fn_type->kind == SPL_TYPE_FN);
        // n->fn_decl.attr_list; TODO @extern(vm) using extern attribute
        spl_ir_func_ref_t fn_ref =
            spl_ir_builder_fn_begin(ir, n->fn_decl.name, tid, (spl_ir_attr_t){0} /*TODO*/);
        Assert(fn_ref != 0);
        spl_ir_builder_block_begin(ir);

        map_put(ctx->def2ir_map, n->resolved_def_id, fn_ref);
        vec_for(n->fn_decl.param_list, i) {
            spl_ast_node_ref_t nref = vec_at(n->fn_decl.param_list, i);
            spl_ast_node_t *n = spl_ast_node(sema->ast, nref);
            spl_type_id_t tid = tid_from_def(sema, n->resolved_def_id);
            spl_ir_node_ref_t idx_ref = spl_ir_const_comptime_int(ir, sema, i);
            spl_ir_node_ref_t var_ref = spl_ir_builder_control_param(ir, tid, idx_ref);
            map_put(ctx->def2ir_map, n->resolved_def_id, var_ref);
        }

        transit(ctx, n->fn_decl.block);
        if (!ir->terminated) {
            // FIXME
            spl_ir_builder_control_ret(ir, n->fn_decl.type_expr, 0);
        }
        spl_ir_builder_block_end(ir);
        spl_ir_builder_fn_end(ir);

    } break;
    case SPL_AST_TYPE_DECL: {
        transit(ctx, n->type_decl.type_expr);
    } break;
    case SPL_AST_VAR_DECL: {
        spl_type_id_t tid = tid_from_def(sema, n->resolved_def_id);
        spl_ir_node_ref_t var_ref =
            spl_ir_builder_mem_alloca(ir, tid, spl_ir_const_comptime_int(ir, sema, 1));
        map_put(ctx->def2ir_map, n->resolved_def_id, var_ref);
        if (n->var_const_decl.expr) {
            spl_ir_node_ref_t expr_ref = transit_expr(ctx, n->var_const_decl.expr, false);
            if (expr_ref != 0) { // 0 means undefined
                spl_ir_builder_mem_store(ir, tid, var_ref, expr_ref);
            }
        }
    } break;
    case SPL_AST_CONST_DECL: {
        TODO();
    } break;
    case SPL_AST_MEMBER_DECL:
        TODO();
        break;
    case SPL_AST__COMPTIME_STMT:
    case SPL_AST__DIRECTIVE_BLOCK:
    case SPL_AST_PARAM_DECL: {
        UNREACHABLE();
    } break;
    case SPL_AST_ATTR_ITEM:
    case SPL_AST_ARGG_INIT_ITEM: {
        TODO();
    } break;
    case SPL_AST_IF_STATEMENT: {
        spl_ir_node_ref_t ref = transit_expr(ctx, n->if_statement.expr, false);
        ref = spl_ir_builder_control_br(ir, ref);
        spl_ir_builder_block_end(ir);

        usize true_block = spl_ir_builder_block_begin(ir);
        transit_expr(ctx, n->if_statement.if_block, false);
        usize true_ref = spl_ir_builder_control_jmp(ir);
        spl_ir_builder_block_end(ir);

        usize false_block = spl_ir_builder_block_begin(ir);
        transit_expr(ctx, n->if_statement.else_block, false);
        usize false_ref = spl_ir_builder_control_jmp(ir);
        spl_ir_builder_block_end(ir);

        usize end_block = spl_ir_builder_block_begin(ir);
        spl_ir_builder_control_jmp_fixup(ir, true_ref, end_block);
        spl_ir_builder_control_jmp_fixup(ir, false_ref, end_block);
        spl_ir_builder_control_br_fixup(ir, ref, true_block, false_block);
    } break;
    case SPL_AST_IFVAR_STATEMENT: {
        TODO();
    } break;
    case SPL_AST_WHILE_STATEMENT: {
        spl_ir_node_ref_t end_ref = spl_ir_builder_control_jmp(ir);
        spl_ir_builder_block_end(ir);

        usize cond_block = spl_ir_builder_block_begin(ir);
        spl_ir_node_ref_t cond_expr_ref = transit_expr(ctx, n->while_statement.expr, false);
        spl_ir_node_ref_t cond_ref = spl_ir_builder_control_br(ir, cond_expr_ref);
        spl_ir_builder_block_end(ir);

        usize true_block = spl_ir_builder_block_begin(ir);
        transit_expr(ctx, n->while_statement.while_block, false);
        usize true_ref = spl_ir_builder_control_jmp(ir);
        spl_ir_builder_block_end(ir);

        usize end_block = spl_ir_builder_block_begin(ir);
        spl_ir_builder_control_jmp_fixup(ir, true_ref, cond_block);
        spl_ir_builder_control_br_fixup(ir, cond_ref, true_block, end_block);
    } break;
    case SPL_AST_LOOP_STATEMENT: {
        usize end_ref = spl_ir_builder_control_jmp(ir);
        spl_ir_builder_block_end(ir);

        usize loop_block = spl_ir_builder_block_begin(ir);
        spl_ir_node_ref_t loop_expr_ref = transit_expr(ctx, n->loop_statement.loop_block, false);
        usize loop_ref = spl_ir_builder_control_jmp(ir);
        spl_ir_builder_block_end(ir);

        usize end_block = spl_ir_builder_block_begin(ir);
        spl_ir_builder_control_jmp_fixup(ir, end_ref, loop_block);
        spl_ir_builder_control_jmp_fixup(ir, loop_ref, loop_block);
    } break;
    case SPL_AST_FOR_STATEMENT:
    case SPL_AST_MATCH_STATEMENT:
        TODO();
        break;
    case SPL_AST_RET_STATEMENT: {
        if (n->ret_statement.expr) {
            ret = transit(ctx, n->ret_statement.expr);
        }
        ret = spl_ir_builder_control_ret(ir, n->resolved_type_id, ret);
    } break;
    case SPL_AST_BREAK_STATEMENT:
    case SPL_AST_CONTINUE_STATEMENT:
    case SPL_AST_DEFER_STATEMENT:
    case SPL_AST_TRY_STATEMENT:
    case SPL_AST_CATCH_STATEMENT:
    case SPL_AST_ERRDEFER_STATEMEMT:
    case SPL_AST_EXPR_STATEMENT:
        TODO();
        break;
    case SPL_AST_PACKED_EXPR:
    case SPL_AST_ASSIGN_EXPR:
    case SPL_AST_ASSIGN_ADD_EXPR:
    case SPL_AST_ASSIGN_SUB_EXPR:
    case SPL_AST_ASSIGN_MUL_EXPR:
    case SPL_AST_ASSIGN_DIV_EXPR:
    case SPL_AST_ASSIGN_MOD_EXPR:
    case SPL_AST_ASSIGN_AND_EXPR:
    case SPL_AST_ASSIGN_OR_EXPR:
    case SPL_AST_ASSIGN_XOR_EXPR:
    case SPL_AST_ASSIGN_LSHIFT_EXPR:
    case SPL_AST_ASSIGN_RSHIFT_EXPR:
    case SPL_AST_BOOL_OR_EXPR:
    case SPL_AST_BOOL_AND_EXPR:
    case SPL_AST_BIT_OR_EXPR:
    case SPL_AST_BIT_XOR_EXPR:
    case SPL_AST_BIT_AND_EXPR:
    case SPL_AST_CMP_EQ_EXPR:
    case SPL_AST_CMP_NE_EXPR:
    case SPL_AST_CMP_LE_EXPR:
    case SPL_AST_CMP_GE_EXPR:
    case SPL_AST_CMP_LT_EXPR:
    case SPL_AST_CMP_GT_EXPR:
    case SPL_AST_RANGE_EXPR:
    case SPL_AST_LSHIFT_EXPR:
    case SPL_AST_RSHIFT_EXPR:
    case SPL_AST_ADD_EXPR:
    case SPL_AST_SUB_EXPR:
    case SPL_AST_MUL_EXPR:
    case SPL_AST_DIV_EXPR:
    case SPL_AST_MOD_EXPR:
    case SPL_AST_MINUS_EXPR:
    case SPL_AST_NOT_EXPR:
    case SPL_AST_BIT_NOT_EXPR:
    case SPL_AST_ADDRESS_EXPR:
    case SPL_AST_CALL_EXPR:
    case SPL_AST_FIELD_EXPR:
    case SPL_AST_DEREF_EXPR:
    case SPL_AST_INDEX_EXPR:
    case SPL_AST_SLICE_EXPR:
    case SPL_AST_AS_EXPR:
    case SPL_AST_EXPR_INTEGER_LIT:
    case SPL_AST_EXPR_FLOAT_LIT:
    case SPL_AST_EXPR_CHAR_LIT:
    case SPL_AST_EXPR_STRING_LIT:
    case SPL_AST_EXPR_TRUE:
    case SPL_AST_EXPR_FALSE:
    case SPL_AST_EXPR_NULL:
    case SPL_AST_EXPR_UNDEFINED:
    case SPL_AST_EXPR_IDENT:
    case SPL_AST_ARGGREGATE_INIT:
    case SPL_AST_EXPR_EXPR:
    case SPL_AST_ARRAY_LIT:
    case SPL_AST_BUILTIN_EXPR:
    case SPL_AST_BLOCK_EXPR: {
        return transit_expr(ctx, ref, false);
    }
    case SPL_AST_BASE_TYPE_FN:
    case SPL_AST_BASE_TYPE_PATH:
    case SPL_AST_TYPE_POINTER:
    case SPL_AST_TYPE_ARRAY:
    case SPL_AST_TYPE_SLICE:
    case SPL_AST_TYPE_STRUCT:
    case SPL_AST_TYPE_UNION:
    case SPL_AST_TYPE_SHAPE:
    case SPL_AST_TYPE_ENUM:
    case SPL_AST_TYPE_VOID:
    case SPL_AST_TYPE_BOOL:
    case SPL_AST_TYPE_OPAQUE:
    case SPL_AST_TYPE_I8:
    case SPL_AST_TYPE_U8:
    case SPL_AST_TYPE_I16:
    case SPL_AST_TYPE_U16:
    case SPL_AST_TYPE_I32:
    case SPL_AST_TYPE_U32:
    case SPL_AST_TYPE_I64:
    case SPL_AST_TYPE_U64:
    case SPL_AST_TYPE_ISIZE:
    case SPL_AST_TYPE_USIZE:
    case SPL_AST_TYPE__F32:
    case SPL_AST_TYPE__F64:
    case SPL_AST_TYPE_ANY:
    case SPL_AST_TYPE_IDENT:
    case SPL_AST_COUNT: {
        UNREACHABLE();
        break;
    }
    }
    return ret;
}

static usize hash_int(spl_def_id_t k) { return MAP_HASH_INT(k); }
static int cmp_int(spl_def_id_t k1, spl_def_id_t k2) { return MAP_CMP_INT(k1, k2); }

void spl_ast2ir_init(spl_ast2ir_t *ctx, spl_ir_builder_t *ir, const spl_sema_t *sema) {
    ctx->sema = sema;
    ctx->ir = ir;
    map_init(ctx->def2ir_map, hash_int, cmp_int);
    ctx->err_count = 0;
}

void spl_ast2ir_drop(spl_ast2ir_t *ctx) { map_free(ctx->def2ir_map); }

void spl_ast2ir_run(spl_ast2ir_t *ctx) { transit(ctx, ctx->sema->ast->root); }