const std = @import("std");

const Equation = struct { total: u64, numbers: []u64 };

const Input = []Equation;

pub fn parse(allocator: std.mem.Allocator, data: []const u8) !Input {
    var lines = std.mem.split(u8, data, "\n");
    var equations = std.ArrayList(Equation).init(allocator);
    while (lines.next()) |line| {
        if (line.len == 0) continue;
        var split = std.mem.split(u8, line, " ");
        const total = try std.fmt.parseInt(
            u64,
            std.mem.trimRight(u8, split.next() orelse return error.InvalidInput, ":"),
            10,
        );
        var numbers = std.ArrayList(u64).init(allocator);
        while (split.next()) |number| {
            try numbers.append(try std.fmt.parseInt(u64, number, 10));
        }
        try equations.append(.{ .total = total, .numbers = numbers.items });
    }
    return equations.items;
}

const test_input =
    \\190: 10 19
    \\3267: 81 40 27
    \\83: 17 5
    \\156: 15 6
    \\7290: 6 8 6 15
    \\161011: 16 10 13
    \\192: 17 8 14
    \\21037: 9 7 18 13
    \\292: 11 6 16 20
    \\
;

test "part1" {
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();

    const output = try part1(arena.allocator(), try parse(arena.allocator(), test_input));
    std.debug.print("got {}\n", .{output});
    std.debug.assert(output == 3749);
}

fn can_become(total: u64, numbers: []u64, target: u64) bool {
    // std.debug.print("can_become({}, {any}, {})\n", .{ total, numbers, target });
    if (numbers.len == 0) return total == target;
    const last = numbers.len - 1;
    return total >= numbers[last] and can_become(total - numbers[last], numbers[0..last], 0) or
        total % numbers[last] == 0 and can_become(total / numbers[last], numbers[0..last], 1);
}

pub fn part1(_: std.mem.Allocator, input: Input) !u64 {
    var sum: u64 = 0;
    for (input) |equation| {
        if (can_become(equation.total, equation.numbers, 0))
            sum += equation.total;
        // std.debug.print("\n", .{});
    }

    return sum;
}

test "part2" {
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();

    const output = try part2(arena.allocator(), try parse(arena.allocator(), test_input));
    std.debug.print("got {}\n", .{output});
    std.debug.assert(output == 11387);
}

fn can_become2(total: u64, numbers: []u64, target: u64) bool {
    // std.debug.print("can_become({}, {any}, {})\n", .{ total, numbers, target });
    if (numbers.len == 0) return total == target;
    const last = numbers.len - 1;
    // Addition
    return (total >= numbers[last] and can_become2(total - numbers[last], numbers[0..last], 0)) or
        // Multiplication
        (total % numbers[last] == 0 and can_become2(total / numbers[last], numbers[0..last], 1)) or blk: {
        // Concatenation
        const digit_count = std.fmt.count("{}", .{numbers[last]});
        const concat_mult = std.math.pow(u64, 10, digit_count);
        break :blk (total >= numbers[last] and (total - numbers[last]) % concat_mult == 0 and can_become2((total - numbers[last]) / concat_mult, numbers[0..last], 0));
    };
}

pub fn part2(_: std.mem.Allocator, input: Input) !u64 {
    var sum: u64 = 0;
    for (input) |equation| {
        if (can_become2(equation.total, equation.numbers, 0))
            sum += equation.total;
        // std.debug.print("\n", .{});
    }

    return sum;
}