import "compiler/ast.cup"
import "std/file.cup"

let gen_out_file: File*;
let gen_label_counter = 0;

fn emit_asm4(msg1: char*, msg2: char*, msg3: char*, msg4: char*) {
    fwrite(gen_out_file, msg1, strlen(msg1));
    fwrite(gen_out_file, msg2, strlen(msg2));
    fwrite(gen_out_file, msg3, strlen(msg3));
    fwrite(gen_out_file, msg4, strlen(msg4));
}

fn emit_asm3(msg1: char*, msg2: char*, msg3: char*) {
    fwrite(gen_out_file, msg1, strlen(msg1));
    fwrite(gen_out_file, msg2, strlen(msg2));
    fwrite(gen_out_file, msg3, strlen(msg3));
}

fn emit_asm2(msg1: char*, msg2: char*) {
    fwrite(gen_out_file, msg1, strlen(msg1));
    fwrite(gen_out_file, msg2, strlen(msg2));
}

fn emit_asm(msg: char*) {
    fwrite(gen_out_file, msg, strlen(msg));
}

fn emit_num(num: int) {
    fputu(gen_out_file, num);
}

fn generate_syscall(num: int) {
    emit_asm("    mov rax, "); emit_num(num); emit_asm("\n");
    emit_asm("    syscall\n");
}

fn generate_lvalue_into_rax(node: Node*) {
    if (node.typ == AST_LOCAL_VAR) {
        let offset = node.d.variable.offset;
        emit_asm("    mov rax, rbp\n");
        emit_asm("    sub rax, "); emit_num(offset); emit_asm("\n");
    } else {
        die2("Unsupported type in generate_lvalue_into_rax: ", node_type_to_string(node.typ));
    }
}

