#include "aoc.h"

typedef struct {
    int id; // -1 if free space
    int len;
    int offset;
} Block;

static void
dump(int *ids, int len) {
    for (int i = 0; i < len; i++) {
        if (ids[i] == -1) {
            printf(".");
        }
        else {
            printf("%d", ids[i] % 10);
        }
    }
    puts("");
}

static void
disk_from_blocks(i32 *disk, Block *blocks, usize block_count) {
    for (usize i = 0; i < block_count; i++) {
        Block *block = blocks + i;
        for (int j = 0; j < block->len; j++) {
            disk[j] = block->id;
        }
        disk += block->len;
    }
}

int main(int argc, char **argv) {
    (void) argc;
    Arena *arena = make_arena(Megabytes(10));
    str input = str_trim(read_file(arena, argv[1]));

    enum {
        STATE_BLOCK,
        STATE_SPACE,
    } state = STATE_BLOCK;
    int block_id = 0;

    DYNAMIC_ARRAY(Block) blocks = {0};

    isize total_size = 0;
    for (usize i = 0; i < input.len; ++i) {
        int len = input.data[i] - '0';
        *push(&blocks, arena) = (Block) {
            .id = (state == STATE_BLOCK) ? block_id : -1, 
            .len = len, 
            .offset = total_size
        };
        total_size += len;
        if (state == STATE_BLOCK) {
            block_id++;
            state = STATE_SPACE;
        }
        else {
            state = STATE_BLOCK;
        }
    }
    ASSERT(state == STATE_SPACE);

    i32 *disk = ARENA_ALLOC_ARRAY(arena, i32, total_size);
    disk_from_blocks(disk, blocks.data, blocks.len);

    isize to = 0, from = total_size - 1;
    for (;;) {
        while (to <= from && disk[to] != -1) {
            to++;
        }
        while (from > to && disk[from] == -1) {
            from--;
        }
        if (to >= from) {
            break;
        }
        ASSERT((disk[to] == -1) && (disk[from] != -1));
        disk[to] = disk[from];
        disk[from] = -1;
        from--;
        to++;
    }

    i64 part_1 = 0;
    for (isize i = 0; i < total_size && disk[i] != -1; i++) {
        part_1 += i * disk[i];
    }
    printf("%ld\n", part_1);

    // restore original layout
    disk_from_blocks(disk, blocks.data, blocks.len);

    for (isize from_block_index = blocks.len - 1; from_block_index >= 0; from_block_index -= 2) {
        Block *from = blocks.data + from_block_index;
        ASSERT(from->id != -1);

        for (isize to_block_index = 1; to_block_index < from_block_index; to_block_index += 2) {
            Block *to = blocks.data + to_block_index;
            ASSERT(to->id == -1);
            if (from->len <= to->len) {
                for (isize i = 0; i < from->len; i++) {
                    disk[to->offset + i] = from->id;
                    disk[from->offset + i] = -1;
                }
                to->offset += from->len;
                to->len -= from->len;
                break;
            }
        }
    }

    i64 part_2 = 0;
    for (isize i = 0; i < total_size; i++) {
        if (disk[i] != -1) {
            part_2 += i * disk[i];
        }
    }
    printf("%ld\n", part_2);
}