#include "ir.h"
#include "../frontend/frontend.h"

typedef struct ASTNode ASTNode;

// 上下文结构，记录生成过程中的状态
typedef struct {
    ir_func_t* current_func;    // 当前处理的函数
    ir_bblock_t* current_block; // 当前基本块
    uint32_t vreg_counter;      // 虚拟寄存器计数器
} IRGenContext;
IRGenContext ctx;
ir_prog_t prog;
ir_type_t type_i32 = {
    .tag = IR_TYPE_INT32,
};
static inline void init_ir_node_t(ir_node_t* node) {
    vector_init(node->used_by);
}

static inline ir_node_t* new_ir_node_t() {
    ir_node_t* node = xmalloc(sizeof(ir_node_t));
    init_ir_node_t(node);
}

ir_node_t* emit_instr(ir_bblock_t* block) {
    if (block == NULL) block = ctx.current_block;
    ir_node_t *node = new_ir_node_t();
    vector_push(block->instrs, node);
    return vector_at(block->instrs, block->instrs.size - 1);
}

void emit_br(ir_node_t cond, const char* true_lable, const char* false_lable) {
    ir_node_t br = {
        .tag = IR_NODE_RET,
        .data = {
        }
    };
    // emit_instr(br, NULL);
}

ir_node_t* gen_ir_expr(ASTNode* node) {
    switch (node->type) {
        case NT_TERM_VAL: {
            ir_node_t* ir = new_ir_node_t();
            *ir = (ir_node_t) {
                .tag = IR_NODE_CONST_INT,
                .data.const_int = {
                    .val = node->syms.tok.constant.i,
                },
            };
            return ir;
        }
        case NT_TERM_IDENT: {
            ir_node_t* decl = node->syms.decl_node->decl_val.data;
            return decl;
        }
        case NT_TERM_CALL: {
            // TODO
            ir_node_t* ir = new_ir_node_t();
            *ir = (ir_node_t) {
                .tag = IR_NODE_CALL,
                .data.call = {
                    .callee = NULL,
                },
            };
            vector_init(ir->data.call.args);
            return ir;
        }
        default:
            goto NEXT;
    }
    return NULL;
NEXT:
    ir_node_t* lhs = gen_ir_expr(node->expr.left);
    ir_node_t* rhs = node->expr.right ? gen_ir_expr(node->expr.right) : NULL;

    if (node->type == NT_COMMA) {
        return rhs;
    }

    ir_node_t* instr = emit_instr(NULL);
    vector_push(lhs->used_by, instr);
    if (rhs) { vector_push(rhs->used_by, instr); }

    ir_node_t* ret;
#define BINOP(operand) do { \
        *instr = (ir_node_t){ \
            .tag = IR_NODE_OP, \
            .data.op = { \
                .op = operand, \
                .lhs = lhs, \
                .rhs = rhs, \
            }, \
        }; \
        ret = instr; \
    } while (0)

    switch (node->type) {
        case NT_ADD     :// (expr) + (expr)
            BINOP(IR_OP_ADD);
            break;
        case NT_SUB     :// (expr) - (expr)
            BINOP(IR_OP_SUB);
            break;
        case NT_MUL     :// (expr) * (expr)
            BINOP(IR_OP_MUL);
            break;
        case NT_DIV     :// (expr) / (expr)
            BINOP(IR_OP_DIV);
            break;
        case NT_MOD     :// (expr) % (expr)
            BINOP(IR_OP_MOD);
            break;
        case NT_AND     :// (expr) & (expr)
            BINOP(IR_OP_AND);
            break;
        case NT_OR      :// (expr) | (expr)
            BINOP(IR_OP_OR);
            break;
        case NT_XOR     :// (expr) ^ (expr)
            BINOP(IR_OP_XOR);
            break;
        case NT_BIT_NOT :// ~ (expr)
            // TODO
            // BINOP(IR_OP_NOT);
            break;
        case NT_L_SH    :// (expr) << (expr)
            BINOP(IR_OP_SHL);
            break;
        case NT_R_SH    :// (expr) >> (expr)
            BINOP(IR_OP_SHR); // Shift right logical.
            // TODO
            // BINOP(IR_OP_SAR); // Shift right arithmetic.
            break;
        case NT_EQ      :// (expr) == (expr)
            BINOP(IR_OP_EQ);
            break;
        case NT_NEQ     :// (expr) != (expr)
            BINOP(IR_OP_NEQ);
            break;
        case NT_LE      :// (expr) <= (expr)
            BINOP(IR_OP_LE);
            break;
        case NT_GE      :// (expr) >= (expr)
            BINOP(IR_OP_GE);
            break;
        case NT_LT      :// (expr) < (expr)
            BINOP(IR_OP_LT);
            break;
        case NT_GT      :// (expr) > (expr)
            BINOP(IR_OP_GE);
            break;
        case NT_AND_AND :// (expr) && (expr)
            break;
        case NT_OR_OR   :// (expr) || (expr)
            break;
        case NT_NOT     :// ! (expr)
            ir_node_t* zero = xmalloc(sizeof(ir_node_t));
            *zero = (ir_node_t){
                .tag = IR_NODE_CONST_INT,
                .data.const_int = {
                    .val = 0,
                },
            };
            *instr = (ir_node_t){
                .tag = IR_NODE_OP,
                .data.op = {
                    .op = IR_OP_EQ,
                    .lhs = zero,
                    .rhs = lhs,
                },
            };
            ret = instr;
            break;
        case NT_ASSIGN  :// (expr) = (expr)
            *instr = (ir_node_t){
                .tag = IR_NODE_STORE,
                .data.store = {
                    .target = lhs,
                    .value = rhs,
                },
            };
            ret = rhs;
            break;
        // case NT_COND    : // (expr) ? (expr) : (expr)
        default:
            // TODO self error msg
            error("Unsupported IR generation for AST node type %d", node->type);
            break;
    }
    return ret;
}

