#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "codegen.h"
#include "ast.h"

CodeBuffer* codebuf_new(void) {
    CodeBuffer *buf = malloc(sizeof(CodeBuffer));
    buf->capacity = 65536;
    buf->data = malloc(buf->capacity);
    buf->size = 0;
    return buf;
}

void codebuf_free(CodeBuffer *buf) {
    if (!buf) return;
    if (buf->data) free(buf->data);
    free(buf);
}

void codebuf_write_byte(CodeBuffer *buf, unsigned char byte) {
    if (buf->size >= buf->capacity) {
        buf->capacity *= 2;
        buf->data = realloc(buf->data, buf->capacity);
    }
    buf->data[buf->size++] = byte;
}

void codebuf_write_uint32(CodeBuffer *buf, unsigned int val) {
    for (int i = 0; i < 4; i++) codebuf_write_byte(buf, (val >> (i * 8)) & 0xFF);
}

void codebuf_write_uint64(CodeBuffer *buf, unsigned long long val) {
    for (int i = 0; i < 8; i++) codebuf_write_byte(buf, (val >> (i * 8)) & 0xFF);
}

/* ==================== Instruction emitters ==================== */

static void emit_movz(CodeBuffer *buf, int reg, int imm, int shift) {
    codebuf_write_uint32(buf, 0xD2800000 | ((imm & 0xFFFF) << 5) | ((shift / 16) << 21) | reg);
}

static void emit_movk(CodeBuffer *buf, int reg, int imm, int shift) {
    codebuf_write_uint32(buf, 0xF2800000 | ((imm & 0xFFFF) << 5) | ((shift / 16) << 21) | reg);
}

static void emit_mov_imm64(CodeBuffer *buf, int reg, unsigned long long imm) {
    if (imm == 0) {
        codebuf_write_uint32(buf, 0xAA1F03E0 | reg);
        return;
    }
    emit_movz(buf, reg, imm & 0xFFFF, 0);
    if (imm >> 16) emit_movk(buf, reg, (imm >> 16) & 0xFFFF, 16);
    if (imm >> 32) emit_movk(buf, reg, (imm >> 32) & 0xFFFF, 32);
    if (imm >> 48) emit_movk(buf, reg, (imm >> 48) & 0xFFFF, 48);
}

static void emit_mov_reg(CodeBuffer *buf, int rd, int rm) {
    codebuf_write_uint32(buf, 0xAA000000 | (rm << 16) | (31 << 5) | rd);
}

static void emit_add(CodeBuffer *buf, int rd, int rn, int rm) {
    codebuf_write_uint32(buf, 0x8B000000 | (rm << 16) | (rn << 5) | rd);
}

static void emit_sub(CodeBuffer *buf, int rd, int rn, int rm) {
    codebuf_write_uint32(buf, 0xCB000000 | (rm << 16) | (rn << 5) | rd);
}

static void emit_mul(CodeBuffer *buf, int rd, int rn, int rm) {
    codebuf_write_uint32(buf, 0x9B007C00 | (rm << 16) | (rn << 5) | rd);
}

static void emit_sdiv(CodeBuffer *buf, int rd, int rn, int rm) {
    codebuf_write_uint32(buf, 0x9AC00C00 | (rm << 16) | (rn << 5) | rd);
}

static void emit_msub(CodeBuffer *buf, int rd, int rn, int rm, int ra) {
    codebuf_write_uint32(buf, 0x9B008000 | (rm << 16) | (ra << 10) | (rn << 5) | rd);
}

static void emit_neg(CodeBuffer *buf, int rd, int rm) {
    codebuf_write_uint32(buf, 0xCB000000 | (rm << 16) | (31 << 5) | rd);
}

static void emit_str_imm(CodeBuffer *buf, int rt, int rn, int offset) {
    codebuf_write_uint32(buf, 0xF9000000 | ((offset / 8) << 10) | (rn << 5) | rt);
}

static void emit_ldr_imm(CodeBuffer *buf, int rt, int rn, int offset) {
    codebuf_write_uint32(buf, 0xF9400000 | ((offset / 8) << 10) | (rn << 5) | rt);
}

static void emit_strb_imm(CodeBuffer *buf, int rt, int rn, int offset) {
    codebuf_write_uint32(buf, 0x39000000 | ((offset & 0xFFF) << 10) | (rn << 5) | rt);
}

