const coral = @import("coral");

const tokens = @import("./tokens.zig");

const vm = @import("./vm.zig");

pub const Chunk = struct {
	vm_state: *vm.State,
	arity: usize,
	name_string: *vm.ObjectInstance,
	interned_instances: *vm.ObjectInstance,
	bytecode_buffer: Buffer,

	const Buffer = coral.stack.Dense(u8);

	const InstanceStack = coral.stack.Dense(*vm.ObjectInstance);

	pub fn deinit(self: *Chunk) void {
		self.bytecode_buffer.deinit();
	}

	fn emit_byte(self: *Chunk, byte: u8) !void {
		return self.bytecode_buffer.push_one(byte);
	}

	fn emit_float(self: *Chunk, value: Float) !void {
		return self.bytecode_buffer.push_all(coral.io.bytes_of(&value));
	}

	fn emit_integer(self: *Chunk, value: Integer) !void {
		return self.bytecode_buffer.push_all(coral.io.bytes_of(&value));
	}

	fn emit_opcode(self: *Chunk, opcode: Opcode) !void {
		return self.bytecode_buffer.push_one(@enumToInt(opcode));
	}

	pub fn fetch_byte(self: Chunk, cursor: *usize) ?u8 {
		if (cursor.* >= self.bytecode_buffer.values.len) return null;

		defer cursor.* += 1;

		return self.bytecode_buffer.values[cursor.*];
	}

	pub fn fetch_float(self: Chunk, cursor: *usize) ?Float {
		const operand_size = @sizeOf(Float);
		const updated_cursor = cursor.* + operand_size;

		if (updated_cursor > self.bytecode_buffer.values.len) return null;

		var operand_bytes align(@alignOf(Float)) = [_]u8{0} ** operand_size;

		coral.io.copy(&operand_bytes, self.bytecode_buffer.values[cursor.* .. updated_cursor]);

		cursor.* = updated_cursor;

		return @bitCast(Float, operand_bytes);
	}

	pub fn fetch_integer(self: Chunk, cursor: *usize) ?Integer {
		const operand_size = @sizeOf(Integer);
		const updated_cursor = cursor.* + operand_size;

		if (updated_cursor > self.bytecode_buffer.values.len) return null;

		var operand_bytes align(@alignOf(Integer)) = [_]u8{0} ** operand_size;

		coral.io.copy(&operand_bytes, self.bytecode_buffer.values[cursor.* .. updated_cursor]);

		cursor.* = updated_cursor;

		return @bitCast(Integer, operand_bytes);
	}

	pub fn fetch_interned(self: Chunk, cursor: *usize) ?*vm.ObjectInstance {
		const interned_value = self.vm_state.get(self.interned_instances, .{
			.integer = self.fetch_integer(cursor) orelse return null
		}) catch return null;

		coral.debug.assert(interned_value == .instance);

		return interned_value.instance;
	}

	pub fn fetch_opcode(self: Chunk, cursor: *usize) ?Opcode {
		return @intToEnum(Opcode, self.fetch_byte(cursor) orelse return null);
	}

	pub fn init(vm_state: *vm.State, name: []const u8) !Chunk {
		const assumed_average_bytecode_size = 1024;
		var bytecode_buffer = try Buffer.init(vm_state.allocator, assumed_average_bytecode_size);

		errdefer bytecode_buffer.deinit();

		return Chunk{
			.name_string = try vm_state.new_string(name),
			.interned_instances = try vm_state.new_array(),

			.vm_state = vm_state,
			.bytecode_buffer = bytecode_buffer,
			.arity = 0,
		};
	}

	fn intern_chunk(self: *Chunk, chunk: *Chunk) !vm.Integer {
		var constant_slot = @as(vm.Integer, 0);
		const interned_count = try self.vm_state.count(self.interned_instances);

		while (constant_slot < interned_count) : (constant_slot += 1) {
			const interned_value = try self.vm_state.get(self.interned_instances, .{.integer = constant_slot});

			coral.debug.assert(interned_value == .instance);

			switch (self.vm_state.userdata_of(interned_value.instance).*) {
				.chunk => |interned_chunk| if (interned_chunk == chunk) return constant_slot,
				else => continue,
			}
		}

		try self.vm_state.set(self.interned_instances, .{.integer = constant_slot}, .{
			.instance = try self.vm_state.new(.{.chunk = chunk}, .{
				.caller = struct {
					fn call(state: *vm.State, instance: *vm.ObjectInstance, _: *vm.ObjectInstance, arguments: []const vm.Value) vm.RuntimeError!vm.Value {
						const instance_chunk =
							@ptrCast(*Chunk, @alignCast(@alignOf(Chunk), state.userdata_of(instance).native));

						coral.debug.assert(instance_chunk.vm_state == state);

						return vm.execute_chunk(instance_chunk, arguments);
					}
				}.call,
			})
		});

		return constant_slot;
	}

	fn intern_string(self: *Chunk, string: []const u8) !vm.Integer {
		var constant_slot = @as(vm.Integer, 0);
		const interned_count = try self.vm_state.count(self.interned_instances);

		while (constant_slot < interned_count) : (constant_slot += 1) {
			const interned_value = try self.vm_state.get(self.interned_instances, .{.integer = constant_slot});

			coral.debug.assert(interned_value == .instance);

			switch (self.vm_state.userdata_of(interned_value.instance).*) {
				.string => |interned_string| if (coral.io.equals(interned_string, string)) return constant_slot,
				else => continue,
			}
		}

		try self.vm_state.set(self.interned_instances, .{.integer = constant_slot}, .{
			.instance = try self.vm_state.new_string(string)
		});

		return constant_slot;
	}

	pub fn parse(self: *Chunk, script_tokenizer: *tokens.Tokenizer) ParseError!void {
		errdefer self.reset();

		self.reset();

		var parser = Parser{.tokenizer = script_tokenizer};

		while (true) {
			parser.step() catch |step_error| switch (step_error) {
				error.UnexpectedEnd => return,
			};

			if (!(try parser.parse_statement(self))) break;
		}

		while (true) {
			parser.step() catch |step_error| switch (step_error) {
				error.UnexpectedEnd => return,
			};

			try parser.current_token.expect(.newline);
		}
	}

	pub fn reset(self: *Chunk) void {
		self.bytecode_buffer.clear();
	}
};

