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

#include "stb_ds.h"

typedef struct {
	char *key;
	int value;
} var_map;

static var_map *variables = NULL;
static int stack_offset = 0;
static int label_counter = 0;
static int *break_stack = NULL;

void gen_expr(FILE *fp, ast_node *expr);
void gen_unary(FILE *fp, ast_node *expr);
void gen_statement_list(FILE *fp, ast_node *stmt);
int get_var_offset(char *name, usize name_len);
int get_var_offset_sized(char *name, usize name_len, usize size);

void gen_binary(FILE *fp, ast_node *expr)
{
	switch (expr->expr.binary.operator) {
		case OP_PLUS:
			gen_expr(fp, expr->expr.binary.left);
			fprintf(fp, "mov %%rax, %%rcx\n");
			gen_expr(fp, expr->expr.binary.right);
			fprintf(fp, "add %%rcx, %%rax\n");
			break;
		case OP_MINUS:
			gen_expr(fp, expr->expr.binary.left);
			fprintf(fp, "mov %%rax, %%rcx\n");
			gen_expr(fp, expr->expr.binary.right);
			fprintf(fp, "sub %%rax, %%rcx\n");
			fprintf(fp, "mov %%rcx, %%rax\n");
			break;
		case OP_MUL:
			gen_expr(fp, expr->expr.binary.left);
			fprintf(fp, "mov %%rax, %%rcx\n");
			gen_expr(fp, expr->expr.binary.right);
			fprintf(fp, "imul %%rcx, %%rax\n");
			break;
		case OP_DIV:
			gen_expr(fp, expr->expr.binary.left);
			fprintf(fp, "mov %%rax, %%rcx\n");
			gen_expr(fp, expr->expr.binary.right);
			fprintf(fp, "mov %%rax, %%rbx\n");
			fprintf(fp, "mov %%rcx, %%rax\n");
			fprintf(fp, "cqo\n");
			fprintf(fp, "idiv %%rbx\n");
			break;
		case OP_MOD:
			gen_expr(fp, expr->expr.binary.left);
			fprintf(fp, "mov %%rax, %%rcx\n");
			gen_expr(fp, expr->expr.binary.right);
			fprintf(fp, "mov %%rax, %%rbx\n");
			fprintf(fp, "mov %%rcx, %%rax\n");
			fprintf(fp, "cqo\n");
			fprintf(fp, "idiv %%rbx\n");
			fprintf(fp, "mov %%rdx, %%rax\n");
			break;
		case OP_BOR:
			gen_expr(fp, expr->expr.binary.left);
			fprintf(fp, "mov %%rax, %%rcx\n");
			gen_expr(fp, expr->expr.binary.right);
			fprintf(fp, "or %%rcx, %%rax\n");
			break;
		case OP_BAND:
			gen_expr(fp, expr->expr.binary.left);
			fprintf(fp, "mov %%rax, %%rcx\n");
			gen_expr(fp, expr->expr.binary.right);
			fprintf(fp, "and %%rcx, %%rax\n");
			break;
		case OP_BXOR:
			gen_expr(fp, expr->expr.binary.left);
			fprintf(fp, "mov %%rax, %%rcx\n");
			gen_expr(fp, expr->expr.binary.right);
			fprintf(fp, "xor %%rcx, %%rax\n");
			break;
		case OP_EQ:
			gen_expr(fp, expr->expr.binary.left);
			fprintf(fp, "mov %%rax, %%rcx\n");
			gen_expr(fp, expr->expr.binary.right);
			fprintf(fp, "cmp %%rax, %%rcx\n");
			fprintf(fp, "sete %%al\n");
			fprintf(fp, "movzx %%al, %%rax\n");
			break;
		case OP_NEQ:
			gen_expr(fp, expr->expr.binary.left);
			fprintf(fp, "mov %%rax, %%rcx\n");
			gen_expr(fp, expr->expr.binary.right);
			fprintf(fp, "cmp %%rax, %%rcx\n");
			fprintf(fp, "setne %%al\n");
			fprintf(fp, "movzx %%al, %%rax\n");
			break;
		case OP_LT:
			gen_expr(fp, expr->expr.binary.left);
			fprintf(fp, "mov %%rax, %%rcx\n");
			gen_expr(fp, expr->expr.binary.right);
			fprintf(fp, "cmp %%rax, %%rcx\n");
			fprintf(fp, "setl %%al\n");
			fprintf(fp, "movzx %%al, %%rax\n");
			break;
		case OP_GT:
			gen_expr(fp, expr->expr.binary.left);
			fprintf(fp, "mov %%rax, %%rcx\n");
			gen_expr(fp, expr->expr.binary.right);
			fprintf(fp, "cmp %%rax, %%rcx\n");
			fprintf(fp, "setg %%al\n");
			fprintf(fp, "movzx %%al, %%rax\n");
			break;
		case OP_LE:
			gen_expr(fp, expr->expr.binary.left);
			fprintf(fp, "mov %%rax, %%rcx\n");
			gen_expr(fp, expr->expr.binary.right);
			fprintf(fp, "cmp %%rax, %%rcx\n");
			fprintf(fp, "setle %%al\n");
			fprintf(fp, "movzx %%al, %%rax\n");
			break;
		case OP_GE:
			gen_expr(fp, expr->expr.binary.left);
			fprintf(fp, "mov %%rax, %%rcx\n");
			gen_expr(fp, expr->expr.binary.right);
			fprintf(fp, "cmp %%rax, %%rcx\n");
			fprintf(fp, "setge %%al\n");
			fprintf(fp, "movzx %%al, %%rax\n");
			break;
		case OP_AND:
			gen_expr(fp, expr->expr.binary.left);
			fprintf(fp, "test %%rax, %%rax\n");
			fprintf(fp, "setne %%al\n");
			fprintf(fp, "movzx %%al, %%rcx\n");
			gen_expr(fp, expr->expr.binary.right);
			fprintf(fp, "test %%rax, %%rax\n");
			fprintf(fp, "setne %%al\n");
			fprintf(fp, "movzx %%al, %%rax\n");
			fprintf(fp, "and %%rcx, %%rax\n");
			break;
		case OP_OR:
			gen_expr(fp, expr->expr.binary.left);
			fprintf(fp, "test %%rax, %%rax\n");
			fprintf(fp, "setne %%al\n");
			fprintf(fp, "movzx %%al, %%rcx\n");
			gen_expr(fp, expr->expr.binary.right);
			fprintf(fp, "test %%rax, %%rax\n");
			fprintf(fp, "setne %%al\n");
			fprintf(fp, "movzx %%al, %%rax\n");
			fprintf(fp, "or %%rcx, %%rax\n");
			break;
		case OP_ASSIGN: {
			if (expr->expr.binary.left->type != NODE_IDENTIFIER) {
				fprintf(fp, "# ERROR: left side of assignment must be identifier\n");
				break;
			}
			gen_expr(fp, expr->expr.binary.right);
			int offset = get_var_offset(expr->expr.binary.left->expr.string.start,
			                            expr->expr.binary.left->expr.string.len);
			fprintf(fp, "mov %%rax, -%d(%%rbp)\n", offset);
			break;
		}
		case OP_ASSIGN_PTR: {
			gen_expr(fp, expr->expr.binary.left);
			fprintf(fp, "push %%rax\n");
			gen_expr(fp, expr->expr.binary.right);
			fprintf(fp, "pop %%rcx\n");
			fprintf(fp, "mov %%rax, (%%rcx)\n");
			break;
		}
		case OP_PLUS_EQ: {
			if (expr->expr.binary.left->type != NODE_IDENTIFIER) {
				fprintf(fp, "# ERROR: left side of assignment must be identifier\n");
				break;
			}
			int offset = get_var_offset(expr->expr.binary.left->expr.string.start,
			                            expr->expr.binary.left->expr.string.len);
			fprintf(fp, "mov -%d(%%rbp), %%rcx\n", offset);
			gen_expr(fp, expr->expr.binary.right);
			fprintf(fp, "add %%rcx, %%rax\n");
			fprintf(fp, "mov %%rax, -%d(%%rbp)\n", offset);
			break;
		}
		case OP_MINUS_EQ: {
			if (expr->expr.binary.left->type != NODE_IDENTIFIER) {
				fprintf(fp, "# ERROR: left side of assignment must be identifier\n");
				break;
			}
			int offset = get_var_offset(expr->expr.binary.left->expr.string.start,
			                            expr->expr.binary.left->expr.string.len);
			fprintf(fp, "mov -%d(%%rbp), %%rcx\n", offset);
			gen_expr(fp, expr->expr.binary.right);
			fprintf(fp, "sub %%rax, %%rcx\n");
			fprintf(fp, "mov %%rcx, %%rax\n");
			fprintf(fp, "mov %%rax, -%d(%%rbp)\n", offset);
			break;
		}
		case OP_MUL_EQ: {
			if (expr->expr.binary.left->type != NODE_IDENTIFIER) {
				fprintf(fp, "# ERROR: left side of assignment must be identifier\n");
				break;
			}
			int offset = get_var_offset(expr->expr.binary.left->expr.string.start,
			                            expr->expr.binary.left->expr.string.len);
			fprintf(fp, "mov -%d(%%rbp), %%rcx\n", offset);
			gen_expr(fp, expr->expr.binary.right);
			fprintf(fp, "imul %%rcx, %%rax\n");
			fprintf(fp, "mov %%rax, -%d(%%rbp)\n", offset);
			break;
		}
		case OP_DIV_EQ: {
			if (expr->expr.binary.left->type != NODE_IDENTIFIER) {
				fprintf(fp, "# ERROR: left side of assignment must be identifier\n");
				break;
			}
			int offset = get_var_offset(expr->expr.binary.left->expr.string.start,
			                            expr->expr.binary.left->expr.string.len);
			fprintf(fp, "mov -%d(%%rbp), %%rcx\n", offset);
			gen_expr(fp, expr->expr.binary.right);
			fprintf(fp, "mov %%rax, %%rbx\n");
			fprintf(fp, "mov %%rcx, %%rax\n");
			fprintf(fp, "cqo\n");
			fprintf(fp, "idiv %%rbx\n");
			fprintf(fp, "mov %%rax, -%d(%%rbp)\n", offset);
			break;
		}
		case OP_MOD_EQ: {
			if (expr->expr.binary.left->type != NODE_IDENTIFIER) {
				fprintf(fp, "# ERROR: left side of assignment must be identifier\n");
				break;
			}
			int offset = get_var_offset(expr->expr.binary.left->expr.string.start,
			                            expr->expr.binary.left->expr.string.len);
			fprintf(fp, "mov -%d(%%rbp), %%rcx\n", offset);
			gen_expr(fp, expr->expr.binary.right);
			fprintf(fp, "mov %%rax, %%rbx\n");
			fprintf(fp, "mov %%rcx, %%rax\n");
			fprintf(fp, "cqo\n");
			fprintf(fp, "idiv %%rbx\n");
			fprintf(fp, "mov %%rdx, -%d(%%rbp)\n", offset);
			break;
		}
		case OP_BOR_EQ: {
			if (expr->expr.binary.left->type != NODE_IDENTIFIER) {
				fprintf(fp, "# ERROR: left side of assignment must be identifier\n");
				break;
			}
			int offset = get_var_offset(expr->expr.binary.left->expr.string.start,
			                            expr->expr.binary.left->expr.string.len);
			fprintf(fp, "mov -%d(%%rbp), %%rcx\n", offset);
			gen_expr(fp, expr->expr.binary.right);
			fprintf(fp, "or %%rcx, %%rax\n");
			fprintf(fp, "mov %%rax, -%d(%%rbp)\n", offset);
			break;
		}
		case OP_BAND_EQ: {
			if (expr->expr.binary.left->type != NODE_IDENTIFIER) {
				fprintf(fp, "# ERROR: left side of assignment must be identifier\n");
				break;
			}
			int offset = get_var_offset(expr->expr.binary.left->expr.string.start,
			                            expr->expr.binary.left->expr.string.len);
			fprintf(fp, "mov -%d(%%rbp), %%rcx\n", offset);
			gen_expr(fp, expr->expr.binary.right);
			fprintf(fp, "and %%rcx, %%rax\n");
			fprintf(fp, "mov %%rax, -%d(%%rbp)\n", offset);
			break;
		}
		case OP_BXOR_EQ: {
			if (expr->expr.binary.left->type != NODE_IDENTIFIER) {
				fprintf(fp, "# ERROR: left side of assignment must be identifier\n");
				break;
			}
			int offset = get_var_offset(expr->expr.binary.left->expr.string.start,
			                            expr->expr.binary.left->expr.string.len);
			fprintf(fp, "mov -%d(%%rbp), %%rcx\n", offset);
			gen_expr(fp, expr->expr.binary.right);
			fprintf(fp, "xor %%rcx, %%rax\n");
			fprintf(fp, "mov %%rax, -%d(%%rbp)\n", offset);
			break;
		}
		case OP_RSHIFT_EQ: {
			if (expr->expr.binary.left->type != NODE_IDENTIFIER) {
				fprintf(fp, "# ERROR: left side of assignment must be identifier\n");
				break;
			}
			int offset = get_var_offset(expr->expr.binary.left->expr.string.start,
			                            expr->expr.binary.left->expr.string.len);
			fprintf(fp, "mov -%d(%%rbp), %%rax\n", offset);
			fprintf(fp, "push %%rax\n");
			gen_expr(fp, expr->expr.binary.right);
			fprintf(fp, "mov %%rax, %%rcx\n");
			fprintf(fp, "pop %%rax\n");
			fprintf(fp, "sar %%cl, %%rax\n");
			fprintf(fp, "mov %%rax, -%d(%%rbp)\n", offset);
			break;
		}
		case OP_LSHIFT_EQ: {
			if (expr->expr.binary.left->type != NODE_IDENTIFIER) {
				break;
			}
			int offset = get_var_offset(expr->expr.binary.left->expr.string.start,
			                            expr->expr.binary.left->expr.string.len);
			fprintf(fp, "mov -%d(%%rbp), %%rax\n", offset);
			fprintf(fp, "push %%rax\n");
			gen_expr(fp, expr->expr.binary.right);
			fprintf(fp, "mov %%rax, %%rcx\n");
			fprintf(fp, "pop %%rax\n");
			fprintf(fp, "shl %%cl, %%rax\n");
			fprintf(fp, "mov %%rax, -%d(%%rbp)\n", offset);
			break;
		}
	}
}