static void emit_ldrb_imm(CodeBuffer *buf, int rt, int rn, int offset) {
    codebuf_write_uint32(buf, 0x39400000 | ((offset & 0xFFF) << 10) | (rn << 5) | rt);
}

static void emit_ldrb_reg(CodeBuffer *buf, int rt, int rn, int rm) {
    codebuf_write_uint32(buf, 0x38606800 | (rm << 16) | (rn << 5) | rt);
}

static void emit_adr(CodeBuffer *buf, int reg, int offset) {
    int immlo = offset & 0x3;
    int immhi = (offset >> 2) & 0x7FFFF;
    codebuf_write_uint32(buf, 0x10000000 | (immlo << 29) | (immhi << 5) | reg);
}

static void emit_svc(CodeBuffer *buf) {
    codebuf_write_uint32(buf, 0xD4000001);
}

static void emit_ret(CodeBuffer *buf) {
    codebuf_write_uint32(buf, 0xD65F03C0);
}

static void emit_cmp_reg(CodeBuffer *buf, int rn, int rm) {
    codebuf_write_uint32(buf, 0xEB000000 | (rm << 16) | (rn << 5) | 31);
}

static void emit_cmp_imm(CodeBuffer *buf, int rn, int imm) {
    if (imm >= 0 && imm <= 4095) {
        codebuf_write_uint32(buf, 0xF100001F | ((imm & 0xFFF) << 10) | (rn << 5));
    } else {
        emit_mov_imm64(buf, 15, (unsigned long long)imm);
        emit_cmp_reg(buf, rn, 15);
    }
}

static void emit_cset(CodeBuffer *buf, int rd, int cond) {
    static const int inv[] = {1,0,3,2,5,4,7,6,9,8,11,10,13,12,14,15};
    int c = inv[cond & 0xF];
    codebuf_write_uint32(buf, 0x9A800400 | (c << 12) | (31 << 16) | (31 << 5) | rd);
}


/* ==================== Label system ==================== */

#define MAX_LABELS 512
#define MAX_PATCHES 2048

typedef struct {
    int defined;
    int offset;
    int patches[MAX_PATCHES];
    int patch_count;
    int is_cond;
    int cond;
} Label;

static Label g_labels[MAX_LABELS];
static int g_label_count = 0;

static int label_new(void) {
    int id = g_label_count++;
    g_labels[id].defined = 0;
    g_labels[id].offset = 0;
    g_labels[id].patch_count = 0;
    return id;
}

static void label_define(CodeBuffer *buf, int id) {
    g_labels[id].defined = 1;
    g_labels[id].offset = buf->size;
    for (int i = 0; i < g_labels[id].patch_count; i++) {
        int pos = g_labels[id].patches[i];
        int instr_offset = pos;
        int rel = (g_labels[id].offset - instr_offset) / 4;
        unsigned int instr;
        memcpy(&instr, buf->data + pos, 4);
        if ((instr & 0xFC000000) == 0x94000000) {
            instr = 0x94000000 | (rel & 0x3FFFFFF);
        } else if ((instr & 0xFC000000) == 0x14000000) {
            instr = 0x14000000 | (rel & 0x3FFFFFF);
        } else if ((instr & 0xFF000000) == 0x34000000) {
            instr = (instr & 0xFF00001F) | ((rel & 0x7FFFF) << 5);
        } else if ((instr & 0xFF000000) == 0x35000000) {
            instr = (instr & 0xFF00001F) | ((rel & 0x7FFFF) << 5);
        } else {
            instr = 0x54000000 | ((rel & 0x7FFFF) << 5) | (instr & 0xF);
        }
        memcpy(buf->data + pos, &instr, 4);
    }
    g_labels[id].patch_count = 0;
}

static void emit_b(CodeBuffer *buf, int label_id) {
    if (g_labels[label_id].defined) {
        int rel = (g_labels[label_id].offset - buf->size) / 4;
        codebuf_write_uint32(buf, 0x14000000 | (rel & 0x3FFFFFF));
    } else {
        g_labels[label_id].patches[g_labels[label_id].patch_count++] = buf->size;
        codebuf_write_uint32(buf, 0x14000000);
    }
}

