pub fn day9_part1(input: &str) -> String {
    let diskmap = parse(input);
    let mut blocks = diskmap_to_blocks(&diskmap);
    compact(&mut blocks);
    let sum = checksum(&blocks);
    sum.to_string()
}

pub fn day9_part2(input: &str) -> String {
    let diskmap = parse(input);
    let compacted_diskmap = defrag_compact(&diskmap);
    let blocks = compacted_diskmap_to_blocks(&compacted_diskmap);
    let sum = checksum(&blocks);
    sum.to_string()
}

#[derive(Copy, Clone, PartialEq, Eq, Debug)]
struct File {
    id: u64,
    length: usize,
    free_space: usize,
}

fn compacted_diskmap_to_blocks(compacted_diskmap: &[File]) -> Vec<Option<u64>> {
    compacted_diskmap
        .iter()
        .flat_map(|file| {
            [
                vec![Some(file.id); file.length],
                vec![None; file.free_space],
            ]
        })
        .flatten()
        .collect()
}

fn defrag_compact(diskmap: &[usize]) -> Vec<File> {
    let files = diskmap_to_files(diskmap);
    let mut compacted_files = files.clone();

    for file in files.clone().into_iter().rev() {
        let Some(insertable_index) = compacted_files
            .iter()
            .enumerate()
            .find(|(_, insertable_file)| insertable_file.free_space >= file.length)
            .map(|(i, _)| i)
        else {
            continue;
        };

        let prev_free_space = compacted_files[insertable_index].free_space;
        compacted_files[insertable_index].free_space = 0;
        compacted_files.insert(insertable_index + 1, file);
        compacted_files[insertable_index + 1].free_space = prev_free_space - file.length;

        let nullable_index = compacted_files
            .iter()
            .enumerate()
            .rev()
            .find(|(_, nullable_file)| nullable_file.id == file.id)
            .map(|(i, _)| i)
            .unwrap();
        compacted_files[nullable_index - 1].free_space +=
            compacted_files[nullable_index].length + compacted_files[nullable_index].free_space;
        compacted_files.remove(nullable_index);
    }

    compacted_files
}

fn diskmap_to_files(diskmap: &[usize]) -> Vec<File> {
    let mut id = 0u64;
    let mut blocks = vec![];

    for digit in diskmap {
        if id % 2 == 0 {
            blocks.push(File {
                id: id / 2,
                length: *digit,
                free_space: 0,
            });
        } else {
            blocks.last_mut().unwrap().free_space = *digit;
        }
        id += 1
    }

    blocks
}

fn checksum(blocks: &[Option<u64>]) -> u64 {
    blocks
        .iter()
        .enumerate()
        .map(|(index, block)| if let Some(block) = block {index as u64 * block} else {0})
        .sum::<u64>()
}

fn diskmap_to_blocks(diskmap: &[usize]) -> Vec<Option<u64>> {
    let mut id = 0u64;
    let mut blocks = vec![];

    for digit in diskmap {
        if id % 2 == 0 {
            blocks.extend_from_slice(&vec![Some(id / 2); *digit]);
        } else {
            blocks.extend_from_slice(&vec![None; *digit]);
        }
        id += 1
    }

    blocks
}

fn compact(blocks: &mut [Option<u64>]) {
    let mut start_index = 0;
    let mut end_index = blocks.len() - 1;

    while start_index < end_index {
        if blocks[start_index].is_some() {
            start_index += 1;
        } else if blocks[end_index].is_none() {
            end_index -= 1
        } else {
            blocks[start_index] = blocks[end_index];
            blocks[end_index] = None;
        }
    }
}

fn parse(input: &str) -> Vec<usize> {
    input
        .chars()
        .map(|char| char.to_digit(10).unwrap() as usize)
        .collect()
}