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,
        declare_var: DeclareVar,

        pub const DeclareVar = struct {
            ident: Lexer.Location,
            value: *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)}),
                .declare_var => |declare_var| writer.print("let {s} = {};", .{ declare_var.ident.getIdent(file_source), declare_var.value.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 _: []const u8, _: std.fmt.FormatOptions, writer: anytype) !void {
                    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 _: []const u8, _: std.fmt.FormatOptions, writer: anytype) !void {
            const self, const file_source = data;
            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 {
    switch (lexer.peek().type) {
        .let => {
            const let = lexer.next();
            const ident = try mustEat(lexer, .identifier);
            _ = try mustEat(lexer, .equal);
            const value = try expression(allocator, lexer);
            const semicolon = try mustEat(lexer, .semicolon);
            return .{
                .loc = let.loc.combine(semicolon.loc),
                .type = .{ .declare_var = .{ .ident = ident.loc, .value = value } },
            };
        },
        else => {
            var expr = try expression(allocator, lexer);
            const semicolon = lexer.next();
            if (semicolon.type != .semicolon) 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 parseTerms(allocator, lexer);
}

pub fn parseTerms(allocator: Allocator, lexer: *Lexer) !*Expr {
    var lhs = try parseInvocations(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 parseInvocations(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 parseInvocations(allocator: Allocator, lexer: *Lexer) !*Expr {
    var proc = try parsePrimaryExpr(allocator, lexer);
    while (true) {
        if (lexer.peek().type != .left_paren) break;
        _ = lexer.next();
        const arg = try expression(allocator, lexer);
        _ = try mustEat(lexer, .right_paren);
        proc = try allocate(Expr, allocator, .{
            .loc = proc.loc.combine(arg.loc),
            .type = .{ .call = .{ .proc = proc, .arg = arg } },
        });
    }
    return proc;
}

pub fn parsePrimaryExpr(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 mustEat(lexer: *Lexer, ty: Lexer.Token.Type) !Lexer.Token {
    const token = lexer.next();
    if (token.type != ty) {
        std.debug.print("Expected {}. Got {}\n", .{ ty, token.type });
        return error.UnexpectedToken;
    }
    return token;
}

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