use std::{collections::VecDeque, time::Instant};

use aoc_2024::{load, print_res, tinyvec::TinyVec};
use bstr::{BString, ByteSlice};
use color_eyre::eyre::ensure;

type Parsed = VecDeque<u8>;

#[inline(never)]
pub fn parsing(input: &BString) -> color_eyre::Result<Parsed> {
    let parsed: Parsed = input
        .trim()
        .iter()
        .map(|&n| {
            ensure!(n.is_ascii_digit());
            Ok(n - b'0')
        })
        .collect::<Result<_, _>>()?;
    // The last entry is a filled block
    ensure!(parsed.len() % 2 == 1);

    Ok(parsed)
}

#[inline(never)]
pub fn part1(mut input: Parsed) {
    let mut checksum = 0;

    let mut front_id = 0;
    let mut back_id = input.len() / 2;
    let mut current_back: Option<(usize, u8)> = None;

    let mut cursor = 0;

    while !input.is_empty() {
        let front_len = input.pop_front().unwrap();

        for _ in 0..front_len {
            checksum += cursor * front_id;
            cursor += 1;
        }
        front_id += 1;

        let mut front_hole = match input.pop_front() {
            None => break,
            Some(h) => h,
        };
        while front_hole != 0 {
            let (id, mut amount) = match current_back {
                None if input.is_empty() => break,
                Some(_) if input.is_empty() => todo!(),
                Some((id, amount)) => (id, amount),
                None => {
                    let back_len = input.pop_back().unwrap();
                    input.pop_back(); // Back hole (if any)

                    let fetched = (back_id, back_len);
                    back_id -= 1;
                    fetched
                }
            };

            let moved = std::cmp::min(front_hole, amount);
            front_hole -= moved;
            amount -= moved;

            for _ in 0..moved {
                checksum += cursor * id;
                cursor += 1;
            }

            if amount != 0 {
                current_back = Some((id, amount));
            } else {
                current_back = None;
            }
        }
    }

    if let Some((id, amount)) = current_back {
        for _ in 0..amount {
            checksum += cursor * id;
            cursor += 1;
        }
    }

    print_res!("Checksum: {checksum}")
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
struct File {
    width: u8,
    id: [u8; 2],
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum Entry {
    Hole { width: u8 },
    File(File),
}

impl Entry {
    fn file(self) -> File {
        match self {
            Entry::Hole { .. } => panic!("Entry is a hole"),
            Entry::File(file) => file,
        }
    }

    fn width(self) -> u8 {
        match self {
            Entry::Hole { width } => width,
            Entry::File(file) => file.width,
        }
    }
}

#[derive(Debug)]
struct FileSystem {
    blocks: Vec<TinyVec<Entry>>,
}

impl FileSystem {
    fn iter(&self) -> impl DoubleEndedIterator<Item = &Entry> {
        self.blocks.iter().flat_map(|v| v.iter())
    }

    fn iter_mut(&mut self) -> impl DoubleEndedIterator<Item = &mut Entry> {
        self.blocks.iter_mut().flat_map(|v| v.iter_mut())
    }

    fn file_mut(&mut self, id: u16) -> Option<&mut Entry> {
        self.iter_mut().rev().find(
            |x| matches!(x, Entry::File(File { id: block_id, .. }) if id.to_le_bytes() == *block_id),
        )
    }

    fn find_hole(&mut self, src: File) -> Option<(usize, &mut TinyVec<Entry>)> {
        for block in &mut self.blocks {
            match block
                .iter()
                .enumerate()
                .take_while(|(_, x)| **x != Entry::File(src))
                .find(|(_, x)| matches!(**x, Entry::Hole { width } if width >= src.width))
            {
                None => continue,
                Some((idx, _)) => return Some((idx, block)),
            }
        }

        None
    }

    fn checksum(&self) -> usize {
        let mut idx = 0;
        let mut checksum = 0;

        for &entry in self.iter() {
            let (id, width) = match entry {
                Entry::Hole { width } => (None, width),
                Entry::File(file) => (Some(u16::from_le_bytes(file.id)), file.width),
            };

            for _ in 0..width {
                if let Some(id) = id {
                    checksum += idx * (id as usize);
                }
                idx += 1;
            }
        }

        checksum
    }
}

impl std::fmt::Display for FileSystem {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        for &entry in self.iter() {
            let (char, width) = match entry {
                Entry::Hole { width } => ('.', width),
                Entry::File(File { width, id }) => (
                    if id[1] == 0 && id[0] <= 9 {
                        char::from_u32('0' as u32 + id[0] as u32).unwrap()
                    } else {
                        '*'
                    },
                    width,
                ),
            };
            for _ in 0..width {
                write!(f, "{}", char)?;
            }
        }
        writeln!(f)
    }
}

#[inline(never)]
pub fn part2(input: Parsed) {
    let mut id = 0u16;
    let mut fs = FileSystem {
        blocks: input
            .iter()
            .enumerate()
            .map(|(idx, &n)| {
                let mut v = TinyVec::new();
                match idx % 2 == 0 {
                    true => {
                        v.push(Entry::File(File {
                            width: n,
                            id: id.to_le_bytes(),
                        }));
                        id += 1;
                    }
                    false => v.push(Entry::Hole { width: n }),
                };
                v
            })
            .collect(),
    };

    for id in (0..id).rev() {
        let file = fs.file_mut(id).unwrap().file();

        let Some((hole, block)) = fs.find_hole(file) else {
            continue;
        };

        let remaining = block[hole].width() - file.width;

        block[hole] = Entry::File(file);
        if remaining > 0 {
            if hole == block.len() - 1 {
                block.push(Entry::Hole { width: remaining });
            } else {
                todo!()
            }
        }

        *fs.file_mut(id).unwrap() = Entry::Hole { width: file.width }
    }

    print_res!("Total checksum: {}", fs.checksum());
}

pub fn main() -> color_eyre::Result<()> {
    let context = load()?;

    let start = Instant::now();
    let parsed = parsing(&context.input)?;
    let elapsed = humantime::format_duration(start.elapsed());

    let start = Instant::now();
    if context.part == 1 {
        part1(parsed);
    } else {
        part2(parsed);
    }
    let elapsed_part = humantime::format_duration(start.elapsed());

    println!("  Parsing: {elapsed}");
    println!("  Solving: {elapsed_part}");

    Ok(())
}