/*
 * Copyright (C) 2025  Johnny Richard <johnny@johnnyrichard.com>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */
#include <assert.h>
#include <stdio.h>
#include <stdlib.h>
#include "array.h"
#include "stack.h"
#include "vm.h"

void
vm_init(vm_t *vm)
{
	assert(vm);
	vm->stack = (stack_t) {
		.capacity = STACK_CAPACITY,
		.size = 0,
		.data = calloc(sizeof(int), STACK_CAPACITY)
	};
	vm->call_stack = (stack_t) {
		.capacity = 3000,
		.size = 0,
		.data = calloc(sizeof(int), 3000)
	};
	vm->heap = calloc(sizeof(int), HEAP_CAPACITY);
	vm->labels = array(label_t);
}

void
vm_run(vm_t *vm, inst_t *insts)
{
	size_t program_size = array_length(insts);

	// resolve lables
	for (size_t ip = 0; ip < program_size; ++ip) {
		if (insts[ip].type == INST_LABEL) {
			int name = insts[ip].operand;
			label_t label = { .name = name, .index = ip };
			array_append(vm->labels, label);
		}
	}

	// resolve jumps
	for (size_t ip = 0; ip < program_size; ++ip) {
		switch (insts[ip].type) {
			case INST_CALL:
			case INST_JMP:
			case INST_JMPN:
			case INST_JMPZ: {
				int label_name = insts[ip].operand;
				size_t i = 0;
				size_t labels_size = array_length(vm->labels);
				for (; i < labels_size; ++i) {
					if (vm->labels[i].name == label_name) {
						insts[ip].operand = vm->labels[i].index;
						break;
					}
				}
				assert(i < labels_size);
			}
			default: break;
		}
	}

	// run
	for (size_t ip = 0; ip < program_size; ++ip) {
		switch (insts[ip].type) {
			case INST_PUSH: {
				stack_push(&vm->stack, insts[ip].operand);
				break;
			}
			case INST_DUP: {
				stack_push(&vm->stack, vm->stack.data[vm->stack.size - 1]);
				break;
			}
			case INST_COPY: {
				int operand = insts[ip].operand;
				stack_push(&vm->stack, vm->stack.data[vm->stack.size - 1 - operand]);
				break;
			}
			case INST_SWAP: {
				int a = stack_pop(&vm->stack);
				int b = stack_pop(&vm->stack);
				stack_push(&vm->stack, a);
				stack_push(&vm->stack, b);
				break;
			}
			case INST_DROP: {
				stack_pop(&vm->stack);
				break;
			}
			case INST_SLIDE: {
				int operand = insts[ip].operand;
				int top = stack_pop(&vm->stack);
				while (operand-- > 0) {
					stack_pop(&vm->stack);
				}
				stack_push(&vm->stack, top);
				break;
			}
			case INST_ADD: {
				int lhs = stack_pop(&vm->stack);
				int rhs = stack_pop(&vm->stack);
				stack_push(&vm->stack, lhs + rhs);
				break;
			}
			case INST_SUB: {
				int rhs = stack_pop(&vm->stack);
				int lhs = stack_pop(&vm->stack);
				stack_push(&vm->stack, lhs - rhs);
				break;
			}
			case INST_MUL: {
				int rhs = stack_pop(&vm->stack);
				int lhs = stack_pop(&vm->stack);
				stack_push(&vm->stack, lhs * rhs);
				break;
			}
			case INST_DIV: {
				int rhs = stack_pop(&vm->stack);
				int lhs = stack_pop(&vm->stack);
				stack_push(&vm->stack, lhs / rhs);
				break;
			}
			case INST_MOD: {
				int rhs = stack_pop(&vm->stack);
				int lhs = stack_pop(&vm->stack);
				stack_push(&vm->stack, lhs % rhs);
				break;
			}
			case INST_STORE: {
				int value = stack_pop(&vm->stack);
				int address = stack_pop(&vm->stack);
				vm->heap[address] = value;
				break;
			}
			case INST_LOAD: {
				int address = stack_pop(&vm->stack);
				stack_push(&vm->stack, vm->heap[address]);
				break;
			}
			case INST_CALL: {
				stack_push(&vm->call_stack, ip + 1);
				ip = insts[ip].operand;
				break;
			}
			case INST_RET: {
				ip = stack_pop(&vm->call_stack);
				break;
			}
			case INST_JMP: {
				ip = insts[ip].operand;
				break;
			}
			case INST_JMPZ: {
				int value = stack_pop(&vm->stack);
				if (value == 0) {
					ip = insts[ip].operand;
				}
				break;
			}
			case INST_JMPN: {
				int value = stack_pop(&vm->stack);
				if (value < 0) {
					ip = insts[ip].operand;
				}
				break;
			}
			case INST_PRINTI: {
				int value = stack_pop(&vm->stack);
				printf("%d", value);
				break;
			}
			case INST_PRINTC: {
				int value = stack_pop(&vm->stack);
				printf("%c", value);
				break;
			}
			case INST_READI: {
				int address = stack_pop(&vm->stack);
				scanf("%d", vm->heap + address);
				break;
			}
			case INST_READC: {
				int address = stack_pop(&vm->stack);
				vm->heap[address] = getchar();
				break;
			}
			case INST_END: {
				return;
			}
			case INST_LABEL: {
				break;
			}
			default: {
				assert(0 && "Unsupported instruction");
				break;
			}
		}
	}
}