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

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

fn Fmt(T: type) type {
    return std.fmt.Formatter(struct {
        fn format(
            data: struct { T, []const u8, usize },
            comptime _: []const u8,
            _: std.fmt.FormatOptions,
            writer: anytype,
        ) !void {
            const self, const file_source, const indent = data;
            return self.format(writer, file_source, indent);
        }
    }.format);
}

pub fn fmt(tree: anytype, source: []const u8, indent: usize) Fmt(@TypeOf(tree)) {
    return .{ .data = .{ tree, source, indent } };
}

pub const Block = struct {
    loc: Lexer.Location,
    stmts: []Stmt,

    fn format(self: Block, writer: anytype, source: []const u8, indent: usize) !void {
        try writer.writeAll("{\n");
        for (self.stmts) |stmt| {
            try writer.print("{}\n", .{fmt(stmt, source, indent + 4)});
        }
        try writer.writeByteNTimes(' ', indent);
        try writer.writeAll("}");
    }
};

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

    pub const Type = union(enum) {
        expr: *const Expr,
        assign_var: AssignVar,
        block: Block,

        pub const AssignVar = struct {
            ident: Lexer.Location,
            is_decl: bool,
            value: *const Expr,
        };
    };

    fn format(self: Stmt, writer: anytype, source: []const u8, indent: usize) !void {
        try writer.writeByteNTimes(' ', indent);
        return switch (self.type) {
            .expr => |expr| writer.print("{}", .{fmt(expr, source, indent)}),
            .block => |b| writer.print("{}", .{fmt(b, source, indent)}),
            .assign_var => |assign_var| writer.print("{s}{s} = {}", .{
                if (assign_var.is_decl) "let " else "",
                assign_var.ident.getIdent(source),
                fmt(assign_var.value, source, indent),
            }),
        };
    }
};

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

    pub const Type = union(enum) {
        integer_literal,
        bin_op: BinOp,
        call: Call,
        identifier,
        @"if": If,

        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 const If = struct {
            cond: *const Expr,
            then: Block,
            @"else": ?Block,
        };
    };

    fn format(self: Expr, writer: anytype, source: []const u8, indent: usize) !void {
        switch (self.type) {
            .integer_literal => try writer.print("{}", .{self.loc.getInt(source)}),
            .bin_op => |bin_op| {
                try writer.print("{} {} {}", .{ fmt(bin_op.lhs, source, indent), bin_op.op, fmt(bin_op.rhs, source, indent) });
            },
            .call => |call| {
                try writer.print("{}({})", .{ fmt(call.proc, source, indent), fmt(call.arg, source, indent) });
            },
            .identifier => try writer.print("{s}", .{self.loc.getIdent(source)}),
            .@"if" => |@"if"| {
                try writer.print("if {} {}", .{
                    fmt(@"if".cond, source, indent),
                    fmt(@"if".then, source, indent),
                });
                if (@"if".@"else") |@"else"| {
                    try writer.print(" else {}", .{fmt(@"else", source, indent)});
                }
            },
        }
    }
};

const ParseError = error{ OutOfMemory, ExpectedRightParen, UnexpectedToken, InvalidAssignTarget };

pub fn block(allocator: Allocator, lexer: *Lexer) !Block {
    const left_curly = try mustEat(lexer, .left_curly);
    var stmts: std.ArrayList(Stmt) = .init(allocator);
    while (lexer.peek().type != .right_curly) {
        try stmts.append(try statement(allocator, lexer));
    }
    const right_curly = try mustEat(lexer, .right_curly);
    return .{
        .loc = left_curly.loc.combine(right_curly.loc),
        .stmts = try stmts.toOwnedSlice(),
    };
}

pub fn statement(allocator: Allocator, lexer: *Lexer) ParseError!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);
            return .{
                .loc = let.loc.combine(value.loc),
                .type = .{ .assign_var = .{ .ident = ident.loc, .is_decl = true, .value = value } },
            };
        },
        .left_curly => {
            const b = try block(allocator, lexer);
            return .{
                .loc = b.loc,
                .type = .{ .block = b },
            };
        },
        else => {
            const lhs = try expression(allocator, lexer);
            if (lexer.peek().type == .equal) {
                _ = try mustEat(lexer, .equal);
                const value = try expression(allocator, lexer);
                if (lhs.type != .identifier) {
                    std.debug.print("Invalid assign target. Found '{s}', expected an identifier.", .{@tagName(lhs.type)});
                    return error.InvalidAssignTarget;
                }
                return .{
                    .loc = lhs.loc,
                    .type = .{ .assign_var = .{ .ident = lhs.loc, .is_decl = false, .value = value } },
                };
            }
            return .{
                .loc = lhs.loc,
                .type = .{ .expr = lhs },
            };
        },
    }
}

pub fn expression(allocator: Allocator, lexer: *Lexer) ParseError!*Expr {
    return parseTerms(allocator, lexer);
}

pub fn parseTerms(allocator: Allocator, lexer: *Lexer) !*Expr {
    var lhs = try parseIf(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 parseIf(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 parseIf(allocator: Allocator, lexer: *Lexer) !*Expr {
    switch (lexer.peek().type) {
        .@"if" => {
            const @"if" = lexer.next();
            const cond = try expression(allocator, lexer);
            const then = try block(allocator, lexer);
            const @"else" = if (lexer.peek().type == .@"else") blk: {
                _ = lexer.next();
                break :blk try block(allocator, lexer);
            } else null;
            return try allocate(Expr, allocator, .{
                .loc = @"if".loc.combine((@"else" orelse then).loc),
                .type = .{ .@"if" = .{
                    .cond = cond,
                    .then = then,
                    .@"else" = @"else",
                } },
            });
        },
        else => return try parseInvocations(allocator, lexer),
    }
}

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 '{s}'\n", .{@tagName(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;
}