const std = @import("std");
const Allocator = std.mem.Allocator;

const root = @import("root");
const Lexer = root.Lexer;
const Token = root.Lexer.Token;

pub const Expr = struct {
    loc: Lexer.Location,
    type: Type,

    pub const Type = union(enum) {
        integer_literal,
        bin_op: BinOp,
        invalid: Token,
        call: Call,
        identifier,

        pub const BinOp = struct {
            lhs: *const Expr,
            rhs: *const Expr,
            op: Op,

            const Op = enum {
                plus,
                minus,

                pub fn format(self: Op, comptime fmt: []const u8, options: std.fmt.FormatOptions, writer: anytype) !void {
                    _ = fmt;
                    _ = options;
                    try writer.writeByte(switch (self) {
                        .plus => '+',
                        .minus => '-',
                    });
                }
            };
        };

        pub const Call = struct {
            proc: *const Expr,
            arg: *const Expr,
        };
    };

    pub fn format(self: Expr, comptime fmt: []const u8, options: std.fmt.FormatOptions, writer: anytype) !void {
        _ = fmt;
        _ = options;
        switch (self.type) {
            .integer_literal => try writer.print("<int>", .{}),
            .bin_op => |bin_op| try writer.print("({} {} {})", .{ bin_op.lhs, bin_op.op, bin_op.rhs }),
            .invalid => try writer.print("<invalid>", .{}),
            .call => |call| try writer.print("({} {})", .{ call.proc, call.arg }),
            .identifier => try writer.print("<identifier>", .{}),
        }
    }
};

pub fn expression(allocator: Allocator, lexer: *Lexer) error{OutOfMemory}!*Expr {
    return addExpr(allocator, lexer);
}

pub fn addExpr(allocator: Allocator, lexer: *Lexer) !*Expr {
    var lhs = try callExpr(allocator, lexer);
    while (true) {
        const op: Expr.Type.BinOp.Op = switch (lexer.peek().type) {
            .plus => .plus,
            .minus => .minus,
            else => break,
        };
        _ = lexer.next();

        const rhs = try callExpr(allocator, lexer);
        lhs = try allocate(allocator, .{
            .loc = lhs.loc.combine(rhs.loc),
            .type = .{ .bin_op = .{ .lhs = lhs, .op = op, .rhs = rhs } },
        });
    }
    return lhs;
}

pub fn callExpr(allocator: Allocator, lexer: *Lexer) !*Expr {
    var proc = try primaryExpr(allocator, lexer);
    while (true) {
        switch (lexer.peek().type) {
            .left_paren, .integer_literal => {},
            else => break,
        }
        const arg = try primaryExpr(allocator, lexer);
        proc = try allocate(allocator, .{
            .loc = proc.loc.combine(arg.loc),
            .type = .{ .call = .{ .proc = proc, .arg = arg } },
        });
    }
    return proc;
}

pub fn primaryExpr(allocator: Allocator, lexer: *Lexer) !*Expr {
    const token = lexer.next();
    // std.debug.print("term {}\n", .{token});
    return allocate(allocator, switch (token.type) {
        .left_paren => {
            const res = expression(allocator, lexer);
            const right_paren = lexer.next();
            if (right_paren.type != .right_paren)
                return allocate(allocator, .{
                    .loc = right_paren.loc,
                    .type = .{ .invalid = right_paren },
                });
            return res;
        },
        .integer_literal => .{ .loc = token.loc, .type = .integer_literal },
        .identifier => .{ .loc = token.loc, .type = .identifier },
        else => .{ .loc = token.loc, .type = .{ .invalid = token } },
    });
}

fn allocate(allocator: Allocator, expr: Expr) !*Expr {
    const res = try allocator.create(Expr);
    res.* = expr;
    return res;
}