int get_var_offset(char *name, usize name_len)
{
	char *var_name = strndup(name, name_len);
	ptrdiff_t idx = shgeti(variables, var_name);

	if (idx >= 0) {
		free(var_name);
		return variables[idx].value;
	}

	stack_offset += 8;
	shput(variables, var_name, stack_offset);
	return stack_offset;
}

int get_var_offset_sized(char *name, usize name_len, usize size)
{
	char *var_name = strndup(name, name_len);
	ptrdiff_t idx = shgeti(variables, var_name);

	if (idx >= 0) {
		free(var_name);
		return variables[idx].value;
	}

	usize aligned_size = (size + 7) & ~7;
	stack_offset += aligned_size;
	shput(variables, var_name, stack_offset);
	return stack_offset;
}

void gen_statement_list(FILE *fp, ast_node *stmt)
{
	if (!stmt) return;

	if (stmt->type == NODE_UNIT) {
		ast_node *current = stmt;
		while (current && current->type == NODE_UNIT) {
			if (current->expr.unit_node.expr) {
				gen_expr(fp, current->expr.unit_node.expr);
			}
			current = current->expr.unit_node.next;
		}
	} else {
		gen_expr(fp, stmt);
	}
}

void gen_unary(FILE *fp, ast_node *expr)
{
	switch (expr->expr.unary.operator) {
		case UOP_MINUS:
			gen_expr(fp, expr->expr.unary.right);
			fprintf(fp, "neg %%rax\n");
			break;
		case UOP_NOT:
			gen_expr(fp, expr->expr.unary.right);
			fprintf(fp, "test %%rax, %%rax\n");
			fprintf(fp, "sete %%al\n");
			fprintf(fp, "movzx %%al, %%rax\n");
			break;
		case UOP_INCR:
			if (expr->expr.unary.right->type != NODE_IDENTIFIER) {
				fprintf(fp, "# ERROR: increment requires identifier\n");
				break;
			}
			int offset_incr = get_var_offset(expr->expr.unary.right->expr.string.start,
			                                  expr->expr.unary.right->expr.string.len);
			fprintf(fp, "mov -%d(%%rbp), %%rax\n", offset_incr);
			fprintf(fp, "inc %%rax\n");
			fprintf(fp, "mov %%rax, -%d(%%rbp)\n", offset_incr);
			break;
		case UOP_DECR:
			if (expr->expr.unary.right->type != NODE_IDENTIFIER) {
				fprintf(fp, "# ERROR: decrement requires identifier\n");
				break;
			}
			int offset_decr = get_var_offset(expr->expr.unary.right->expr.string.start,
			                                  expr->expr.unary.right->expr.string.len);
			fprintf(fp, "mov -%d(%%rbp), %%rax\n", offset_decr);
			fprintf(fp, "dec %%rax\n");
			fprintf(fp, "mov %%rax, -%d(%%rbp)\n", offset_decr);
			break;
		case UOP_REF:
			if (expr->expr.unary.right->type != NODE_IDENTIFIER) {
				fprintf(fp, "# ERROR: address-of requires identifier\n");
				break;
			}
			int offset_ref = get_var_offset(expr->expr.unary.right->expr.string.start,
			                                 expr->expr.unary.right->expr.string.len);
			fprintf(fp, "lea -%d(%%rbp), %%rax\n", offset_ref);
			break;
		case UOP_DEREF:
			gen_expr(fp, expr->expr.unary.right);
			fprintf(fp, "mov (%%rax), %%rax\n");
			break;
	}
}

