use {
    anyhow::{
        anyhow,
        Context,
        Result,
    },
    regex::Regex,
    std::{
        collections::HashMap,
        env::args,
        path::Path,
    },
};

type Stack = Vec<char>;
type Stacks = HashMap<usize, Stack>;

#[derive(Debug)]
struct Rearrangment {
    count: usize,
    from: usize,
    to: usize,
}

impl Rearrangment {
    fn new(count: usize, from: usize, to: usize) -> Self {
        Self {
            count,
            from,
            to,
        }
    }
}

fn read_input<T: AsRef<Path>>(filename: T) -> Result<(Stacks, Vec<Rearrangment>)> {
    let re = Regex::new(r"move (?P<count>[0-9]+) from (?P<from>[0-9]+) to (?P<to>[0-9]+)")?;

    let content = std::fs::read_to_string(filename)?;
    let (picture, lines) = content.split_once("\n\n").ok_or_else(|| anyhow!("Invalid input"))?;

    let mut stacks = HashMap::new();
    for layer in picture.split('\n') {
        for (num, chunk) in layer.as_bytes().chunks(4).enumerate() {
            let byte = chunk.get(1).ok_or_else(|| anyhow!("Could not parse picture."))?;
            let ch = *byte as char;
            if ('A'..='Z').contains(&ch) {
                stacks.entry(num + 1).and_modify(|e: &mut Vec<char>| e.insert(0, ch))
                    .or_insert_with(|| vec![ch]);
            }
        }
    }

    let rearrangements = lines.split('\n').filter(|line| !line.is_empty()).map(|line| {
        let caps = re.captures(line).ok_or_else(|| anyhow!("Regex matching failed."))?;

        match (caps.name("count"), caps.name("from"), caps.name("to")) {
            (Some(count), Some(from), Some(to)) => Ok(Rearrangment::new(count.as_str().parse()?,
                from.as_str().parse()?, to.as_str().parse()?)),
            _ => Err(anyhow!("Failed to parse rearrangements.")),
        }
    }).collect::<Result<Vec<_>>>()?;

    Ok((stacks, rearrangements))
}

fn part1(input_stacks: &Stacks, rearrangements: &[Rearrangment]) -> Result<String> {
    let mut stacks = input_stacks.clone();
    for rearrangement in rearrangements {
        let count = rearrangement.count;
        let from = rearrangement.from;
        let to = rearrangement.to;
        for _ in 0..count {
            let source = stacks.get_mut(&from)
                .ok_or_else(|| anyhow!("Missing stack {from}"))?;
            if let Some(elem) = source.pop() {
                let destination = stacks.get_mut(&to)
                    .ok_or_else(|| anyhow!("Missing stack {to}"))?;
                destination.push(elem);
            }
        }
    }

    let mut keys: Vec<_> = stacks.keys().cloned().collect();
    keys.sort();

    Ok(keys.iter().filter_map(|key| {
        let stack = stacks.get_mut(key).unwrap_or_else(|| unreachable!());
        stack.pop()
    }).collect())
}

fn part2(input_stacks: &Stacks, rearrangements: &[Rearrangment]) -> Result<String> {
    let mut stacks = input_stacks.clone();
    for rearrangement in rearrangements {
        let count = rearrangement.count;
        let from = rearrangement.from;
        let to = rearrangement.to;
        let source = stacks.get_mut(&from)
            .ok_or_else(|| anyhow!("Missing stack {from}"))?;

        let mut elems = source.drain(source.len()-count..).collect();
        let destination = stacks.get_mut(&to)
            .ok_or_else(|| anyhow!("Missing stack {to}"))?;
        destination.append(&mut elems);
    }

    let mut keys: Vec<_> = stacks.keys().cloned().collect();
    keys.sort();

    Ok(keys.iter().filter_map(|key| {
        let stack = stacks.get_mut(key).unwrap_or_else(|| unreachable!());
        stack.pop()
    }).collect())
}

fn main() -> Result<()> {
    let ( do_part_1, do_part_2 ) = aoc::do_parts();

    let filename = args().nth(1).ok_or_else(|| anyhow!("Missing input filename"))?;
    let (stacks, rearrangements) = read_input(filename).context("Could not read input")?;
    if do_part_1 {
        let solution = part1(&stacks, &rearrangements).context("No solution for part 1")?;
        println!("Part1, solution found to be: {}", solution);
    }
    if do_part_2 {
        let solution = part2(&stacks, &rearrangements).context("No solution for part 2")?;
        println!("Part2, solution found to be: {}", solution);
    }
    Ok(())
}