#include "vm/vm.h"
#include "vm/print.h"
#include "parse/parse.h"
#include "parse/lex.h"
#include "io.h"
#include "bitset.h"
#include "bytecode.h"
#include "loader.h"

#include <stdio.h>
#include <string.h>

#ifdef __unix__
#define USE_READLINE
#include <unistd.h>
#include <readline/readline.h>
#include <readline/history.h>
#endif

static int do_print_bytecode = 0;
static int do_print_tokens = 0;
static int do_step = 0;
static int do_serialize_bytecode = 0;
static int do_repl = 0;
static char *input_filename = "-";

static int parse_text(FILE *inf, struct l2_io_writer *w) {
	// Init lexer with its input reader
	struct l2_io_file_reader r;
	r.r.read = l2_io_file_read;
	r.f = inf;
	struct l2_lexer lexer;
	l2_lexer_init(&lexer, &r.r);
	if (do_print_tokens) {
		lexer.do_log_tokens = 1;
	}

	// Init gen with its output writer
	struct l2_generator gen;
	l2_gen_init(&gen, w);

	struct l2_parse_error err;
	if (l2_parse_program(&lexer, &gen, &err) < 0) {
		fprintf(stderr, "Parse error: %s:%i:%i: %s\n",
				input_filename, err.line, err.ch, err.message);
		l2_parse_error_free(&err);
		l2_gen_free(&gen);
		fclose(inf);
		return -1;
	}

	l2_gen_free(&gen);
	return 0;
}

static void step_through(struct l2_vm *vm) {
	printf("=====\n\nInitial state:\n");
	l2_vm_print_state(vm);

	char buf[16];
	while (!vm->halted) {
		size_t iptr = vm->iptr;
		printf("\n======\n\n(%d) Will run instr: ", vm->iptr);
		l2_vm_print_op(vm->ops, vm->opcount, &iptr);
		if (fgets(buf, sizeof(buf), stdin) == NULL) {
			break;
		}

		l2_vm_step(vm);
		l2_vm_print_state(vm);
	}
}

static void repl() {
	struct l2_io_mem_writer w = {
		.w.write = l2_io_mem_write,
	};

	struct l2_io_mem_reader r = {
		.r.read = l2_io_mem_read,
	};

	struct l2_lexer lexer;

	struct l2_generator gen;
	l2_gen_init(&gen, &w.w);

	struct l2_vm vm;
	l2_vm_init(&vm, NULL, 0);

	while (1) {
		char line[4096];
#ifdef USE_READLINE
		char *rline = readline("> ");
		if (rline == NULL) return;
		if (rline[0] == '\0') continue;
		add_history(rline);
		snprintf(line, sizeof(line), "print (%s)", rline);
		free(rline);
#else
		char rline[4096];
		if (fgets(rline, sizeof(rline), stdin) == NULL) return;
		if (rline[0] == '\n' && rline[1] == '\0') continue;
		snprintf(line, sizeof(line), "print (%s)", rline);
#endif

		r.idx = 0;
		r.len = strlen(line);
		r.mem = line;
		l2_lexer_init(&lexer, &r.r);

		size_t prev_len = w.len;

		struct l2_parse_error err;
		if (l2_parse_program(&lexer, &gen, &err) < 0) {
			fprintf(stderr, "Parse error: %s\n -- %s\n", err.message, line);
			l2_parse_error_free(&err);
			w.len = prev_len;
		} else if (w.len > 0) {
			vm.ops = w.mem;
			vm.opcount = w.len / sizeof(l2_word);

			while (vm.iptr < vm.opcount) {
				l2_vm_step(&vm);
			}

			l2_vm_gc(&vm);
		}
	}
}

static void usage(const char *argv0) {
	printf("Usage: %s [options] [input|-]\n", argv0);
	printf("\n");
	printf("Options:\n");
	printf("  --help:            Print this help text\n");
	printf("  --bytecode:        Print the generated bytecode, don't execute\n");
	printf("  --tokens:          Print the tokens as the program is parsed, don't execute\n");
	printf("  --step:            Step through the program\n");
	printf("  --repl:            Start a repl\n");
	printf("  --output,-o <out>: Write bytecode to file\n");
}

