#include <stdlib.h>
#include <string.h>
#include <assert.h>
#include "to_assembly.h"


static bool is_function_label(const char *lbl) {
	return memcmp(lbl, "__", 2) != 0;
}

static bool needs_0_eq_1(const struct irins *ins) {
	return ins->type == INS_ADD || ins->type == INS_SUB || ins->type == INS_MUL ||
				ins->type == INS_DIV || ins->type == INS_MOD || ins->type == INS_NEG ||
				ins->type == INS_NOT;
}

static void fix_function(struct ir *ir, int startidx, struct allocation *ral) {
	int spillspace = 0;

	int of_cap = 32;
	int *offsets = malloc(of_cap * sizeof(int));
	memset(offsets, 0xff, of_cap * sizeof(int));

	if (ir->inss[startidx]->type == INS_LBL && is_function_label(ir->inss[startidx]->name)) {
		startidx++;
	}

	for (int idx = startidx; idx < ir->len; idx++) {
		struct irins *ins = ir->inss[idx];

		if (ins->type == INS_LBL && is_function_label(ins->name)) break;

		if (ins->type == INS_CALLV) {
			ins->type = INS_CALL;
			ir_insert_before(ir, idx + 1, irins_make_01(INS_MOV, ins->r0, ref_reg(REG_A)));
			idx--;
			continue;
		}
		if (ins->type == INS_RETV) {
			ir_insert_before(ir, idx, irins_make_01(INS_MOV, ref_reg(REG_A), ins->r1));
			ins->type = INS_RET;
			idx--;
			continue;
		}

		// Apply the register allocation
		bool haveref[3];
		irins_which_refs(ins, haveref);
		for (int ri = 0; ri < 3; ri++) {
			struct ref *ref = &ins->three_refs[ri];
			if (haveref[ri] && ref->type == REF_REG && ref->reg >= 0) {
				if (ral->allocs[ref->reg].spill) {
					int old_cap = of_cap;
					while (ref->reg >= of_cap) {
						of_cap *= 2;
						offsets = realloc(offsets, of_cap * sizeof(int));
					}
					memset(offsets + old_cap * sizeof(int), 0xff, (of_cap - old_cap) * sizeof(int));
					if (offsets[ref->reg] == -1) offsets[ref->reg] = ++spillspace;
					*ref = ref_mem(REG_BP, -offsets[ref->reg], REFREL_ZERO);
				} else {
					ref->reg = ral->allocs[ref->reg].reg;
				}
			}
		}

		if (ins->type == INS_MUL) {
			if (!ref_equal(ins->r1, ref_reg(REG_A))) {
				ir_insert_before(ir, idx, irins_make_01(INS_MOV, ref_reg(REG_A), ins->r1));
				ins->r1 = ref_reg(REG_A);
				idx++;
			}
			if (!ref_equal(ins->r0, ref_reg(REG_A))) {
				ir_insert_before(ir, idx + 1, irins_make_01(INS_MOV, ins->r0, ref_reg(REG_A)));
				ins->r0 = ref_reg(REG_A);
			}
		} else if (needs_0_eq_1(ins) && !ref_equal(ins->r0, ins->r1)) {
			ir_insert_before(ir, idx, irins_make_01(INS_MOV, ins->r0, ins->r1));
			ins->r1 = ins->r0;
		}
	}

	free(offsets);

	// TODO: Really, nothing should be relative to BP, instead everything should use SP. That
	// would completely eliminate this BP dance anyway.

	// if (spillspace == 0) return;

	ir_insert_before(ir, startidx + 0, irins_make_1(INS_PUSH, ref_reg(REG_BP)));
	ir_insert_before(ir, startidx + 1, irins_make_01(INS_MOV, ref_reg(REG_BP), ref_reg(REG_SP)));
	if (spillspace != 0) {
		ir_insert_before(ir, startidx + 2,
				irins_make_012(INS_SUB, ref_reg(REG_SP), ref_reg(REG_SP), ref_imm(spillspace)));
	}

	for (int idx = startidx + 1; idx < ir->len; idx++) {
		struct irins *ins = ir->inss[idx];

		if (ins->type == INS_LBL && is_function_label(ins->name)) break;

		if (ins->type == INS_RET || ins->type == INS_BRK) {
			ir_insert_before(ir, idx++, irins_make_01(INS_MOV, ref_reg(REG_SP), ref_reg(REG_BP)));
			ir_insert_before(ir, idx++, irins_make_0(INS_POP, ref_reg(REG_BP)));
			continue;
		}
	}
}

void to_assembly(struct ir *ir, struct allocation *ral) {
	if (ir->len > 0 && (ir->inss[0]->type != INS_LBL || !is_function_label(ir->inss[0]->name))) {
		fix_function(ir, 0, ral);
	}

	for (int idx = 0; idx < ir->len; idx++) {
		if (ir->inss[idx]->type == INS_LBL && is_function_label(ir->inss[idx]->name)) {
			fix_function(ir, idx, ral);
		}
	}

	for (int idx = 0; idx < ir->len; idx++) {
		struct irins *ins = ir->inss[idx];

		bool haveref[3];
		irins_which_refs(ins, haveref);
		for (int ri = 0; ri < 3; ri++) {
			struct ref *ref = &ins->three_refs[ri];

			if (haveref[ri] && ref->type == REF_MEM) {
				switch (ref->rel) {
					case REFREL_ZERO: break;
					case REFREL_DATA: ref->offset += 0x0200; break;
					case REFREL_HEAP: ref->offset += 0x5000; fprintf(stderr, "WARNING: Using hard-coded heap start address\n"); break;
					default: assert(false);
				}
				ref->rel = REFREL_ZERO;
			}
		}
	}
}