fn generate_expr_into_rax(node: Node*) {
    if (node.typ == AST_LITERAL) {
        if (node.etyp.typ == TYPE_INT) {
            emit_asm("    mov rax, "); emit_num(node.d.literal.as_int); emit_asm("\n");
        } else {
            die("Unsupported literal type in generate_expr_into_rax");
        }
    } else if (node.typ == AST_CONDITIONAL) {
        let label = ++gen_label_counter;
        generate_expr_into_rax(node.d.conditional.cond);
        emit_asm("    cmp rax, 0\n");
        emit_asm("    je .cond_els"); emit_num(label); emit_asm("\n");
        generate_expr_into_rax(node.d.conditional.then);
        emit_asm("    jmp .cond"); emit_num(label); emit_asm("\n");
        emit_asm(".cond_els"); emit_num(label); emit_asm(":\n");
        generate_expr_into_rax(node.d.conditional.els);
        emit_asm(".cond"); emit_num(label); emit_asm(":\n");

    } else if (node.typ == AST_PLUS) {
        generate_expr_into_rax(node.d.binary.rhs);
        emit_asm("    push rax\n");
        generate_expr_into_rax(node.d.binary.lhs);
        emit_asm("    pop rbx\n");
        emit_asm("    add rax, rbx\n");

    } else if (node.typ == AST_MINUS) {
        generate_expr_into_rax(node.d.binary.rhs);
        emit_asm("    push rax\n");
        generate_expr_into_rax(node.d.binary.lhs);
        emit_asm("    pop rbx\n");
        emit_asm("    sub rax, rbx\n");

    } else if (node.typ == AST_DIV) {
        generate_expr_into_rax(node.d.binary.rhs);
        emit_asm("    push rax\n");
        generate_expr_into_rax(node.d.binary.lhs);
        emit_asm("    pop rbx\n");
        emit_asm("    cqo\n");
        emit_asm("    idiv rbx\n");

    } else if (node.typ == AST_MOD) {
        generate_expr_into_rax(node.d.binary.rhs);
        emit_asm("    push rax\n");
        generate_expr_into_rax(node.d.binary.lhs);
        emit_asm("    pop rbx\n");
        emit_asm("    cqo\n");
        emit_asm("    idiv rbx\n");
        emit_asm("    mov rax, rdx\n");

    } else if (node.typ == AST_MUL) {
        generate_expr_into_rax(node.d.binary.rhs);
        emit_asm("    push rax\n");
        generate_expr_into_rax(node.d.binary.lhs);
        emit_asm("    pop rbx\n");
        emit_asm("    imul rbx\n");

    } else if (node.typ == AST_EQ) {
        generate_expr_into_rax(node.d.binary.rhs);
        emit_asm("    push rax\n");
        generate_expr_into_rax(node.d.binary.lhs);
        emit_asm("    pop rbx\n");
        emit_asm("    cmp rax, rbx\n");
        emit_asm("    sete al\n");
        emit_asm("    movzx rax, al\n");

    } else if (node.typ == AST_NEQ) {
        generate_expr_into_rax(node.d.binary.rhs);
        emit_asm("    push rax\n");
        generate_expr_into_rax(node.d.binary.lhs);
        emit_asm("    pop rbx\n");
        emit_asm("    cmp rax, rbx\n");
        emit_asm("    setne al\n");
        emit_asm("    movzx rax, al\n");

    } else if (node.typ == AST_LT) {
        generate_expr_into_rax(node.d.binary.rhs);
        emit_asm("    push rax\n");
        generate_expr_into_rax(node.d.binary.lhs);
        emit_asm("    pop rbx\n");
        emit_asm("    cmp rax, rbx\n");
        emit_asm("    setl al\n");
        emit_asm("    movzx rax, al\n");

    } else if (node.typ == AST_LEQ) {
        generate_expr_into_rax(node.d.binary.rhs);
        emit_asm("    push rax\n");
        generate_expr_into_rax(node.d.binary.lhs);
        emit_asm("    pop rbx\n");
        emit_asm("    cmp rax, rbx\n");
        emit_asm("    setle al\n");
        emit_asm("    movzx rax, al\n");

    } else if (node.typ == AST_GT) {
        generate_expr_into_rax(node.d.binary.rhs);
        emit_asm("    push rax\n");
        generate_expr_into_rax(node.d.binary.lhs);
        emit_asm("    pop rbx\n");
        emit_asm("    cmp rax, rbx\n");
        emit_asm("    setg al\n");
        emit_asm("    movzx rax, al\n");

    } else if (node.typ == AST_GEQ) {
        generate_expr_into_rax(node.d.binary.rhs);
        emit_asm("    push rax\n");
        generate_expr_into_rax(node.d.binary.lhs);
        emit_asm("    pop rbx\n");
        emit_asm("    cmp rax, rbx\n");
        emit_asm("    setge al\n");
        emit_asm("    movzx rax, al\n");

    } else if (node.typ == AST_BWAND) {
        generate_expr_into_rax(node.d.binary.rhs);
        emit_asm("    push rax\n");
        generate_expr_into_rax(node.d.binary.lhs);
        emit_asm("    pop rbx\n");
        emit_asm("    and rax, rbx\n");

    } else if (node.typ == AST_BWOR) {
        generate_expr_into_rax(node.d.binary.rhs);
        emit_asm("    push rax\n");
        generate_expr_into_rax(node.d.binary.lhs);
        emit_asm("    pop rbx\n");
        emit_asm("    or rax, rbx\n");

    } else if (node.typ == AST_XOR) {
        generate_expr_into_rax(node.d.binary.rhs);
        emit_asm("    push rax\n");
        generate_expr_into_rax(node.d.binary.lhs);
        emit_asm("    pop rbx\n");
        emit_asm("    xor rax, rbx\n");

    } else if (node.typ == AST_BWINV) {
        generate_expr_into_rax(node.d.unary);
        emit_asm("    not rax\n");

    } else if (node.typ == AST_NEG) {
        generate_expr_into_rax(node.d.unary);
        emit_asm("    neg rax\n");

    } else if (node.typ == AST_NOT) {
        generate_expr_into_rax(node.d.unary);
        emit_asm("    cmp rax, 0\n");
        emit_asm("    sete al\n");
        emit_asm("    movzx rax, al\n");

    } else if (node.typ == AST_OR) {
        // With short circuiting!
        let label = ++gen_label_counter;
        generate_expr_into_rax(node.d.binary.lhs,);
        // If left is true, we can short-circuit
        emit_asm("    cmp rax, 0\n");
        emit_asm("    je .or_r"); emit_num(label); emit_asm("\n");
        emit_asm("    mov rax, 1\n");
        emit_asm("    jmp .or_e"); emit_num(label); emit_asm("\n");
        emit_asm(".or_r"); emit_num(label); emit_asm(":\n");
        generate_expr_into_rax(node.d.binary.rhs);
        // Booleanize the result
        emit_asm("    cmp rax, 0\n");
        emit_asm("    setne al\n");
        emit_asm(".or_e"); emit_num(label); emit_asm(":\n");

    } else if (node.typ == AST_AND) {
        let label = ++gen_label_counter;
        generate_expr_into_rax(node.d.binary.lhs);
        // If left is false, we can short-circuit
        emit_asm("    cmp rax, 0\n");
        emit_asm("    jne .and_r"); emit_num(label); emit_asm("\n");
        emit_asm("    mov rax, 0\n");
        emit_asm("    jmp .and_e"); emit_num(label); emit_asm("\n");
        emit_asm(".and_r"); emit_num(label); emit_asm(":\n");
        generate_expr_into_rax(node.d.binary.rhs);
        // Booleanize the result
        emit_asm("    cmp rax, 0\n");
        emit_asm("    setne al\n");
        emit_asm(".and_e"); emit_num(label); emit_asm(":\n");

    } else if (is_lvalue(node.typ)) {
        generate_lvalue_into_rax(node);
        emit_asm("    mov rax, [rax]\n");

    } else if (node.typ == AST_ASSIGN) {
        putsln("...");
        generate_lvalue_into_rax(node.d.assign.lhs);
        emit_asm("    push rax\n");
        generate_expr_into_rax(node.d.assign.rhs);
        emit_asm("    pop rbx\n");
        emit_asm("    mov [rbx], rax\n");

    } else {
        die2("Unsupported node type in generate_expr_into_rax: ", node_type_to_string(node.typ));
    }
}

