const std = @import("std");

const Input = struct {
    compressed: []Entry,
    uncompressed: std.ArrayList(?u32),

    const Entry = struct { taken: u8, free: u8 };
};

pub fn parse(allocator: std.mem.Allocator, data: []const u8) !Input {
    var compressed = std.ArrayList(Input.Entry).init(allocator);
    var i: usize = 0;
    while (i < data.len) : (i += 2) {
        try compressed.append(.{ .taken = data[i] - '0', .free = if (i + 1 == data.len or data[i + 1] == '\n') 0 else data[i + 1] - '0' });
    }

    var uncompressed = std.ArrayList(?u32).init(allocator);
    for (0.., compressed.items) |j, entry| {
        try uncompressed.appendNTimes(@intCast(j), entry.taken);
        try uncompressed.appendNTimes(null, entry.free);
    }
    return .{ .compressed = compressed.items, .uncompressed = uncompressed };
}

const test_input = "2333133121414131402";

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 == 1928);
}

pub fn part1(allocator: std.mem.Allocator, input: Input) !u64 {
    var filesystem = std.ArrayList(u32).init(allocator);
    var unmoved: u32 = 0;
    var moved: u32 = @intCast(input.compressed.len);
    var moved_left: u8 = 0;
    while (unmoved < moved) {
        try filesystem.appendNTimes(unmoved, input.compressed[unmoved].taken);
        for (0..input.compressed[unmoved].free) |_| {
            while (moved_left == 0) {
                moved -= 1;
                moved_left = input.compressed[moved].taken;
            }
            try filesystem.append(moved);
            moved_left -= 1;
        }
        unmoved += 1;
    }
    try filesystem.appendNTimes(moved, moved_left);

    // for (filesystem.items) |item| {
    //     std.debug.print("{}", .{item});
    // }
    // std.debug.print("\n", .{});

    var sum: u64 = 0;
    for (0.., filesystem.items) |pos, id| {
        sum += pos * @as(u64, @intCast(id));
    }

    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 == 2858);
}

pub fn part2(allocator: std.mem.Allocator, input: Input) !u32 {
    var filesystem = std.ArrayList(u32).init(allocator);

    _ = .{ filesystem, input };
    filesystem = undefined;

    var sum: u64 = 0;
    for (0.., filesystem.items) |pos, id| {
        sum += pos * @as(u64, @intCast(id));
    }

    return sum;
}