#include "asm.h"
#include "ir.h"

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

struct fixups {
    struct fixups* next;
    ip disp;
};

struct label_info {
    const struct frame* frame;
    size_t argc;
    ip definition;
    struct fixups* fixups;
};

struct labels {
    struct labels* prev;
    struct label_info info;
};

struct frame {
    struct frame* prev;
    size_t depth;
    struct labels* labels;
};
static struct frame* current_frame;

struct storage {
    enum { STORE_REG, STORE_STACK } type;
    union { reg reg; size_t off; };
};

static size_t top_of_stack;

// Enter a new stack frame.
static void enter(void) {
    struct frame* next = malloc(sizeof(struct frame));
    next->prev = current_frame;
    next->depth = top_of_stack;
    current_frame = next;
}

// Leave the current stack frame.
static void leave(void) {
    struct frame* next = current_frame;
    current_frame = next->prev;
    top_of_stack = next->depth;
    struct labels* labels = next->labels;
    free(next);
    while (labels != NULL) {
        struct labels* next_label = labels->prev;
        free(labels);
        labels = next_label;
    }
}

// Allocate registers or stack space for the arguments.
static void reserve(size_t argc) {
    top_of_stack += argc;
}

static var new_var(void) {
    var var = top_of_stack;
    top_of_stack++;
    inst_push(RA);
    return var;
}

static struct storage storage(var var) {
    struct storage storage;
    if (var == top_of_stack - 1) {
        storage.type = STORE_REG;
        storage.reg = RA;
    } else {
        storage.type = STORE_STACK;
        storage.off = -var * 8 - 16;
    }
    return storage;
}

static void move(var dest, var src) {
    if (dest == src) return;
    struct storage ds = storage(dest);
    struct storage ss = storage(src);
    if (ds.type == STORE_REG && ss.type == STORE_REG) {
        inst_mov(ds.reg, ss.reg);
    } else if (ds.type == STORE_STACK && ss.type == STORE_REG) {
        inst_mov_to_disp(RB, ss.reg, ds.off);
    } else if (ds.type == STORE_REG && ds.type == STORE_STACK) {
        inst_mov_from_disp(ds.reg, RB, ss.off);
    } else {
        // FIXME: DI is scratch register?
        inst_mov_from_disp(DI, RB, ss.off);
        inst_mov_to_disp(RB, DI, ds.off);
    }
}

static void exchange(var x, var y) {
    if (x == y) return;
    assert(0); // UNIMPLEMENTED
}

// Restore the stack and registers to a previous frame,
// in preparation for a jump out of the current frame.
//
// This involves loading spilled variables into registers,
// restoring the stack pointer,
// spilling variables onto the stack to make space for arguments,
// and relocating arguments to the correct registers.
static void restore(const struct frame* frame, size_t argc, var* args) {
    for (size_t i = 0; i < argc; i++) {
        var arg = args[i];
        size_t depth = frame->depth + i;
        if (arg == depth) continue;
        size_t conflict = (size_t) -1;
        for (size_t j = i + 1; j < argc; j++) {
            if (depth == args[j]) {
                conflict = j;
                break;
            }
        }
        if (conflict == (size_t) -1) {
            move(depth, args[i]);
        } else {
            // TODO: an algorithm which produces fewer exchanges
            exchange(args[conflict], args[i]);
            args[conflict] = args[i];
        }
    }
}

label declare_label(size_t argc) {
    struct labels* labels = malloc(sizeof(struct labels));
    labels->prev = current_frame->labels;
    labels->info.frame = current_frame;
    labels->info.argc = argc;
    labels->info.definition = (ip) -1;
    current_frame->labels = labels;
    return &labels->info;
}

void define_label(label label, var* args) {
    assert(label->frame == current_frame);
    label->definition = here;
    struct fixups* fixups = label->fixups;
    while (fixups != NULL) {
        struct fixups* fixup = fixups;
        inst_jump_resolve(fixup->disp, here);
        fixups = fixup->next;
        free(fixup);
    }
    enter();
    reserve(label->argc);
}

void queue_fixup(label label, ip disp) {
    struct fixups* fixup = malloc(sizeof(struct fixups));
    fixup->next = label->fixups;
    fixup->disp = disp;
    label->fixups = fixup;
}

void jump(label label, var* args) {
    restore(label->frame, label->argc, args);
    if (label->definition != (ip) -1) {
        inst_jump(label->definition);
    } else {
        ip disp = inst_jump_unresolved();
        queue_fixup(label, disp);
    }
    leave();
}

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

void jump_if(label label, var cond, var* args) {
    assert(0); // UNIMPLEMENTED
}


static void save(var* vars) {

}

var lit(uint64_t x) {
    var var = new_var();
    struct storage stg = storage(var);
    assert(stg.type == STORE_REG);
    inst_mov_imm(stg.reg, x);
    return var;
}

void syscall(size_t argc, var* args) {
    assert(argc > 0);
    // rax already populated by top of stack
    // FIXME: this won't work forever
    // FIXME: save args in case we don't want to sysexit
    // FIXME: save other registers
    if (argc > 1)
        inst_mov_from_disp(DI, RB, -args[1] * 8 - 16);
    if (argc > 2)
        inst_mov_from_disp(SI, RB, -args[2] * 8 - 16);
    if (argc > 3)
        inst_mov_from_disp(RD, RB, -args[3] * 8 - 16);
    inst_syscall();
}

void init(void) {
    inst_mov(RB, SP);
}