void gen_expr(FILE *fp, ast_node *expr)
{
	switch (expr->type) {
		case NODE_INTEGER:
			fprintf(fp, "mov $%lu, %%rax\n", expr->expr.integer);
			break;
		case NODE_FLOAT: {
			// TODO: do not truncate
			i64 int_val = (i64)expr->expr.flt;
			fprintf(fp, "mov $%ld, %%rax\n", int_val);
			break;
		}
		case NODE_CHAR:
			fprintf(fp, "mov $%d, %%rax\n", (int)(unsigned char)expr->expr.ch);
			break;
		case NODE_BOOL:
			fprintf(fp, "mov $%d, %%rax\n", expr->expr.boolean ? 1 : 0);
			break;
		case NODE_IDENTIFIER: {
			int offset = get_var_offset(expr->expr.string.start, expr->expr.string.len);
			fprintf(fp, "mov -%d(%%rbp), %%rax\n", offset);
			break;
		}
		case NODE_BINARY:
			gen_binary(fp, expr);
			break;
		case NODE_UNARY:
			gen_unary(fp, expr);
			break;
		case NODE_CAST:
			gen_expr(fp, expr->expr.cast.value);
			break;
		case NODE_VAR_DECL: {
			int offset;
			if (expr->expr.var_decl.type && expr->expr.var_decl.type->expr_type) {
				usize var_size = expr->expr.var_decl.type->expr_type->size;
				offset = get_var_offset_sized(expr->expr.var_decl.name,
				                              expr->expr.var_decl.name_len, var_size);
			} else {
				offset = get_var_offset(expr->expr.var_decl.name, expr->expr.var_decl.name_len);
			}

			if (expr->expr.var_decl.value) {
				if (expr->expr.var_decl.value->type == NODE_STRUCT_INIT) {
					ast_node *member_list = expr->expr.var_decl.value->expr.struct_init.members;
					ast_node *current = member_list;

					type *struct_type = expr->expr_type;
					if (!struct_type && expr->expr.var_decl.type) {
						struct_type = expr->expr.var_decl.type->expr_type;
					}

					while (current && current->type == NODE_UNIT) {
						ast_node *assignment = current->expr.unit_node.expr;
						if (assignment && assignment->type == NODE_BINARY &&
						    assignment->expr.binary.operator == OP_ASSIGN) {
							ast_node *field = assignment->expr.binary.left;
							ast_node *value = assignment->expr.binary.right;

							if (field->type == NODE_IDENTIFIER && struct_type &&
							    struct_type->tag == TYPE_STRUCT) {
								char *field_name = strndup(field->expr.string.start,
								                          field->expr.string.len);

								member *m = struct_type->data.structure.members;
								int field_offset = -1;
								while (m) {
									if (m->name_len == field->expr.string.len &&
									    strncmp(m->name, field->expr.string.start, m->name_len) == 0) {
										field_offset = m->offset;
										break;
									}
									m = m->next;
								}

								if (field_offset >= 0) {
									gen_expr(fp, value);

									type *field_type = shget(struct_type->data.structure.member_types, field_name);

									if (field_type && field_type->size == 4) {
										fprintf(fp, "mov %%eax, -%d(%%rbp)\n", offset + field_offset);
									} else if (field_type && field_type->size == 2) {
										fprintf(fp, "mov %%ax, -%d(%%rbp)\n", offset + field_offset);
									} else if (field_type && field_type->size == 1) {
										fprintf(fp, "mov %%al, -%d(%%rbp)\n", offset + field_offset);
									} else {
										fprintf(fp, "mov %%rax, -%d(%%rbp)\n", offset + field_offset);
									}
								}

								free(field_name);
							}
						}
						current = current->expr.unit_node.next;
					}
				} else {
					gen_expr(fp, expr->expr.var_decl.value);
					fprintf(fp, "mov %%rax, -%d(%%rbp)\n", offset);
				}
			}
			break;
		}
		case NODE_RETURN: {
			if (expr->expr.ret.value) {
				gen_expr(fp, expr->expr.ret.value);
			}
			fprintf(fp, "mov %%rbp, %%rsp\n");
			fprintf(fp, "pop %%rbp\n");
			fprintf(fp, "ret\n");
			break;
		}
		case NODE_IF: {
			int label_else = label_counter++;
			int label_end = label_counter++;

			gen_expr(fp, expr->expr.if_stmt.condition);
			fprintf(fp, "test %%rax, %%rax\n");

			if (expr->expr.if_stmt.otherwise) {
				fprintf(fp, "jz .L%d\n", label_else);
			} else {
				fprintf(fp, "jz .L%d\n", label_end);
			}

			gen_statement_list(fp, expr->expr.if_stmt.body);

			if (expr->expr.if_stmt.otherwise) {
				fprintf(fp, "jmp .L%d\n", label_end);
				fprintf(fp, ".L%d:\n", label_else);
				gen_statement_list(fp, expr->expr.if_stmt.otherwise);
			}

			fprintf(fp, ".L%d:\n", label_end);
			break;
		}
		case NODE_WHILE: {
			int label_start = label_counter++;
			int label_end = label_counter++;

			fprintf(fp, ".L%d:\n", label_start);

			gen_expr(fp, expr->expr.whle.condition);
			fprintf(fp, "test %%rax, %%rax\n");
			fprintf(fp, "jz .L%d\n", label_end);

			arrput(break_stack, label_end);

			gen_statement_list(fp, expr->expr.whle.body);

			arrpop(break_stack);

			fprintf(fp, "jmp .L%d\n", label_start);

			fprintf(fp, ".L%d:\n", label_end);
			break;
		}
		case NODE_LABEL: {
			char *label_name = strndup(expr->expr.label.name, expr->expr.label.name_len);
			fprintf(fp, ".L_%s:\n", label_name);
			free(label_name);
			break;
		}
		case NODE_GOTO: {
			char *label_name = strndup(expr->expr.label.name, expr->expr.label.name_len);
			fprintf(fp, "jmp .L_%s\n", label_name);
			free(label_name);
			break;
		}
		case NODE_BREAK: {
			if (arrlen(break_stack) > 0) {
				int loop_end = break_stack[arrlen(break_stack) - 1];
				fprintf(fp, "jmp .L%d\n", loop_end);
			} else {
				fprintf(fp, "# ERROR: break outside of loop\n");
			}
			break;
		}
		case NODE_ACCESS: {
			ast_node *base = expr->expr.access.expr;
			ast_node *member_node = expr->expr.access.member;

			if (base->type == NODE_IDENTIFIER) {
				int base_offset = get_var_offset(base->expr.string.start, base->expr.string.len);

				type *struct_type = base->expr_type;

				if (member_node->type == NODE_IDENTIFIER && struct_type &&
				    struct_type->tag == TYPE_STRUCT) {
					member *m = struct_type->data.structure.members;
					int field_offset = -1;
					while (m) {
						if (m->name_len == member_node->expr.string.len &&
						    strncmp(m->name, member_node->expr.string.start, m->name_len) == 0) {
							field_offset = m->offset;
							break;
						}
						m = m->next;
					}

					if (field_offset >= 0) {
						fprintf(fp, "mov -%d(%%rbp), %%rax\n", base_offset + field_offset);
					} else {
						fprintf(fp, "# ERROR: field not found\n");
					}
				} else {
					fprintf(fp, "# ERROR: not a struct type\n");
				}
			} else {
				fprintf(fp, "# ERROR: complex struct access not implemented\n");
			}
			break;
		}
		case NODE_STRUCT_INIT: {
			fprintf(fp, "# ERROR: struct init outside of variable declaration\n");
			break;
		}
		case NODE_CALL: {
			const char *arg_regs[] = {"%rdi", "%rsi", "%rdx", "%rcx", "%r8", "%r9"};

			int param_count = 0;
			ast_node *param = expr->expr.call.parameters;
			while (param && param->type == NODE_UNIT) {
				param_count++;
				param = param->expr.unit_node.next;
			}

			param = expr->expr.call.parameters;
			int i = 0;
			while (param && param->type == NODE_UNIT) {
				if (param->expr.unit_node.expr) {
					gen_expr(fp, param->expr.unit_node.expr);
					fprintf(fp, "push %%rax\n");
				}
				param = param->expr.unit_node.next;
				i++;
			}

			for (int j = param_count - 1; j >= 0; j--) {
				if (j < 6) {
					fprintf(fp, "pop %s\n", arg_regs[j]);
				} else {
					// TODO: handle more than 6 arguments properly
					fprintf(fp, "pop %%rax\n");
					fprintf(fp, "push %%rax\n");
				}
			}

			fprintf(fp, "call %.*s\n", (int)expr->expr.call.name_len, expr->expr.call.name);

			if (param_count > 6) {
				int stack_args = param_count - 6;
				fprintf(fp, "add $%d, %%rsp\n", stack_args * 8);
			}

			break;
		}
	}
}