void gen_ir_from_ast(struct ASTNode* node) {
    switch (node->type) {
        case NT_ROOT: {
            for (int i = 0; i < node->root.child_size; i ++) {
                gen_ir_from_ast(node->root.children[i]);
            }
        } break;
        case NT_FUNC: {
            ir_func_t *func = xmalloc(sizeof(ir_func_t));
            *func = (ir_func_t) {
                .name = node->func.name->syms.tok.constant.str,
            };
            vector_init(func->bblocks);

            ir_bblock_t *entry = xmalloc(sizeof(ir_bblock_t));
            *entry = (ir_bblock_t) {
                .label = "entry",
            };
            vector_init(entry->instrs);
            vector_push(func->bblocks, entry);

            IRGenContext prev_ctx = ctx;
            ctx = (IRGenContext) {
                .current_func =  func,
                .current_block = vector_at(func->bblocks, 0),
                .vreg_counter = 0,
            };

            gen_ir_from_ast(node->func.body);

            ctx = prev_ctx;
            vector_push(prog.funcs, func);
        } break;
        case NT_STMT_RETURN: {
            ir_node_t* ret = gen_ir_expr(node->return_stmt.expr_stmt);
            ir_node_t* ir = emit_instr(NULL);
            *ir = (ir_node_t) {
                .tag = IR_NODE_RET,
                .data = {
                    .ret = {
                        .ret_val = ret,
                    }
                }
            };
            break;
        }
        case NT_BLOCK: {
            for (int i = 0; i < node->block.child_size; i ++) {
                gen_ir_from_ast(node->block.children[i]);
            }
            break;
        }
        case NT_STMT_IF: {
            ir_node_t *cond = gen_ir_expr(node->if_stmt.cond);

            // xmalloc();
            // ir_bblock_t then_block = {
            // };
            node->if_stmt.if_stmt;
            node->if_stmt.else_stmt;
            break;
        }
        case NT_STMT_WHILE: {
            node->while_stmt.cond;
            node->while_stmt.body;
            break;
        }
        case NT_STMT_DOWHILE: {
            node->do_while_stmt.cond;
            node->do_while_stmt.body;
            break;
        }
        case NT_STMT_FOR: {
            node->for_stmt.init;
            node->for_stmt.cond;
            node->for_stmt.iter;
            node->for_stmt.body;
            break;
        }
        case NT_DECL_VAR: {
            ir_node_t* ret_node = emit_instr(NULL);
            *ret_node = (ir_node_t) {
                .tag = IR_NODE_ALLOC,
                .name = node->decl_val.name->syms.tok.constant.str,
                .type = &type_i32,
            };
            node->decl_val.data = ret_node;
            if (node->decl_val.expr_stmt != NULL) {
                gen_ir_from_ast(node->decl_val.expr_stmt);
            }
            break;
        }
        case NT_STMT_EXPR: {
            gen_ir_expr(node->expr_stmt.expr_stmt);
            break;
        }
        case NT_STMT_EMPTY: {
            break;
        }
        default:
            // TODO: 错误处理
            error("unknown node type");
            break;
    }
}