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

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

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

    pub const Type = union(enum) {
        expr: *const Expr,
    };

    pub fn fmt(self: Stmt, file_source: []const u8) Format {
        return .{ .data = .{ self, file_source } };
    }

    const Format = std.fmt.Formatter(struct {
        fn format(data: struct { Stmt, []const u8 }, comptime f: []const u8, options: std.fmt.FormatOptions, writer: anytype) !void {
            const self, const file_source = data;
            _ = f;
            _ = options;
            return switch (self.type) {
                .expr => |expr| writer.print("{};", .{expr.fmt(file_source)}),
            };
        }
    }.format);
};

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

    pub const Type = union(enum) {
        integer_literal,
        bin_op: BinOp,
        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 f: []const u8, options: std.fmt.FormatOptions, writer: anytype) !void {
                    _ = f;
                    _ = options;
                    try writer.writeByte(switch (self) {
                        .plus => '+',
                        .minus => '-',
                    });
                }
            };
        };

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

    pub fn fmt(self: Expr, file_source: []const u8) Format {
        return .{ .data = .{ self, file_source } };
    }

    const Format = std.fmt.Formatter(struct {
        fn format(data: struct { Expr, []const u8 }, comptime f: []const u8, options: std.fmt.FormatOptions, writer: anytype) !void {
            const self, const file_source = data;
            _ = f;
            _ = options;
            switch (self.type) {
                .integer_literal => try writer.print("{}", .{self.loc.getInt(file_source)}),
                .bin_op => |bin_op| try writer.print("({} {} {})", .{ bin_op.lhs.fmt(file_source), bin_op.op, bin_op.rhs.fmt(file_source) }),
                .call => |call| try writer.print("({} {})", .{ call.proc.fmt(file_source), call.arg.fmt(file_source) }),
                .identifier => try writer.print("{s}", .{self.loc.getIdent(file_source)}),
            }
        }
    }.format);
};

pub fn statements(allocator: Allocator, lexer: *Lexer) ![]Stmt {
    var stmts: std.ArrayList(Stmt) = .init(allocator);
    while (lexer.peek().type != .eof) {
        try stmts.append(try statement(allocator, lexer));
    }
    return try stmts.toOwnedSlice();
}

pub fn statement(allocator: Allocator, lexer: *Lexer) !Stmt {
    var expr = try expression(allocator, lexer);
    const semicolon: Lexer.Token = if (lexer.peek().type == .semicolon) lexer.next() else return error.ExpectedSemicolon;
    return .{
        .loc = expr.loc.combine(semicolon.loc),
        .type = .{ .expr = expr },
    };
}

pub fn expression(allocator: Allocator, lexer: *Lexer) error{ OutOfMemory, ExpectedRightParen, UnexpectedToken }!*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(Expr, 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, .identifier => {},
            else => break,
        }
        const arg = try primaryExpr(allocator, lexer);
        proc = try allocate(Expr, 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();
    return allocate(Expr, allocator, switch (token.type) {
        .left_paren => {
            const res = expression(allocator, lexer);
            const right_paren = lexer.next();
            if (right_paren.type != .right_paren)
                return error.ExpectedRightParen;
            return res;
        },
        .integer_literal => .{ .loc = token.loc, .type = .integer_literal },
        .identifier => .{ .loc = token.loc, .type = .identifier },
        else => |t| {
            std.debug.print("Expected '(', integer literal, or identifier. Got {}\n", .{t});
            return error.UnexpectedToken;
        },
    });
}

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