fn main() {
    let input = include_str!("../input.txt").trim();
    let time = std::time::Instant::now();

    let mut memory = input.chars().enumerate()
        .filter_map(|(i, c)| {
            c.to_digit(10).and_then(|num| {
                if num == 0 {
                    None
                } else if i % 2 == 0 {
                    Some((Some(i as u32 / 2), num))
                } else {
                    Some((None, num))
                }
            })
        })
        .into_iter()
        .collect::<Vec<(Option<u32>, u32)>>();

    let mut i = memory.len() - 1;
    while i > 1 {
        if let (Some(id), files) = memory[i] {
            for j in 0..i {
                if let (None, free_space) = memory[j] {
                    if free_space < files {
                        continue;
                    }

                    if free_space == files {
                        // Exact match
                        memory.swap(i, j);

                    } else if free_space > files {
                        // Partial match
                        memory[j] = (Some(id), files);
                        memory.insert(j + 1, (None, free_space - files));
                        i += 1;
                        memory[i] = (None, files);

                        // Conveniently not needed because of how the data is structured!
                        // if memory[j + 2].0.is_none() {
                        //     memory[j + 1] = (None, free_space - files + memory[j + 2].1);
                        //     memory.remove(j + 1);
                        // }
                    }

                    let mut new_free_space = files;

                    // Merge free space with free space to the left
                    if let (None, prev_space) = memory[i - 1] {
                        new_free_space += prev_space;
                        memory[i] = (None, new_free_space);
                        memory.remove(i - 1);
                        i -= 1;
                    }

                    // Merge free space with free space to the right
                    if let Some((None, next_space)) = memory.get(i + 1) {
                        new_free_space += next_space;
                        memory[i] = (None, new_free_space);
                        memory.remove(i + 1);
                    }

                    break;
                }
            }
        }

        i -= 1;
    }

    let mut sum: i64 = 0;
    let mut index: u32 = 0;

    for data in memory {
        if let (Some(id), len) = data {
            for _ in 0..len {
                sum += (id * index) as i64;
                index += 1;
            }
        } else if let (None, len) = data {
            index += len;
        }
    }

    println!("Time elapsed: {:?}", time.elapsed());
    println!("Sum: {}", sum);
}