static void emit_b_cond(CodeBuffer *buf, int cond, int label_id) {
    if (g_labels[label_id].defined) {
        int rel = (g_labels[label_id].offset - buf->size) / 4;
        codebuf_write_uint32(buf, 0x54000000 | ((rel & 0x7FFFF) << 5) | cond);
    } else {
        g_labels[label_id].patches[g_labels[label_id].patch_count++] = buf->size;
        codebuf_write_uint32(buf, 0x54000000 | cond);
    }
}

static void emit_cbz(CodeBuffer *buf, int rt, int label_id) {
    if (g_labels[label_id].defined) {
        int rel = (g_labels[label_id].offset - buf->size) / 4;
        codebuf_write_uint32(buf, 0x34000000 | ((rel & 0x7FFFF) << 5) | rt);
    } else {
        g_labels[label_id].patches[g_labels[label_id].patch_count++] = buf->size;
        codebuf_write_uint32(buf, 0x34000000 | rt);
    }
}

static void emit_cbnz(CodeBuffer *buf, int rt, int label_id) {
    if (g_labels[label_id].defined) {
        int rel = (g_labels[label_id].offset - buf->size) / 4;
        codebuf_write_uint32(buf, 0x35000000 | ((rel & 0x7FFFF) << 5) | rt);
    } else {
        g_labels[label_id].patches[g_labels[label_id].patch_count++] = buf->size;
        codebuf_write_uint32(buf, 0x35000000 | rt);
    }
}

static void emit_bl(CodeBuffer *buf, int label_id) {
    if (g_labels[label_id].defined) {
        int rel = (g_labels[label_id].offset - buf->size) / 4;
        codebuf_write_uint32(buf, 0x94000000 | (rel & 0x3FFFFFF));
    } else {
        g_labels[label_id].patches[g_labels[label_id].patch_count++] = buf->size;
        codebuf_write_uint32(buf, 0x94000000);
    }
}

/* ==================== Context ==================== */

typedef struct {
    CodeBuffer *text;
    CodeBuffer *rodata;
    CodeBuffer *data;
    VarInfo *vinfo;
    int opt_level;
    int lbl_itoa;
    int lbl_strcmp;
    int lbl_strlen;
    int lbl_exec;
    int lbl_filetest;
    int text_total_size;
} Ctx;

static int rodata_add_string(CodeBuffer *rodata, const char *s, int add_newline) {
    int offset = rodata->size;
    for (int i = 0; s[i]; i++) codebuf_write_byte(rodata, (unsigned char)s[i]);
    if (add_newline) codebuf_write_byte(rodata, '\n');
    codebuf_write_byte(rodata, '\0');
    return offset;
}

static int var_get_index(Ctx *ctx, const char *name) {
    for (int i = 0; i < ctx->vinfo->var_count; i++) {
        if (strcmp(ctx->vinfo->var_names[i], name) == 0) return i;
    }
    if (ctx->vinfo->var_count < 256) {
        int idx = ctx->vinfo->var_count++;
        ctx->vinfo->var_names[idx] = strdup(name);
        ctx->vinfo->var_types[idx] = VARTYPE_INT;
        ctx->vinfo->var_init[idx] = 0;
        ctx->vinfo->var_str_init[idx] = NULL;
        return idx;
    }
    return -1;
}

static void var_store(Ctx *ctx, int idx) {
    emit_str_imm(ctx->text, 0, 19, idx * 8);
}

static void var_load(Ctx *ctx, int idx) {
    emit_ldr_imm(ctx->text, 0, 19, idx * 8);
}

/* ==================== Expression codegen ==================== */

static void codegen_expr(Ctx *ctx, Expr *e) {
    if (!e) { emit_mov_imm64(ctx->text, 0, 0); return; }
    if (e->type == EXPR_INT) {
        emit_mov_imm64(ctx->text, 0, (unsigned long long)e->int_val);
        return;
    }
    if (e->type == EXPR_VAR) {
        int idx = var_get_index(ctx, e->var_name);
        var_load(ctx, idx);
        return;
    }
    codegen_expr(ctx, e->left);
    emit_mov_reg(ctx->text, 21, 0);
    codegen_expr(ctx, e->right);
    emit_mov_reg(ctx->text, 1, 0);
    emit_mov_reg(ctx->text, 0, 21);
    switch (e->op) {
        case '+': emit_add(ctx->text, 0, 0, 1); break;
        case '-': emit_sub(ctx->text, 0, 0, 1); break;
        case '*': emit_mul(ctx->text, 0, 0, 1); break;
        case '/': emit_sdiv(ctx->text, 0, 0, 1); break;
        case '%':
            emit_sdiv(ctx->text, 2, 0, 1);
            emit_msub(ctx->text, 0, 2, 1, 0);
            break;
    }
}