fn generate_block(node: Node*);

fn generate_statement(node: Node*) {
    if (node.typ == AST_RETURN) {
        generate_expr_into_rax(node.d.unary);
        emit_asm("    mov rsp, rbp\n");
        emit_asm("    pop rbp\n");
        emit_asm("    ret\n");
    } else if (node.typ == AST_VARDECL) {
        if (node.d.var_decl.init) {
            generate_expr_into_rax(node.d.var_decl.init);
            let offset = node.d.var_decl.var.offset;
            emit_asm("    mov [rbp-"); emit_num(offset); emit_asm("], rax\n");
        }
    } else if (node.typ == AST_BLOCK) {
        generate_block(node);
    } else if (node.typ == AST_IF) {
        let label = ++gen_label_counter;
        generate_expr_into_rax(node.d.conditional.cond);
        emit_asm("    cmp rax, 0\n");
        emit_asm("    je .els"); emit_num(label); emit_asm("\n");
        generate_block(node.d.conditional.then);
        emit_asm("    jmp .if"); emit_num(label); emit_asm("\n");
        emit_asm(".els"); emit_num(label); emit_asm(":\n");
        if (node.d.conditional.els)
            generate_block(node.d.conditional.els);
        emit_asm(".if"); emit_num(label); emit_asm(":\n");

    } else {
        // Default to a simple expression statement
        generate_expr_into_rax(node);
    }
}

fn generate_block(node: Node*) {
    let n = node.d.block.children.size;
    for (let i = 0; i < n; ++i) {
        generate_statement(node.d.block.children.data[i]);
    }
}

fn generate_function(node: Node*) {
    emit_asm3("global func_", node.d.func.name, "\n");
    emit_asm3("func_", node.d.func.name, ":\n");
    emit_asm("    push rbp\n");
    emit_asm("    mov rbp, rsp\n");
    emit_asm("    sub rsp, "); emit_num(node.d.func.max_locals_size); emit_asm("\n");

    generate_block(node.d.func.body);

    emit_asm("    mov rsp, rbp\n");
    emit_asm("    pop rbp\n");
    emit_asm("    ret\n");
}

fn generate_program(ast: Node*, file: File*) {
    gen_out_file = file;

    let n = ast.d.block.children.size;
    for (let i = 0; i < n; ++i) {
        let node: Node* = ast.d.block.children.data[i];
        if (node.typ == AST_FUNC) {
            generate_function(node);
        } else {
            die("Unknown node type in generate_program");
        }
    }

    if (OS_IS_MACOS) {
        emit_asm("global _main\n");
        emit_asm("_main:\n");
        // Push argv
        emit_asm("    mov rax, rsi\n");
        emit_asm("    push rax\n");
        // Push argc
        emit_asm("    mov rax, rdi\n");
        emit_asm("    push rax\n");
    } else {
        emit_asm("global _start\n");
        emit_asm("_start:\n");

        emit_asm("    mov rbp, rsp\n");
        // // Push argv
        emit_asm("    mov rax, rbp\n");
        emit_asm("    add rax, 8\n");
        emit_asm("    push rax\n");
        // Push argc
        emit_asm("    mov rax, [rbp]\n");
        emit_asm("    push rax\n");
    }

    emit_asm("    call func_main\n");
    emit_asm("    mov rdi, rax\n");
    generate_syscall(SYS_exit);
}