void gen_function(FILE *fp, ast_node *fn)
{
	ast_node *current = fn->expr.function.body;

	stack_offset = 0;
	label_counter = 0;
	shfree(variables);
	variables = NULL;
	arrfree(break_stack);
	break_stack = NULL;

	fprintf(fp, ".global %s\n%s:\n", fn->expr.function.name, fn->expr.function.name);

	fprintf(fp, "push %%rbp\n");
	fprintf(fp, "mov %%rsp, %%rbp\n");
	fprintf(fp, "sub $256, %%rsp\n");

	const char *param_regs[] = {"%rdi", "%rsi", "%rdx", "%rcx", "%r8", "%r9"};
	member *param = fn->expr.function.parameters;
	int param_idx = 0;
	while (param && param_idx < 6) {
		int offset = get_var_offset(param->name, param->name_len);
		fprintf(fp, "mov %s, -%d(%%rbp)\n", param_regs[param_idx], offset);
		param = param->next;
		param_idx++;
	}

	while (current && current->type == NODE_UNIT) {
		if (current->expr.unit_node.expr) {
			gen_expr(fp, current->expr.unit_node.expr);
		}
		current = current->expr.unit_node.next;
	}

	fprintf(fp, "mov %%rbp, %%rsp\n");
	fprintf(fp, "pop %%rbp\n");
	fprintf(fp, "ret\n");
}

void generate(ast_node *node)
{
	FILE *fp = fopen("test.s", "w");

	ast_node *current = node;

	fprintf(fp, ".section .text\n");
	while (current && current->type == NODE_UNIT) {
		if (current->expr.unit_node.expr && current->expr.unit_node.expr->type == NODE_FUNCTION) {
			gen_function(fp, current->expr.unit_node.expr);
		}
		current = current->expr.unit_node.next;
	}

	fclose(fp);

	shfree(variables);
	variables = NULL;
	arrfree(break_stack);
	break_stack = NULL;
}