/* ==================== Test codegen ==================== */

static void codegen_word_str(Ctx *ctx, Word *w) {
    if (w->type == WORD_VAR) {
        int idx = var_get_index(ctx, w->text);
        var_load(ctx, idx);
    } else {
        int off = rodata_add_string(ctx->rodata, w->text, 0);
        int adr_imm = ctx->text_total_size + off - ctx->text->size;
        emit_adr(ctx->text, 0, adr_imm);
    }
}

static void codegen_test(Ctx *ctx, Test *t) {
    if (!t) { emit_mov_imm64(ctx->text, 0, 1); return; }
    if (t->type == TEST_FILE) {
        int off = rodata_add_string(ctx->rodata, t->file_path, 0);
        int adr_imm = ctx->text_total_size + off - ctx->text->size;
        emit_adr(ctx->text, 0, adr_imm);
        emit_mov_imm64(ctx->text, 1, (unsigned long long)t->op);
        emit_bl(ctx->text, ctx->lbl_filetest);
        return;
    }
    if (t->type == TEST_NUM) {
        codegen_expr(ctx, t->num_left);
        emit_mov_reg(ctx->text, 21, 0);
        codegen_expr(ctx, t->num_right);
        emit_mov_reg(ctx->text, 1, 0);
        emit_mov_reg(ctx->text, 0, 21);
        emit_cmp_reg(ctx->text, 0, 1);
        int cond;
        switch (t->op) {
            case 'e': cond = 0; break;
            case 'n': cond = 1; break;
            case 'l': cond = 11; break;
            case 'L': cond = 13; break;
            case 'g': cond = 12; break;
            case 'G': cond = 10; break;
            default: cond = 0;
        }
        emit_cset(ctx->text, 0, cond);
        return;
    }
    if (t->type == TEST_STR) {
        codegen_word_str(ctx, t->str_left);
        emit_mov_reg(ctx->text, 21, 0);
        codegen_word_str(ctx, t->str_right);
        emit_mov_reg(ctx->text, 1, 0);
        emit_mov_reg(ctx->text, 0, 21);
        emit_bl(ctx->text, ctx->lbl_strcmp);
        emit_cmp_imm(ctx->text, 0, 0);
        int cond = (t->op == '=') ? 0 : 1;
        emit_cset(ctx->text, 0, cond);
        return;
    }
}

/* ==================== Statement codegen ==================== */

static void codegen_stmt(Ctx *ctx, Stmt *s);

static void codegen_stmt_list(Ctx *ctx, Stmt *head) {
    for (Stmt *s = head; s; s = s->next) codegen_stmt(ctx, s);
}

static void emit_write_str(Ctx *ctx, const char *s, int add_newline) {
    int len = (int)strlen(s) + (add_newline ? 1 : 0);
    int off = rodata_add_string(ctx->rodata, s, add_newline);
    emit_mov_imm64(ctx->text, 8, 64);
    emit_mov_imm64(ctx->text, 0, 1);
    int adr_imm = ctx->text_total_size + off - ctx->text->size;
    emit_adr(ctx->text, 1, adr_imm);
    emit_mov_imm64(ctx->text, 2, (unsigned long long)len);
    emit_svc(ctx->text);
}