int main(int argc, char **argv) {
	int was_inf_set = 0;
	FILE *inf = stdin;
	FILE *outbc = NULL;

#ifdef __unix__
	do_repl = isatty(0);
#endif

	int dashes = 0;
	for (int i = 1; i < argc; ++i) {
		if (!dashes && strcmp(argv[i], "--help") == 0) {
			usage(argv[0]);
			return 0;
		} else if (!dashes && strcmp(argv[i], "--bytecode") == 0) {
			do_print_bytecode = 1;
		} else if (!dashes && strcmp(argv[i], "--tokens") == 0) {
			do_print_tokens = 1;
		} else if (!dashes && strcmp(argv[i], "--step") == 0) {
			do_step = 1;
		} else if (!dashes && strcmp(argv[i], "--repl") == 0) {
			do_repl = 1;
		} else if (!dashes && (
				strcmp(argv[i], "--output") == 0 || strcmp(argv[i], "-o") == 0)) {
			if (i == argc - 1) {
				fprintf(stderr, "%s expects an argument\n", argv[i]);
				return 1;
			}

			do_serialize_bytecode = 1;
			i += 1;
			if (strcmp(argv[i], "-") == 0) {
				outbc = stdout;
			} else {
				outbc = fopen(argv[i], "w");
				if (outbc == NULL) {
					perror(argv[i]);
					return 1;
				}
			}
		} else if (strcmp(argv[i], "--") == 0) {
			dashes = 1;
		} else if (strcmp(argv[i], "-") == 0) {
			do_repl = 0;
			input_filename = "-";
			inf = stdin;
		} else if (!was_inf_set) {
			do_repl = 0;
			input_filename = argv[i];
			inf = fopen(argv[i], "r");
			was_inf_set = 1;
			if (inf == NULL) {
				perror(argv[i]);
				return 1;
			}
		} else {
			printf("Unexpected argument: %s\n", argv[i]);
			usage(argv[0]);
			return 1;
		}
	}

	if (do_repl) {
		repl();
		printf("\n");
		return 0;
	}

	struct l2_io_mem_writer bytecode_writer = {
		.w.write = l2_io_mem_write,
	};

	int headerbyte = fgetc(inf);
	if (headerbyte == EOF || ungetc(headerbyte, inf) == EOF) {
		fprintf(stderr, "%s: Reading file failed.\n", input_filename);
		return 1;
	}

	// Detect whether input is compiled bytecode or not
	// (compile bytecode starts with (ESC) 'l' '2' 'c')
	unsigned char header[4];
	if (
			headerbyte == 0x1b &&
			fread(header, 1, 4, inf) >= 4 &&
			header[0] == 0x1b && header[1] == 0x6c &&
			header[2] == 0x32 && header[3] == 0x63) {
		if (l2_bc_load(inf, &bytecode_writer.w) < 0) {
			return 1;
		}
	} else if (headerbyte == 0x1b) {
		fprintf(
				stderr, "%s: Corrupt file? Start byte is 0x1b (ESC)"
				"but the header is wrong\n",
				input_filename);
		return 1;
	} else {
		if (parse_text(inf, &bytecode_writer.w) < 0) {
			return 1;
		}
	}
	fclose(inf);

	if (do_print_bytecode) {
		l2_vm_print_bytecode(bytecode_writer.mem, bytecode_writer.len / sizeof(l2_word));
	}

	if (do_serialize_bytecode) {
		l2_bc_serialize(
				outbc, bytecode_writer.mem,
				bytecode_writer.len / sizeof(l2_word));
	}

	if (do_print_bytecode || do_print_tokens || do_serialize_bytecode) {
		free(bytecode_writer.mem);
		return 0;
	}

	struct l2_vm vm;
	l2_vm_init(&vm, bytecode_writer.mem, bytecode_writer.len / sizeof(l2_word));

	if (do_step) {
		step_through(&vm);
	} else {
		l2_vm_run(&vm);
	}

	l2_vm_free(&vm);
	free(bytecode_writer.mem);
}