pub const Float = f32;

pub const Integer = i32;

pub const Opcode = enum(u8) {
	pop,
	push_nil,
	push_true,
	push_false,
	push_zero,
	push_integer,
	push_float,
	push_array,
	push_table,
	push_interned,

	not,
	neg,
	add,
	sub,
	div,
	mul,

	call,
	get_index,
	set_index,
	get_global,
	set_global,
	get_local,
	set_local,
};

pub const ParseError = vm.RuntimeError || Parser.StepError || tokens.Token.ExpectError || error {
	OutOfMemory,
	IntOverflow,
	UndefinedLocal,
};

const Parser = struct {
	tokenizer: *tokens.Tokenizer,
	current_token: tokens.Token = .newline,
	previous_token: tokens.Token = .newline,
	locals: SmallStack(Local, .{.name = "", .depth = 0}) = .{},

	const Local = struct {
		name: []const u8,
		depth: u16,

		const empty = Local{ .name = "", .depth = 0 };
	};

	const Operator = enum {
		not,
		negate,
		add,
		subtract,
		divide,
		multiply,

		const Self = @This();

		fn opcode(self: Self) Opcode {
			return switch (self) {
				.not => .not,
				.negate => .neg,
				.add => .add,
				.subtract => .sub,
				.multiply => .mul,
				.divide => .div,
			};
		}

		fn precedence(self: Self) isize {
			return switch (self) {
				.not => 13,
				.negate => 13,
				.add => 11,
				.subtract => 11,
				.divide => 12,
				.multiply => 12,
			};
		}
	};

	const OperatorStack = SmallStack(Operator, .not);

	const StepError = error {
		UnexpectedEnd,
	};

	fn declare_local(self: *Parser, name: []const u8) !void {
		return self.locals.push(.{
			.name = name,
			.depth = 0,
		});
	}

	const operator_tokens = &.{.symbol_assign, .symbol_plus,
		.symbol_dash, .symbol_asterisk, .symbol_forward_slash, .symbol_paren_left, .symbol_comma};

	fn parse_closure(self: *Parser, parent_chunk: *Chunk) ParseError!void {
		const closure_chunk = parent_chunk.vm_state.allocator.allocate_one(Chunk) orelse return error.OutOfMemory;

		errdefer parent_chunk.vm_state.allocator.deallocate(closure_chunk);

		closure_chunk.* = try Chunk.init(parent_chunk.vm_state, switch (self.previous_token) {
			.local_identifier => |identifier| identifier,
			.symbol_assign, .symbol_paren_left, .symbol_comma => "<closure>",
			else => return error.UnexpectedToken,
		});

		errdefer closure_chunk.deinit();

		try self.step();
		try self.current_token.expect(.symbol_paren_left);

		while (true) {
			try self.step();

			switch (self.current_token) {
				.symbol_paren_right => break,

				.local_identifier => {
					try self.declare_local(self.current_token.local_identifier);
					try self.step();
				},

				else => return error.UnexpectedToken,
			}

			closure_chunk.arity += 1;

			switch (self.current_token) {
				.symbol_paren_right => break,
				.symbol_comma => continue,
				else => return error.UnexpectedToken,
			}
		}

		try self.step();
		try self.current_token.expect(.symbol_brace_left);

		// TODO: Create new callframe.

		while (true) {
			try self.step();

			switch (self.current_token) {
				.symbol_brace_right => break,
				else => if (try self.parse_statement(closure_chunk)) continue,
			}

			while (true) {
				self.step() catch |step_error| switch (step_error) {
					error.UnexpectedEnd => return,
				};

				switch (self.current_token) {
					.newline => continue,
					.symbol_brace_right => break,
					else => return error.UnexpectedToken,
				}
			}
		}

		try parent_chunk.emit_opcode(.push_interned);
		try parent_chunk.emit_integer(try parent_chunk.intern_chunk(closure_chunk));
		try parent_chunk.emit_opcode(.push_closure);
	}

	fn parse_expression(self: *Parser, chunk: *Chunk) ParseError!void {
		var operators = OperatorStack{};
		var local_depth = @as(usize, 0);

		while (true) {
			switch (self.current_token) {
				.keyword_nil => {
					try self.previous_token.expect_any(operator_tokens);
					try chunk.emit_opcode(.push_nil);

					self.step() catch |step_error| switch (step_error) {
						error.UnexpectedEnd => return,
					};
				},

				.keyword_true => {
					try self.previous_token.expect_any(operator_tokens);
					try chunk.emit_opcode(.push_true);

					self.step() catch |step_error| switch (step_error) {
						error.UnexpectedEnd => return,
					};
				},

				.keyword_false => {
					try self.previous_token.expect_any(operator_tokens);
					try chunk.emit_opcode(.push_false);

					self.step() catch |step_error| switch (step_error) {
						error.UnexpectedEnd => return,
					};
				},

				.integer_literal => |literal| {
					try self.previous_token.expect_any(operator_tokens);

					const value = coral.utf8.parse_signed(@bitSizeOf(i64), literal)
						catch |parse_error| switch (parse_error) {
							error.BadSyntax => unreachable,
							error.IntOverflow => return error.IntOverflow,
						};

					if (value == 0) {
						try chunk.emit_opcode(.push_zero);
					} else {
						try chunk.emit_opcode(.push_integer);
						try chunk.emit_integer(value);
					}

					try self.step();
				},

				.real_literal => |literal| {
					try self.previous_token.expect_any(operator_tokens);

					try chunk.emit_float(coral.utf8.parse_float(@bitSizeOf(f64), literal) catch |parse_error| {
						switch (parse_error) {
							// Already validated to be a real by the tokenizer so this cannot fail, as real syntax is a
							// subset of float syntax.
							error.BadSyntax => unreachable,
						}
					});

					try self.step();
				},

				.string_literal => |literal| {
					try self.previous_token.expect_any(operator_tokens);
					try chunk.emit_opcode(.push_interned);
					try chunk.emit_integer(try chunk.intern_string(literal));
					try self.step();
				},

				.global_identifier, .local_identifier => {
					try self.previous_token.expect_any(&.{.symbol_assign, .symbol_plus,
						.symbol_dash, .symbol_asterisk, .symbol_forward_slash, .symbol_period});

					try self.step();
				},

				.symbol_bang => {
					try self.previous_token.expect_any(operator_tokens);
					try operators.push(.not);
					try self.step();

					local_depth = 0;
				},

				.symbol_plus => {
					try self.parse_operator(chunk, &operators, .add);

					local_depth = 0;
				},

				.symbol_dash => {
					try self.parse_operator(chunk, &operators, .subtract);

					local_depth = 0;
				},

				.symbol_asterisk => {
					try self.parse_operator(chunk, &operators, .multiply);

					local_depth = 0;
				},

				.symbol_forward_slash => {
					try self.parse_operator(chunk, &operators, .divide);

					local_depth = 0;
				},

				.symbol_arrow => {
					try self.parse_closure(chunk);

					local_depth = 0;
				},

				.symbol_period => {
					switch (self.previous_token) {
						.global_identifier => |identifier| {
							try chunk.emit_opcode(.get_global);
							try chunk.emit_integer(try chunk.intern_string(identifier));
						},

						.local_identifier => |identifier| {
							if (local_depth == 0) {
								try chunk.emit_byte(self.resolve_local(identifier) orelse {
									return error.UndefinedLocal;
								});
							} else {
								try chunk.emit_opcode(.get_index);
								try chunk.emit_integer(try chunk.intern_string(identifier));
							}
						},

						else => return error.UnexpectedToken,
					}

					try self.step();

					local_depth += 1;
				},

				.symbol_paren_left => {
					switch (self.previous_token) {
						.local_identifier => |identifier| {
							if (local_depth == 0) {
								try chunk.emit_byte(self.resolve_local(identifier) orelse {
									return error.UndefinedLocal;
								});
							} else {
								try chunk.emit_opcode(.get_index);
								try chunk.emit_integer(try chunk.intern_string(identifier));
							}
						},

						.global_identifier => |identifier| {
							try chunk.emit_opcode(.get_global);
							try chunk.emit_integer(try chunk.intern_string(identifier));
						},

						else => {
							try self.parse_expression(chunk);
							try self.previous_token.expect(.symbol_paren_right);
							try self.step();

							local_depth = 0;

							continue;
						},
					}

					local_depth += 1;

					var argument_count = @as(Integer, 0);

					while (true) {
						try self.step();

						try switch (self.current_token) {
							.symbol_paren_right => break,
							else => self.parse_expression(chunk),
						};

						argument_count += 1;

						switch (self.current_token) {
							.symbol_paren_right => break,
							.symbol_comma => continue,
							else => return error.UnexpectedToken,
						}
					}

					try chunk.emit_opcode(.call);
					try chunk.emit_integer(argument_count);
					try self.step();

					local_depth = 0;
				},

				.symbol_brace_left => {
					const is_call_argument = switch (self.previous_token) {
						.local_identifier, .global_identifier => true,
						else => false,
					};

					var field_count = @as(Integer, 0);

					while (true) {
						try self.step();

						switch (self.current_token) {
							.newline => continue,

							.local_identifier => {
								// Create local copy of identifier because step() will overwrite captures.
								const interned_identifier =
									try chunk.intern_string(self.current_token.local_identifier);

								try chunk.emit_opcode(.push_interned);
								try chunk.emit_integer(interned_identifier);
								try self.step();

								switch (self.current_token) {
									.symbol_assign => {
										try self.step();
										try self.parse_expression(chunk);

										field_count += 1;
									},

									.symbol_brace_right => {
										try chunk.emit_opcode(.push_interned);
										try chunk.emit_integer(interned_identifier);

										field_count += 1;

										break;
									},

									.symbol_comma => {
										try chunk.emit_opcode(.push_interned);
										try chunk.emit_integer(interned_identifier);

										field_count += 1;
									},

									else => return error.UnexpectedToken,
								}
							},

							.symbol_brace_right => break,
							else => return error.UnexpectedToken,
						}
					}

					try chunk.emit_opcode(.push_table);
					try chunk.emit_integer(field_count);

					if (is_call_argument) {
						try chunk.emit_opcode(.call);
						try chunk.emit_integer(1);
					}

					self.step() catch |step_error| switch (step_error) {
						error.UnexpectedEnd => return,
					};
				},

				else => {
					try self.previous_token.expect_any(&.{.keyword_nil, .keyword_true, .keyword_false, .integer_literal,
						.real_literal, .string_literal, .global_identifier, .local_identifier, .symbol_brace_right,
							.symbol_paren_right});

					while (operators.pop()) |operator| {
						try chunk.emit_opcode(operator.opcode());
					}

					return;
				},
			}
		}
	}

	fn parse_operator(self: *Parser, chunk: *Chunk, operators: *OperatorStack, rhs_operator: Operator) ParseError!void {
		try self.previous_token.expect_any(operator_tokens);

		while (operators.pop()) |lhs_operator| {
			if (rhs_operator.precedence() < lhs_operator.precedence()) break try operators.push(lhs_operator);

			try chunk.emit_opcode(lhs_operator.opcode());
		}

		try operators.push(rhs_operator);
		try self.step();
	}

	fn parse_statement(self: *Parser, chunk: *Chunk) ParseError!bool {
		var local_depth = @as(usize, 0);

		while (true) {
			switch (self.current_token) {
				.newline => self.step() catch |step_error| switch (step_error) {
					error.UnexpectedEnd => return true,
				},

				.keyword_return => {
					try self.previous_token.expect(.newline);

					self.step() catch |step_error| switch (step_error) {
						error.UnexpectedEnd => return true,
					};

					try self.parse_expression(chunk);

					return false;
				},

				.local_identifier => {
					try self.previous_token.expect_any(&.{.newline, .symbol_period});
					try self.step();
				},

				.global_identifier => {
					try self.previous_token.expect(.newline);
					try self.step();
				},

				.symbol_period => switch (self.previous_token) {
					.global_identifier => {
						// Create local copy of identifier because step() will overwrite captures.
						const identifier = self.previous_token.global_identifier;

						try self.step();
						try self.current_token.expect(.local_identifier);
						try chunk.emit_opcode(.get_global);
						try chunk.emit_integer(try chunk.intern_string(identifier));

						local_depth += 1;
					},

					.local_identifier => {
						// Create local copy of identifier because step() will overwrite captures.
						const identifier = self.previous_token.local_identifier;

						try self.step();
						try self.current_token.expect(.local_identifier);

						if (local_depth == 0) {
							try chunk.emit_byte(self.resolve_local(identifier) orelse {
								return error.UndefinedLocal;
							});
						} else {
							try chunk.emit_opcode(.get_index);
							try chunk.emit_integer(try chunk.intern_string(identifier));
						}

						local_depth += 1;
					},

					else => return error.UnexpectedToken,
				},

				.symbol_assign => {
					try self.previous_token.expect(.local_identifier);

					const identifier = self.previous_token.local_identifier;

					if (local_depth == 0) {
						if (self.resolve_local(identifier)) |local_slot| {
							try chunk.emit_opcode(.set_local);
							try chunk.emit_byte(local_slot);
						} else {
							try self.declare_local(identifier);
						}
					} else {
						try chunk.emit_opcode(.set_index);
						try chunk.emit_integer(try chunk.intern_string(identifier));
					}

					try self.step();
					try self.parse_expression(chunk);

					local_depth = 0;
				},

				.symbol_arrow => {
					try self.parse_closure(chunk);

					local_depth = 0;
				},

				.symbol_paren_left => {
					switch (self.previous_token) {
						.local_identifier => |identifier| {
							if (local_depth == 0) {
								try chunk.emit_byte(self.resolve_local(identifier) orelse {
									return error.UndefinedLocal;
								});
							} else {
								try chunk.emit_opcode(.get_index);
								try chunk.emit_integer(try chunk.intern_string(identifier));
							}
						},

						.global_identifier => |identifier| {
							try chunk.emit_opcode(.get_global);
							try chunk.emit_integer(try chunk.intern_string(identifier));
						},

						else => return error.UnexpectedToken,
					}

					var argument_count = @as(Integer, 0);

					while (true) {
						try self.step();

						try switch (self.current_token) {
							.symbol_paren_right => break,
							else => self.parse_expression(chunk),
						};

						argument_count += 1;

						switch (self.current_token) {
							.symbol_paren_right => break,
							.symbol_comma => continue,
							else => return error.UnexpectedToken,
						}
					}

					try chunk.emit_opcode(.call);
					try chunk.emit_integer(argument_count);
					try chunk.emit_opcode(.pop);

					self.step() catch |step_error| switch (step_error) {
						error.UnexpectedEnd => return true,
					};

					local_depth = 0;
				},

				else => return error.UnexpectedToken,
			}
		}
	}

	fn resolve_local(self: *Parser, name: []const u8) ?u8 {
		var count = @as(u8, self.locals.buffer.len);

		while (count != 0) {
			const index = count - 1;

			if (coral.io.equals(name, self.locals.buffer[index].name)) return index;

			count = index;
		}

		return null;
	}

	fn step(self: *Parser) StepError!void {
		self.previous_token = self.current_token;
		self.current_token = self.tokenizer.next() orelse return error.UnexpectedEnd;

		@import("std").debug.print("{s}\n", .{self.current_token.text()});
	}
};

fn SmallStack(comptime Element: type, comptime default: Element) type {
	const maximum = 255;

	return struct {
		buffer: [maximum]Element = [_]Element{default} ** maximum,
		count: u8 = 0,

		const Self = @This();

		fn peek(self: Self) ?Element {
			if (self.count == 0) return null;

			return self.buffer[self.count - 1];
		}

		fn pop(self: *Self) ?Element {
			if (self.count == 0) return null;

			self.count -= 1;

			return self.buffer[self.count];
		}

		fn push(self: *Self, element: Element) !void {
			if (self.count == maximum) return error.OutOfMemory;

			self.buffer[self.count] = element;
			self.count += 1;
		}
	};
}

const SymbolTable = coral.table.Hashed(coral.table.string_key, usize);