static void codegen_stmt(Ctx *ctx, Stmt *s) {
    switch (s->type) {
        case STMT_ASSIGN: {
            int idx = var_get_index(ctx, s->var_name);
            if (s->arith_val) {
                ctx->vinfo->var_types[idx] = VARTYPE_INT;
                codegen_expr(ctx, s->arith_val);
                var_store(ctx, idx);
            } else {
                ctx->vinfo->var_types[idx] = VARTYPE_STR;
                int off = rodata_add_string(ctx->rodata, s->str_val ? s->str_val : "", 0);
                int adr_imm = ctx->text_total_size + off - ctx->text->size;
                emit_adr(ctx->text, 0, adr_imm);
                var_store(ctx, idx);
            }
            break;
        }
        case STMT_ECHO: {
            for (int i = 0; i < s->echo_argc; i++) {
                Word *w = s->echo_args[i];
                if (i > 0) emit_write_str(ctx, " ", 0);
                if (w->type == WORD_VAR) {
                    int idx = var_get_index(ctx, w->text);
                    if (ctx->vinfo->var_types[idx] == VARTYPE_INT) {
                        var_load(ctx, idx);
                        emit_bl(ctx->text, ctx->lbl_itoa);
                    } else {
                        var_load(ctx, idx);
                        emit_mov_reg(ctx->text, 3, 0);
                        emit_bl(ctx->text, ctx->lbl_strlen);
                        emit_mov_reg(ctx->text, 2, 0);
                        emit_mov_reg(ctx->text, 1, 3);
                        emit_mov_imm64(ctx->text, 8, 64);
                        emit_mov_imm64(ctx->text, 0, 1);
                        emit_svc(ctx->text);
                    }
                } else {
                    emit_write_str(ctx, w->text, 0);
                }
            }
            emit_write_str(ctx, "", 1);
            break;
        }
        case STMT_IF: {
            int else_lbl = label_new();
            int end_lbl = label_new();
            codegen_test(ctx, s->cond);
            emit_cmp_imm(ctx->text, 0, 0);
            emit_b_cond(ctx->text, 0, else_lbl);
            codegen_stmt_list(ctx, s->then_body);
            emit_b(ctx->text, end_lbl);
            label_define(ctx->text, else_lbl);
            codegen_stmt_list(ctx, s->else_body);
            label_define(ctx->text, end_lbl);
            break;
        }
        case STMT_FOR: {
            int loop_lbl = label_new();
            int end_lbl = label_new();
            int idx = var_get_index(ctx, s->for_var);
            ctx->vinfo->var_types[idx] = VARTYPE_INT;
            emit_mov_imm64(ctx->text, 20, 0);
            label_define(ctx->text, loop_lbl);
            emit_cmp_imm(ctx->text, 20, s->for_list_count);
            emit_b_cond(ctx->text, 2, end_lbl);
            int *check_lbls = malloc(s->for_list_count * sizeof(int));
            for (int i = 0; i < s->for_list_count; i++) check_lbls[i] = label_new();
            int body_lbl = label_new();
            for (int i = 0; i < s->for_list_count; i++) {
                if (i > 0) label_define(ctx->text, check_lbls[i - 1]);
                emit_cmp_imm(ctx->text, 20, i);
                if (i < s->for_list_count - 1) {
                    emit_b_cond(ctx->text, 1, check_lbls[i]);
                } else {
                    emit_b_cond(ctx->text, 1, body_lbl);
                }
                emit_mov_imm64(ctx->text, 0, (unsigned long long)atoll(s->for_list[i]));
                var_store(ctx, idx);
                emit_b(ctx->text, body_lbl);
            }
            label_define(ctx->text, body_lbl);
            free(check_lbls);
            codegen_stmt_list(ctx, s->body);
            emit_mov_imm64(ctx->text, 10, 1);
            emit_add(ctx->text, 20, 20, 10);
            emit_b(ctx->text, loop_lbl);
            label_define(ctx->text, end_lbl);
            break;
        }
        case STMT_WHILE: {
            int loop_lbl = label_new();
            int end_lbl = label_new();
            label_define(ctx->text, loop_lbl);
            codegen_test(ctx, s->cond);
            emit_cmp_imm(ctx->text, 0, 0);
            emit_b_cond(ctx->text, 0, end_lbl);
            codegen_stmt_list(ctx, s->body);
            emit_b(ctx->text, loop_lbl);
            label_define(ctx->text, end_lbl);
            break;
        }
        case STMT_EXEC: {
            int argc = s->exec_argc;
            int str_offs[256];
            for (int i = 0; i < argc; i++) {
                str_offs[i] = rodata_add_string(ctx->rodata, s->exec_argv[i], 0);
            }
            int argv_off = ctx->rodata->size;
            for (int i = 0; i < argc; i++) {
                unsigned long long vaddr = 0x401000ULL + ctx->text_total_size + str_offs[i];
                codebuf_write_uint64(ctx->rodata, vaddr);
            }
            codebuf_write_uint64(ctx->rodata, 0);
            int cmd_adr = ctx->text_total_size + str_offs[0] - ctx->text->size;
            emit_adr(ctx->text, 0, cmd_adr);
            int argv_adr = ctx->text_total_size + argv_off - ctx->text->size;
            emit_adr(ctx->text, 1, argv_adr);
            emit_bl(ctx->text, ctx->lbl_exec);
            break;
        }
        case STMT_EXIT: {
            codegen_expr(ctx, s->exit_expr);
            emit_mov_imm64(ctx->text, 8, 93);
            emit_svc(ctx->text);
            break;
        }
    }
}

