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

typedef struct ASTNode ASTNode;

// 上下文结构，记录生成过程中的状态
typedef struct {
    ir_func_t* cur_func;    // 当前处理的函数
    ir_bblock_t* cur_block; // 当前基本块
} 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) {
    node->name = NULL;
    node->type = NULL;
    vector_init(node->used_by);
}

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

static inline ir_bblock_t* new_irbblock(const char* name) {
    ir_bblock_t* block = xmalloc(sizeof(ir_bblock_t));
    block->label = name;
    vector_init(block->instrs);
    return block;
}

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

ir_node_t* emit_br(ir_node_t* cond, ir_bblock_t* trueb, ir_bblock_t* falseb) {
    ir_node_t* br = emit_instr(NULL);
    *br = (ir_node_t) {
        .tag = IR_NODE_BRANCH,
        .data.branch = {
            .cond = cond,
            .true_bblock = trueb,
            .false_bblock = falseb,
        }
    };
    return br;
}

ir_node_t* gen_ir_expr(ASTNode* node) {
    switch (node->type) {
        case NT_TERM_VAL: {
            ir_node_t* ir = new_irnode();
            *ir = (ir_node_t) {
                .tag = IR_NODE_CONST_INT,
                .data.const_int = {
                    .val = node->syms.tok.val.i,
                },
            };
            return ir;
        }
        case NT_TERM_IDENT: {
            ir_node_t* decl = node->syms.decl_node->decl_val.data;
            return decl;
        }
        case NT_TERM_CALL: {
            ir_node_t* ir = emit_instr(NULL);
            *ir = (ir_node_t) {
                .tag = IR_NODE_CALL,
                .data.call = {
                    .callee = node->call.func_decl->decl_func.def->func.data,
                },
            };
            vector_init(ir->data.call.args);
            for (int i = 0; i < node->call.params->params.params.size; i++) {
                vector_push(ir->data.call.args, \
                    gen_ir_expr(node->call.params->params.params.data[i]));
            }
            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;
}
static ir_func_t* new_irfunc(const char* name) {
    ir_func_t *func = xmalloc(sizeof(ir_func_t));
    
    vector_init(func->bblocks);
    vector_init(func->params);
    *func = (ir_func_t) {
        .name = name,
        // TODO typing system
        .type = &type_i32,
    };
    return func;
}

static void gen_ir_func(ASTNode* node, ir_func_t* func) {
    assert(node->type == NT_FUNC);
    ir_bblock_t *entry = new_irbblock("entry");
    vector_push(func->bblocks, entry);
    
    vector_push(prog.funcs, func);
    IRGenContext prev_ctx = ctx;
    ctx.cur_func = func;
    ctx.cur_block = entry;
    
    ast_node_t* params = node->func.decl->decl_func.params;
    for (int i = 0; i < params->params.params.size; i ++) {
        ir_node_t* decl = emit_instr(entry);
        ast_node_t* param = params->params.params.data[i];
        vector_push(func->params, decl);
        *decl = (ir_node_t) {
            .tag = IR_NODE_ALLOC,
            .name = param->decl_val.name->syms.tok.val.str,
            .type = &type_i32,
        };
        param->decl_val.data = decl;
    }
    gen_ir_from_ast(node->func.body);
    
    ctx = prev_ctx;
}

void gen_ir_from_ast(struct ASTNode* node) {
    switch (node->type) {
        case NT_ROOT: {
            for (int i = 0; i < node->root.children.size; i ++) {
                gen_ir_from_ast(node->root.children.data[i]);
            }
            break;
        }
        case NT_DECL_FUNC: {
            ir_func_t* func = new_irfunc(node->decl_func.name->syms.tok.val.str);
            if (node->decl_func.def == NULL) {
                ast_node_t* def = new_ast_node();
                def->func.body = NULL;
                def->func.decl = node;
                node->decl_func.def = def;
                vector_push(prog.extern_funcs, func);
            }
            node->decl_func.def->func.data = func;
            break;
        }
        case NT_FUNC: {
            gen_ir_func(node, node->func.data);
            break;
        }
        case NT_STMT_RETURN: {
            ir_node_t* ret = NULL;
            if (node->return_stmt.expr_stmt != NULL) {
                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,
                    }
                }
            };
            
            vector_push(ctx.cur_func->bblocks, new_irbblock(NULL));
            break;
        }
        case NT_STMT_BLOCK: {
            gen_ir_from_ast(node->block_stmt.block);
            break;
        }
        case NT_BLOCK: {
            for (int i = 0; i < node->block.children.size; i ++) {
                gen_ir_from_ast(node->block.children.data[i]);
            }
            break;
        }
        case NT_STMT_IF: {
            ir_node_t *cond = gen_ir_expr(node->if_stmt.cond);
            ir_bblock_t* trueb = new_irbblock("true_block");
            ir_bblock_t* falseb = new_irbblock("false_block");
            emit_br(cond, trueb, falseb);

            vector_push(ctx.cur_func->bblocks, trueb);
            ctx.cur_block = trueb;
            gen_ir_from_ast(node->if_stmt.if_stmt);
            ir_node_t* jmp = emit_instr(NULL);
            
            if (node->if_stmt.else_stmt != NULL) {
                vector_push(ctx.cur_func->bblocks, falseb);
                ctx.cur_block = falseb;
                gen_ir_from_ast(node->if_stmt.else_stmt);
                ir_node_t* jmp = emit_instr(NULL);

                ctx.cur_block = new_irbblock("jmp_block");
                vector_push(ctx.cur_func->bblocks, ctx.cur_block);
                *jmp = (ir_node_t) {
                    .tag = IR_NODE_JUMP,
                    .data.jump = {
                        .target_bblock = ctx.cur_block,
                    },
                };
            } else {
                ctx.cur_block = falseb;
            }
            *jmp = (ir_node_t) {
                .tag = IR_NODE_JUMP,
                .data.jump = {
                    .target_bblock = ctx.cur_block,
                },
            };
            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.val.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;
    }
}