type Block = u32;
type FSunpacked = Vec<Option<Block>>;

fn parse_input(input: &str) -> FSunpacked {
    let mut result = vec![];
    let mut is_empty = false;
    let mut blockid = 0;
    for c in input.chars() {
        let Some(nblocks) = c.to_digit(10) else {
            continue;
        };
        if is_empty {
            for _i in 0..nblocks {
                result.push(None);
            }
        } else {
            for _i in 0..nblocks {
                result.push(Some(blockid));
            }
            blockid += 1;
        }
        is_empty = !is_empty;
    }
    result
}

fn compact(v: &mut FSunpacked) {
    let mut i_empty = 0;
    let mut i_full = v.len() - 1;
    'out: loop {
        while v[i_empty].is_some() {
            i_empty += 1;
            if i_full < i_empty {
                break 'out;
            }
        }
        while v[i_full].is_none() {
            i_full -= 1;
            if i_full < i_empty {
                break 'out;
            }
        }
        (v[i_empty], v[i_full]) = (v[i_full], v[i_empty]);
    }
}

fn checksum(v: FSunpacked) -> u64 {
    v.into_iter()
        .enumerate()
        .map(|(i, b)| i as u64 * b.unwrap_or_default() as u64)
        .sum()
}

fn p1(input: &str) -> String {
    let mut unpacked = parse_input(input);
    // println!("{unpacked:?}");
    compact(&mut unpacked);
    // println!("{unpacked:?}");

    let result = checksum(unpacked);
    result.to_string()
}

type FSpacked = Vec<(Option<Block>, u32)>;

fn parse_input2(input: &str) -> FSpacked {
    let mut result = vec![];
    let mut is_empty = false;
    let mut blockid = 0;
    for c in input.chars() {
        // println!("Parsing {c}...");
        let Some(nblocks) = c.to_digit(10) else {
            continue;
        };
        if is_empty {
            result.push((None, nblocks));
        } else {
            result.push((Some(blockid), nblocks));
            blockid += 1;
        }
        is_empty = !is_empty;
    }
    result
}

fn compact2(v: &mut FSpacked) {
    let mut i_empty = 0;
    let mut i_full = v.len() - 1;
    loop {
        while v[i_empty].0.is_some() {
            i_empty += 1;
        }
        while v[i_full].0.is_none() {
            i_full -= 1;
        }
        let fblocks = v[i_full].1;
        let mut j_empty = i_empty;
        while j_empty <= i_full {
            // println!("Can we move {:?} into {:?}?", v[i_full], v[j_empty]);
            match v[j_empty] {
                (None, eblocks) if eblocks >= fblocks => {
                    let diff = eblocks - fblocks;
                    v[j_empty] = v[i_full];
                    v[i_full] = (None, fblocks);
                    if diff != 0 {
                        i_full += 1;
                        j_empty += 1;
                        v.insert(j_empty, (None, diff));
                    }
                    break;
                }
                _ => (),
            }
            j_empty += 1;
        }
        if i_full == 0 {
            break;
        } else {
            i_full -= 1;
        }
    }
}

fn unpack(v: &FSpacked) -> FSunpacked {
    v.iter()
        .flat_map(|(x, count)| (0..*count).map(|_| *x))
        .collect()
}

fn p2(input: &str) -> String {
    let mut packed = parse_input2(input);
    compact2(&mut packed);
    let unpacked = unpack(&packed);
    // println!("{unpacked:?}");

    let result = checksum(unpacked);
    result.to_string()
}

fn main() {
    aoc2024::run_day("9", p1, p2);
}