/* ==================== Runtime subroutines ==================== */

static void emit_itoa_subroutine(Ctx *ctx) {
    label_define(ctx->text, ctx->lbl_itoa);
    codebuf_write_uint32(ctx->text, 0xA9BF7BFD);
    codebuf_write_uint32(ctx->text, 0xD10083FF);
    emit_mov_reg(ctx->text, 4, 0);
    emit_mov_imm64(ctx->text, 5, 0);
    emit_cmp_imm(ctx->text, 4, 0);
    int pos_lbl = label_new();
    emit_b_cond(ctx->text, 10, pos_lbl);
    emit_mov_imm64(ctx->text, 5, 1);
    emit_neg(ctx->text, 4, 4);
    label_define(ctx->text, pos_lbl);
    emit_mov_imm64(ctx->text, 6, 10);
    codebuf_write_uint32(ctx->text, 0x910043E2);
    emit_mov_imm64(ctx->text, 3, 0);
    int loop_lbl = label_new();
    label_define(ctx->text, loop_lbl);
    emit_sdiv(ctx->text, 7, 4, 6);
    emit_msub(ctx->text, 8, 7, 6, 4);
    emit_add(ctx->text, 8, 8, 31);
    emit_mov_imm64(ctx->text, 15, '0');
    emit_add(ctx->text, 8, 8, 15);
    codebuf_write_uint32(ctx->text, 0xD1000442);
    emit_strb_imm(ctx->text, 8, 2, 0);
    emit_add(ctx->text, 3, 3, 31);
    emit_mov_imm64(ctx->text, 15, 1);
    emit_add(ctx->text, 3, 3, 15);
    emit_mov_reg(ctx->text, 4, 7);
    emit_cbnz(ctx->text, 4, loop_lbl);
    emit_cmp_imm(ctx->text, 5, 1);
    int out_lbl = label_new();
    emit_b_cond(ctx->text, 1, out_lbl);
    codebuf_write_uint32(ctx->text, 0xD1000442);
    emit_mov_imm64(ctx->text, 8, '-');
    emit_strb_imm(ctx->text, 8, 2, 0);
    emit_add(ctx->text, 3, 3, 31);
    emit_mov_imm64(ctx->text, 15, 1);
    emit_add(ctx->text, 3, 3, 15);
    label_define(ctx->text, out_lbl);
    emit_mov_imm64(ctx->text, 8, 64);
    emit_mov_imm64(ctx->text, 0, 1);
    emit_mov_reg(ctx->text, 1, 2);
    emit_mov_reg(ctx->text, 2, 3);
    emit_svc(ctx->text);
    codebuf_write_uint32(ctx->text, 0x910083FF);
    codebuf_write_uint32(ctx->text, 0xA8C17BFD);
    emit_ret(ctx->text);
}

static void emit_strcmp_subroutine(Ctx *ctx) {
    label_define(ctx->text, ctx->lbl_strcmp);
    emit_mov_imm64(ctx->text, 2, 0);
    int loop_lbl = label_new();
    int ne_lbl = label_new();
    int eq_lbl = label_new();
    label_define(ctx->text, loop_lbl);
    emit_ldrb_reg(ctx->text, 3, 0, 2);
    emit_ldrb_reg(ctx->text, 4, 1, 2);
    emit_cmp_reg(ctx->text, 3, 4);
    emit_b_cond(ctx->text, 1, ne_lbl);
    emit_cbz(ctx->text, 3, eq_lbl);
    emit_add(ctx->text, 2, 2, 31);
    emit_mov_imm64(ctx->text, 15, 1);
    emit_add(ctx->text, 2, 2, 15);
    emit_b(ctx->text, loop_lbl);
    label_define(ctx->text, ne_lbl);
    emit_mov_imm64(ctx->text, 0, 1);
    emit_ret(ctx->text);
    label_define(ctx->text, eq_lbl);
    emit_mov_imm64(ctx->text, 0, 0);
    emit_ret(ctx->text);
}

