const std = @import("std");

const Input = struct {
    compressed: []Entry,

    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',
        });
    }
    return .{ .compressed = compressed.items };
}

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) !u64 {
    var filesystem_array_list = std.ArrayList(?u32).init(allocator);
    for (0.., input.compressed) |i, entry| {
        try filesystem_array_list.appendNTimes(@as(u32, @intCast(i)), entry.taken);
        try filesystem_array_list.appendNTimes(null, entry.free);
    }
    var filesystem = filesystem_array_list.items;

    var file: u32 = @intCast(input.compressed.len - 1);
    while (true) : (file = std.math.sub(u32, file, 1) catch break) {
        // for (filesystem) |item| {
        //     if (item) |it| std.debug.print("{}", .{it}) else std.debug.print(".", .{});
        // }
        // std.debug.print("\n", .{});

        const index = std.mem.indexOfScalar(?u32, filesystem, file).?;
        const len = input.compressed[file].taken;
        std.debug.assert(std.mem.allEqual(?u32, filesystem[index..][0..len], file));
        const nulls = try allocator.alloc(?u32, len);
        defer allocator.free(nulls);
        @memset(nulls, null);
        const space = std.mem.indexOf(?u32, filesystem, nulls) orelse continue;
        if (space > index) continue;
        @memset(filesystem[index..][0..len], null);
        @memset(filesystem[space..][0..len], file);
    }

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

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

    return sum;
}

pub fn part2OptimizedButWrong(allocator: std.mem.Allocator, input: Input) !u64 {
    var filesystem = std.ArrayList(?u32).init(allocator);
    for (input.compressed) |entry| {
        try filesystem.appendNTimes(null, entry.taken);
        try filesystem.appendNTimes(null, entry.free);
    }

    var file_locations = std.ArrayList(u32).init(allocator);
    var free_slots = std.ArrayList(struct { start: u32, size: u8 }).init(allocator);
    var i: u32 = 0;
    for (input.compressed) |e| {
        try file_locations.append(i);
        try free_slots.append(.{ .start = i + e.taken, .size = e.free });
        i += e.taken + e.free;
    }

    var file_id: u32 = @intCast(input.compressed.len - 1);
    while (true) {
        const file_size = input.compressed[file_id].taken;
        for (free_slots.items) |*slot| {
            if (file_size <= slot.size) {
                for (filesystem.items[slot.start..][0..file_size]) |id| std.debug.assert(id == null);
                @memset(filesystem.items[slot.start..][0..file_size], file_id);
                slot.size -= file_size;
                slot.start += file_size;
                break;
            }
        } else {
            for (filesystem.items[file_locations.items[file_id]..][0..file_size]) |id| std.debug.assert(id == null);
            @memset(filesystem.items[file_locations.items[file_id]..][0..file_size], file_id);
        }
        file_id = std.math.sub(u32, file_id, 1) catch break;
    }

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

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

    return sum;
}