use advent_of_code::strings::char_to_num;
use advent_of_code_macros::include_aoc;

include_aoc!(2024, 09);

type Disk = Vec<Option<u64>>;

fn print_disk(disk: Disk, columns: bool) {
    if columns {
        for i in 1..=disk.len() {
            print!("{:02}", i);
        }
        println!()
    }
    for block in disk {
        match block {
            None => print!("  "),
            Some(x) => print!("{} ", x),
        }
    }
    println!(";")
}

fn layout_to_disk(layout: Vec<u64>) -> Disk {
    let mut disk = Vec::new();
    let mut space = false;
    let mut block;
    let mut blockid = 0;
    for blocksize in layout {
        block = if space {
            None
        } else {
            blockid += 1;
            Some(blockid - 1)
        };
        for _ in 0..blocksize {
            disk.push(block);
        }
        space = !space;
    }
    disk
}

fn calculate_checksum(disk: Disk) -> u64 {
    let iter = disk.iter();
    let mut checksum = 0;
    let mut position = 0;
    for block in iter {
        if block.is_some() {
            checksum += block.unwrap() * position;
        }
        position += 1;
    }
    checksum
}

fn compact_blocks(disk: &mut Disk) {
    let mut lastblock = disk.len() - 1;
    let mut block = 0;
    loop {
        if disk[block].is_none() {
            if disk[lastblock].is_some() {
                disk[block] = disk[lastblock];
                disk[lastblock] = None;
                block += 1;
            }
            lastblock -= 1;
        } else {
            block += 1
        }
        if lastblock <= block {
            break;
        }
    }
}

fn move_blocks(disk: &mut Disk, file_start: usize, length: usize, new_start: usize) {
    for i in 0..=length {
        disk[new_start + i] = disk[file_start - length + i];
        disk[file_start - length + i] = None;
    }
}

fn compact_files(disk: &mut Disk) {
    let mut length;
    let mut fileid;
    let mut spaceblock;
    let mut block;
    let mut endblock = disk.len();
    //print_disk(disk.clone(), true);
    loop {
        endblock -= 1;
        if endblock == 0 {
            break;
        }
        if disk[endblock].is_some() {
            fileid = disk[endblock];
            length = 0;
            while disk[endblock - length] == fileid {
                if endblock - length - 1 == 0 {
                    break;
                }
                length += 1;
            }
            if length == 0 {
                break;
            }
            length -= 1;
            block = 0;
            loop {
                if block + length >= disk.len() || block + length >= endblock {
                    endblock -= length;
                    break;
                }
                spaceblock = disk.get(block..=block + length).unwrap();
                if spaceblock.iter().all(Option::is_none) {
                    move_blocks(disk, endblock, length, block);
                    break;
                }
                block += 1;
            }
        }
        //print_disk(disk.clone(), false);
    }
}

fn main() {
    let mut layout = Vec::with_capacity(DATA.len());
    for char in DATA.chars() {
        layout.push(char_to_num(char));
    }
    let raw_disk = layout_to_disk(layout);
    println!("Layout parsed");
    let mut block_disk = raw_disk.clone();
    let mut file_disk = raw_disk.clone();
    compact_blocks(&mut block_disk);
    let mut cksum = calculate_checksum(block_disk);
    println!("Checksum for the blockcompacted disk is: {}", cksum);
    compact_files(&mut file_disk);
    cksum = calculate_checksum(file_disk);
    println!("Checksum for the filecompacted disk is: {}", cksum,);
}