static void emit_strlen_subroutine(Ctx *ctx) {
    label_define(ctx->text, ctx->lbl_strlen);
    emit_mov_imm64(ctx->text, 1, 0);
    int loop_lbl = label_new();
    int end_lbl = label_new();
    label_define(ctx->text, loop_lbl);
    emit_ldrb_reg(ctx->text, 2, 0, 1);
    emit_cbz(ctx->text, 2, end_lbl);
    emit_mov_imm64(ctx->text, 2, 1);
    emit_add(ctx->text, 1, 1, 2);
    emit_b(ctx->text, loop_lbl);
    label_define(ctx->text, end_lbl);
    emit_mov_reg(ctx->text, 0, 1);
    emit_ret(ctx->text);
}

static void emit_filetest_subroutine(Ctx *ctx) {
    label_define(ctx->text, ctx->lbl_filetest);
    codebuf_write_uint32(ctx->text, 0xA9BF7BFD);
    codebuf_write_uint32(ctx->text, 0xD10403FF);
    emit_mov_reg(ctx->text, 10, 0);
    emit_mov_reg(ctx->text, 11, 1);
    emit_mov_imm64(ctx->text, 0, (unsigned long long)-100);
    emit_mov_reg(ctx->text, 1, 10);
    codebuf_write_uint32(ctx->text, 0x910003E2);
    emit_mov_imm64(ctx->text, 3, 0);
    emit_mov_imm64(ctx->text, 8, 79);
    emit_svc(ctx->text);
    emit_cmp_imm(ctx->text, 0, 0);
    int fail_lbl = label_new();
    emit_b_cond(ctx->text, 1, fail_lbl);
    emit_cmp_imm(ctx->text, 11, 'e');
    int is_e_lbl = label_new();
    emit_b_cond(ctx->text, 0, is_e_lbl);
    codebuf_write_uint32(ctx->text, 0xB94013E0);
    emit_mov_imm64(ctx->text, 15, 0xF000);
    codebuf_write_uint32(ctx->text, 0x8A0F0000);
    emit_cmp_imm(ctx->text, 11, 'f');
    int is_f_lbl = label_new();
    emit_b_cond(ctx->text, 0, is_f_lbl);
    emit_cmp_imm(ctx->text, 0, 0x4000);
    int cset_lbl = label_new();
    emit_b(ctx->text, cset_lbl);
    label_define(ctx->text, is_f_lbl);
    emit_cmp_imm(ctx->text, 0, 0x8000);
    label_define(ctx->text, cset_lbl);
    emit_cset(ctx->text, 0, 0);
    int done_lbl = label_new();
    emit_b(ctx->text, done_lbl);
    label_define(ctx->text, is_e_lbl);
    emit_mov_imm64(ctx->text, 0, 1);
    emit_b(ctx->text, done_lbl);
    label_define(ctx->text, fail_lbl);
    emit_mov_imm64(ctx->text, 0, 0);
    label_define(ctx->text, done_lbl);
    codebuf_write_uint32(ctx->text, 0x910403FF);
    codebuf_write_uint32(ctx->text, 0xA8C17BFD);
    emit_ret(ctx->text);
}

