const std = @import("std");

const Input = std.ArrayList(u64);

const test_input =
    \\125 17
    \\
;

pub fn parse(allocator: std.mem.Allocator, data: []const u8) !Input {
    var it = std.mem.splitAny(u8, data, " \n");
    var numbers = std.ArrayList(u64).init(allocator);
    while (it.next()) |num| {
        if (num.len == 0) continue;
        const n = try std.fmt.parseInt(u64, num, 10);
        try numbers.append(n);
    }
    return numbers;
}

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));
    try std.testing.expectEqual(55312, output);
}

fn stoneCount(allocator: std.mem.Allocator, number: u64, rounds: u8) !u64 {
    if (rounds == 0) return 1;
    if (number == 0) return stoneCount(allocator, 1, rounds - 1);

    var num_str = try std.fmt.allocPrint(allocator, "{}", .{number});
    if (num_str.len % 2 == 0) {
        const a = try std.fmt.parseInt(u64, num_str[0 .. num_str.len / 2], 10);
        const b = try std.fmt.parseInt(u64, num_str[num_str.len / 2 ..], 10);
        allocator.free(num_str);
        return try stoneCount(allocator, a, rounds - 1) +
            try stoneCount(allocator, b, rounds - 1);
    } else {
        allocator.free(num_str);
        return stoneCount(allocator, number * 2024, rounds - 1);
    }
}

pub fn part1(allocator: std.mem.Allocator, input: Input) !u64 {
    var count: u64 = 0;
    for (input.items) |stone| {
        count += try stoneCount(allocator, stone, 25);
    }
    return count;
}

fn stoneCoundCachedImpl(
    allocator: std.mem.Allocator,
    cache: *std.AutoHashMap(u64, [75]u64),
    number: u64,
    rounds: u8,
) error{ ParseIntError, OutOfMemory, Overflow, InvalidCharacter }!u64 {
    if (number == 0) return stoneCountCached(allocator, cache, 1, rounds - 1);

    var num_str = try std.fmt.allocPrint(allocator, "{}", .{number});
    if (num_str.len % 2 == 0) {
        const a = try std.fmt.parseInt(u64, num_str[0 .. num_str.len / 2], 10);
        const b = try std.fmt.parseInt(u64, num_str[num_str.len / 2 ..], 10);
        allocator.free(num_str);
        return try stoneCountCached(allocator, cache, a, rounds - 1) +
            try stoneCountCached(allocator, cache, b, rounds - 1);
    } else {
        allocator.free(num_str);
        return stoneCountCached(allocator, cache, number * 2024, rounds - 1);
    }
}

fn stoneCountCached(
    allocator: std.mem.Allocator,
    cache: *std.AutoHashMap(u64, [75]u64),
    number: u64,
    rounds: u8,
) !u64 {
    if (rounds == 0) return 1;
    if (cache.get(number)) |val| if (val[rounds - 1] != 0) return val[rounds - 1];

    const val = try stoneCoundCachedImpl(allocator, cache, number, rounds);
    const res = try cache.getOrPut(number);
    if (!res.found_existing) @memset(res.value_ptr, 0);
    res.value_ptr[rounds - 1] = val;
    return val;
}

pub fn part2(allocator: std.mem.Allocator, input: Input) !u64 {
    var cache = std.AutoHashMap(u64, [75]u64).init(allocator);
    var count: u64 = 0;
    for (input.items) |stone| {
        count += try stoneCountCached(allocator, &cache, stone, 75);
    }
    return count;
}