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

const coral = @import("coral");

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

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

next: ?*const Self = null,
line: tokens.Line,

kind: union (enum) {
	top_expression: *const Expr,
	@"return": Return,
	declare: Declare,
	@"if": If,
	@"while": While,
},

pub const Declare = struct {
	declaration: *const tree.Declaration,
	initial_expression: *const Expr,
};

pub const If = struct {
	then_expression: *const Expr,
	@"then": *const Self,
	@"else": ?*const Self,
};

pub const Return = struct {
	returned_expression: ?*const Expr,
};

const Self = @This();

pub const While = struct {
	loop_expression: *const Expr,
	loop: *const Self,
};

pub fn parse(root: *tree.Root, stream: *tokens.Stream, environment: *tree.Environment) tree.ParseError!*Self {
	switch (stream.token) {
		.keyword_return => {
			stream.step();

			if (stream.token != .end and stream.token != .newline) {
				return root.create_node(Self{
					.line = stream.line,
					.kind = .{.@"return" = .{.returned_expression = try Expr.parse(root, stream, environment)}},
				});
			}

			if (stream.token != .end and stream.token != .newline) {
				return root.report_error(stream.line, "expected end or newline after return statement", .{});
			}

			return root.create_node(Self{
				.line = stream.line,
				.kind = .{.@"return" = .{.returned_expression = null}},
			});
		},

		.keyword_while => {
			defer stream.skip_newlines();

			stream.step();

			const condition_expression = try Expr.parse(root, stream, environment);

			if (stream.token != .symbol_colon) {
				return root.report_error(stream.line, "expected `:` after `while` statement", .{});
			}

			stream.skip_newlines();

			const first_statement = try parse(root, stream, environment);

			{
				var current_statement = first_statement;

				while (stream.token != .keyword_end) {
					const next_statement = try parse(root, stream, environment);

					current_statement.next = next_statement;
					current_statement = next_statement;
				}
			}

			return root.create_node(Self{
				.line = stream.line,

				.kind = .{
					.@"while" = .{
						.loop = first_statement,
						.loop_expression = condition_expression,
					},
				},
			});
		},

		.keyword_var, .keyword_let => {
			const is_constant = stream.token == .keyword_let;

			stream.skip_newlines();

			const identifier = switch (stream.token) {
				.identifier => |identifier| identifier,
				else => return root.report_error(stream.line, "expected identifier after declaration", .{}),
			};

			stream.skip_newlines();

			if (stream.token != .symbol_equals) {
				return root.report_error(stream.line, "expected `=` after declaration `{identifier}`", .{
					.identifier = identifier,
				});
			}

			stream.skip_newlines();

			return root.create_node(Self{
				.line = stream.line,

				.kind = .{
					.declare = .{
						.initial_expression = try Expr.parse(root, stream, environment),

						.declaration = declare: {
							if (is_constant) {
								break: declare environment.declare_constant(identifier) catch |declaration_error| {
									return root.report_declare_error(stream.line, identifier, declaration_error);
								};
							}

							break: declare environment.declare_variable(identifier) catch |declaration_error| {
								return root.report_declare_error(stream.line, identifier, declaration_error);
							};
						},
					},
				},
			});
		},

		.keyword_if => return parse_branch(root, stream, environment),

		else => return root.create_node(Self{
			.line = stream.line,
			.kind = .{.top_expression = try Expr.parse(root, stream, environment)},
		}),
	}
}

fn parse_branch(root: *tree.Root, stream: *tokens.Stream, environment: *tree.Environment) tree.ParseError!*Self {
	stream.step();

	const expression = try Expr.parse(root, stream, environment);

	if (stream.token != .symbol_colon) {
		return root.report_error(stream.line, "expected `:` after `{token}`", .{.token = stream.token.text()});
	}

	stream.skip_newlines();

	const first_then_statement = try parse(root, stream, environment);
	var current_then_statement = first_then_statement;

	while (true) {
		switch (stream.token) {
			.keyword_end => {
				stream.skip_newlines();

				return root.create_node(Self{
					.line = stream.line,

					.kind = .{
						.@"if" = .{
							.then_expression = expression,
							.@"then" = first_then_statement,
							.@"else" = null,
						},
					},
				});
			},

			.keyword_else => {
				stream.step();

				if (stream.token != .symbol_colon) {
					return root.report_error(stream.line, "expected `:` after `if` statement condition", .{});
				}

				stream.skip_newlines();

				const first_else_statement = try parse(root, stream, environment);
				var current_else_statement = first_else_statement;

				while (stream.token != .keyword_end) {
					const next_statement = try parse(root, stream, environment);

					current_else_statement.next = next_statement;
					current_else_statement = next_statement;
				}

				stream.skip_newlines();

				return root.create_node(Self{
					.line = stream.line,

					.kind = .{
						.@"if" = .{
							.@"else" = first_else_statement,
							.@"then" = first_then_statement,
							.then_expression = expression,
						},
					}
				});
			},

			.keyword_elif => {
				return root.create_node(Self{
					.line = stream.line,

					.kind = .{
						.@"if" = .{
							.@"else" = try parse_branch(root, stream, environment),
							.@"then" = first_then_statement,
							.then_expression = expression,
						},
					},
				});
			},

			else => {
				const next_statement = try parse(root, stream, environment);

				current_then_statement.next = next_statement;
				current_then_statement = next_statement;
			},
		}
	}
}