static void emit_exec_subroutine(Ctx *ctx) {
    label_define(ctx->text, ctx->lbl_exec);
    codebuf_write_uint32(ctx->text, 0xA9BF7BFD);
    codebuf_write_uint32(ctx->text, 0xD10403FF);
    emit_str_imm(ctx->text, 0, 31, 248);
    emit_str_imm(ctx->text, 1, 31, 240);
    const char *dirs[] = {"/home/user/.super_doubao/super-doubao-runtime/workspace/ShellToAssembly_v1.1/fakebin/", "/bin/", "/usr/bin/", "/usr/local/bin/"};
    for (int d = 0; d < 3; d++) {
        codebuf_write_uint32(ctx->text, 0x910203F0);
        int name_off = rodata_add_string(ctx->rodata, dirs[d], 0);
        int adr_imm = ctx->text_total_size + name_off - ctx->text->size;
        emit_adr(ctx->text, 1, adr_imm);
        int copy_lbl = label_new();
        int copy_end = label_new();
        label_define(ctx->text, copy_lbl);
        emit_ldrb_imm(ctx->text, 2, 1, 0);
        emit_cbz(ctx->text, 2, copy_end);
        emit_strb_imm(ctx->text, 2, 16, 0);
        emit_add(ctx->text, 1, 1, 31);
        emit_mov_imm64(ctx->text, 2, 1);
        emit_add(ctx->text, 1, 1, 2);
        emit_add(ctx->text, 16, 16, 31);
        emit_add(ctx->text, 16, 16, 2);
        emit_b(ctx->text, copy_lbl);
        label_define(ctx->text, copy_end);
        emit_ldr_imm(ctx->text, 17, 31, 248);
        int name_copy_lbl = label_new();
        int name_end = label_new();
        label_define(ctx->text, name_copy_lbl);
        emit_ldrb_imm(ctx->text, 2, 17, 0);
        emit_cbz(ctx->text, 2, name_end);
        emit_strb_imm(ctx->text, 2, 16, 0);
        emit_add(ctx->text, 17, 17, 31);
        emit_mov_imm64(ctx->text, 2, 1);
        emit_add(ctx->text, 17, 17, 2);
        emit_add(ctx->text, 16, 16, 31);
        emit_add(ctx->text, 16, 16, 2);
        emit_b(ctx->text, name_copy_lbl);
        label_define(ctx->text, name_end);
        emit_strb_imm(ctx->text, 31, 16, 0);
        codebuf_write_uint32(ctx->text, 0x910203E0);
        emit_ldr_imm(ctx->text, 1, 31, 240);
        emit_mov_imm64(ctx->text, 2, 0);
        emit_mov_imm64(ctx->text, 8, 221);
        emit_svc(ctx->text);
    }
    emit_mov_imm64(ctx->text, 0, 127);
    emit_mov_imm64(ctx->text, 8, 93);
    emit_svc(ctx->text);
}

/* ==================== Main entry ==================== */

static void codegen_pass(Ctx *ctx, Stmt *ast) {
    ctx->lbl_itoa = label_new();
    ctx->lbl_strcmp = label_new();
    ctx->lbl_strlen = label_new();
    ctx->lbl_exec = label_new();
    ctx->lbl_filetest = label_new();
    int main_lbl = label_new();
    emit_movz(ctx->text, 19, 0x0041, 16);
    emit_bl(ctx->text, main_lbl);
    emit_mov_imm64(ctx->text, 8, 93);
    emit_mov_imm64(ctx->text, 0, 0);
    emit_svc(ctx->text);
    label_define(ctx->text, main_lbl);
    codebuf_write_uint32(ctx->text, 0xA9BF7BFD);
    codebuf_write_uint32(ctx->text, 0xD10003FF);
    codegen_stmt_list(ctx, ast);
    codebuf_write_uint32(ctx->text, 0x910003FF);
    codebuf_write_uint32(ctx->text, 0xA8C17BFD);
    emit_ret(ctx->text);
    emit_itoa_subroutine(ctx);
    emit_strcmp_subroutine(ctx);
    emit_strlen_subroutine(ctx);
    emit_filetest_subroutine(ctx);
    emit_exec_subroutine(ctx);
}

int codegen_arm64(Stmt *ast, CodeBuffer *text, CodeBuffer *rodata,
                   CodeBuffer *data, VarInfo *vinfo, int opt_level) {
    CodeBuffer *tmp_text = codebuf_new();
    CodeBuffer *tmp_rodata = codebuf_new();
    Ctx tmp_ctx;
    tmp_ctx.text = tmp_text;
    tmp_ctx.rodata = tmp_rodata;
    tmp_ctx.data = data;
    tmp_ctx.vinfo = vinfo;
    tmp_ctx.opt_level = opt_level;
    tmp_ctx.text_total_size = 0;
    g_label_count = 0;
    codegen_pass(&tmp_ctx, ast);
    int total_size = tmp_text->size;
    codebuf_free(tmp_text);
    codebuf_free(tmp_rodata);

    g_label_count = 0;
    vinfo->var_count = 0;
    Ctx ctx;
    ctx.text = text;
    ctx.rodata = rodata;
    ctx.data = data;
    ctx.vinfo = vinfo;
    ctx.opt_level = opt_level;
    ctx.text_total_size = total_size;
    codegen_pass(&ctx, ast);

    for (int i = 0; i < 2048; i++) codebuf_write_byte(data, 0);
    return 0;
}
