/// This file serves conceptually as the intermediate representation (IR)
/// of the compiler. Compared to "asm", this file is aware of stack frames,
/// control flow blocks and labels, compound types like structs and enums,
/// and register allocation.

#include "asm.h"
#include "format.h"
#include "ir.h"
#include "x86encode.h"

#include <assert.h>
#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define MAX_STACK_FRAMES 32
#define MAX_LABELS 256
#define MAX_FIXUPS 256

struct stack_frame {
    uint32_t depth;
    uint32_t label_depth;
};

struct label {
    uint32_t frame;
    uint32_t argc;
    symbol symbol;
};

static uint32_t stack_depth = 0;
static uint32_t stack_frame = 0;
static struct stack_frame stack_frames[MAX_STACK_FRAMES];
static uint32_t label_depth = 0;
static struct label labels[MAX_LABELS];

void init_ir(var* argc, var* argv, var* env) {
    assert(stack_depth == 0 && stack_frame == 0);
    x86_inst_mov_r64_r64(BP, SP);
    x86_inst_add_r64_imm8(BP, 8 * 3);
    *env = stack_depth++;
    *argv = stack_depth++;
    *argc = stack_depth++;
}

void enter(void) {
    assert(stack_frame < MAX_STACK_FRAMES);
    printf("ENTERING: %i, %i\n", stack_depth, label_depth);
    struct stack_frame frame = { .depth = stack_depth, .label_depth = label_depth };
    stack_frames[stack_frame] = frame;
    stack_frame++;
    // exit label
    declare(0);
}

void leave(var* args) {
    assert(stack_frame > 0);
    struct stack_frame frame = stack_frames[stack_frame - 1];
    stack_depth = frame.depth;
    label_depth = frame.label_depth;
    define(frame.label_depth, args);
    stack_frame--;
}

label declare(uint32_t argc) {
    assert(label_depth < MAX_LABELS);
    symbol sym = new_symbol();
    struct label label = { .frame = stack_frame, .argc = argc, .symbol = sym };
    labels[label_depth] = label;
    return label_depth++;
}

label declare_exit(uint32_t argc) {
    label label = stack_frames[stack_frame - 1].label_depth;
    labels[label].argc = argc;
    return label;
}

void define(label l, var* args) {
    struct label* label = &labels[l];
    printf("DEFINING %i (%i)\n", l, label->argc);
    define_executable_symbol(label->symbol);
    assert(label->frame == stack_frame);
    for (uint32_t i = 0; i < label->argc; i++) {
        args[i] = stack_depth + i;
    }
    stack_depth += label->argc;
}

void load_var(reg reg, var var) {
    // the stack grows downward, so the bottom of the stack, BP, points to nothing;
    // subtracting 8 causes it to point to the first variable, 0.
    // (each variable is 8 bytes.)
    x86_inst_mov_r64_m64_disp(reg, BP, -(var * 8) - 8);
}

var push_var(reg reg) {
    x86_inst_push_r64(reg);
    return stack_depth++;
}

void load_args(struct label* label, var* args) {
    struct stack_frame* dest_frame = &stack_frames[label->frame];
    uint32_t depth_diff = stack_depth - dest_frame->depth;
    if (depth_diff > 0) {
        // FIXME: should be immX!!!
        x86_inst_add_r64_imm8(SP, depth_diff);
    }
    for (uint32_t arg = 0; arg < label->argc; arg++) {
        load_var(AX, args[arg]);
        x86_inst_push_r64(AX);
    }
}

void jump(label l, var* args) {
    struct label* label = &labels[l];
    printf("JUMP %i (%i)\n", l, label->argc);
    load_args(label, args);
    inst_jump(label->symbol);
}

void jump_table(size_t branches, label* labels, var index, var* args) {
    assert(0); // UNIMPLEMENTED
}

void jump_if(label l, var cond, var* args) {
    struct label* label = &labels[l];
    printf("JUMP_IF %i (%i)\n", l, label->argc);
    load_var(BX, cond);
    load_args(label, args);
    inst_jump_if_not_zero(label->symbol, BX);
}

void jump_unless(label l, var cond, var* args) {
    struct label* label = &labels[l];
    printf("JUMP_UNLESS %i (%i)\n", l, label->argc);
    load_var(BX, cond);
    load_args(label, args);
    inst_jump_if_zero(label->symbol, BX);
}

var lit(uint64_t lit) {
    x86_inst_mov_r64_imm(AX, lit);
    x86_inst_push_r64(AX);
    return stack_depth++;
}

var lit_string(char* str) {
    fprintf(stderr, "error: string literals not yet implemented\n");
    exit(1);
}

var add(var addend1, var addend2) {
    load_var(AX, addend1);
    load_var(BX, addend2);
    x86_inst_add_r64_r64(AX, BX);
    return push_var(AX);
}

var sub(var subtrahend, var minuend) {
    // TODO: use modr/m
    load_var(AX, subtrahend);
    load_var(BX, minuend);
    x86_inst_sub_r64_r64(AX, BX);
    return push_var(AX);
}

// Linux system call: https://blog.rchapman.org/posts/Linux_System_Call_Table_for_x86_64/
var syscall(size_t argc, var* args) {
    assert(argc > 0 && argc <= 7);
    switch(argc) {
        case 7:
            load_var(R9, args[6]);
            __attribute__((fallthrough));
        case 6:
            load_var(R8, args[5]);
            __attribute__ ((fallthrough));
        case 5:
            load_var(R10, args[4]);
            __attribute__ ((fallthrough));
        case 4:
            load_var(DX, args[3]);
            __attribute__ ((fallthrough));
        case 3:
            load_var(SI, args[2]);
            __attribute__ ((fallthrough));
        case 2:
            load_var(DI, args[1]);
            __attribute__ ((fallthrough));
        case 1:
            // the system call number, not an argument
            load_var(AX, args[0]);
    }
    // NOTE: syscall clobbers rcx and r11.
    x86_inst_syscall();
    return push_var(AX);
}