use aoc_runner_derive::{aoc as aoc_run, aoc_generator};
use itertools::Itertools;

type Stax = Vec<Vec<char>>;
type Inst = (usize, usize, usize);

enum Pstate {
    Setup,
    Instr,
}

#[aoc_generator(day5)]
fn parse_input_day1(input: &str) -> (Stax, Vec<Inst>) {
    let mut out = Vec::new();
    let mut inst = Vec::new();
    let mut state = Pstate::Setup;
    for _ in 0..9 {
        out.push(Vec::new());
    }

    let re = regex::Regex::new(r"^move (\d+) from (\d+) to (\d+)$").unwrap();
    for line in input.lines() {
        match state {
            Pstate::Setup => {
                if line.trim().starts_with('1') {
                    continue;
                }
                if line.trim().is_empty() {
                    state = Pstate::Instr;
                }
                for (i, cell) in line.chars().chunks(4).into_iter().enumerate() {
                    let cell: String = cell.collect();
                    let cell = cell.trim();
                    if !cell.is_empty() {
                        let cell = cell.strip_prefix('[').unwrap().strip_suffix(']').unwrap();
                        out[i].push(cell.chars().next().unwrap());
                    }
                }
            }
            // parse the instructions
            Pstate::Instr => {
                let ins = re.captures(line).unwrap();
                inst.push((
                    ins[1].parse().unwrap(),
                    ins[2].parse::<usize>().unwrap() - 1,
                    ins[3].parse::<usize>().unwrap() - 1,
                ));
            }
        }
    }
    for stack in out.iter_mut() {
        stack.reverse();
    }
    (out, inst)
}

#[aoc_run(day5, part1)]
fn part1((stax, inst): &(Stax, Vec<Inst>)) -> String {
    let mut stax = stax.to_owned();
    for ins in inst.iter() {
        let (amt, from, to) = ins;
        let from = &mut stax[*from];
        let mut tmp = Vec::with_capacity(*amt);
        for _ in 0..*amt {
            let t = from.pop().unwrap();
            tmp.push(t);
        }
        stax[*to].append(&mut tmp);
    }
    let mut s = String::with_capacity(9);
    for stack in stax {
        s.push(stack.last().copied().unwrap());
    }
    s
}

#[aoc_run(day5, part2)]
fn p2((stax, inst): &(Stax, Vec<Inst>)) -> String {
    let mut stax = stax.to_owned();
    for ins in inst.iter() {
        let (amt, from, to) = ins;
        let from = &mut stax[*from];
        let mut tmp = Vec::with_capacity(*amt);
        for _ in 0..*amt {
            let c = from.pop().unwrap();
            tmp.push(c);
        }
        for c in tmp.into_iter().rev() {
            stax[*to].push(c);
        }
    }
    let mut s = String::with_capacity(9);
    for stack in stax {
        s.push(stack.last().copied().unwrap());
    }
    s
}