use std::str::FromStr;

use aoc_runner_derive::{aoc, aoc_generator};

#[derive(Debug, Clone)]
pub struct Stacks {
    stacks: Vec<Vec<char>>,
}

impl Stacks {
    pub fn apply_instruction(&mut self, ins: &Instruction) {
        for _ in 0..(ins.count) {
            let crte = self.stacks[ins.from - 1].pop().unwrap();
            self.stacks[ins.to - 1].push(crte);
        }
    }

    pub fn apply_instruction_in_order(&mut self, ins: &Instruction) {
        let from_stack = &mut self.stacks[ins.from - 1];
        let crates = from_stack.split_off(from_stack.len() - ins.count);
        self.stacks[ins.to - 1].extend_from_slice(&crates);
    }

    pub fn top_crates(&self) -> String {
        self.stacks
            .iter()
            .map(|stack| stack.last().unwrap())
            .collect()
    }
}

#[derive(Debug, Copy, Clone)]
pub struct InvalidStacks;

impl FromStr for Stacks {
    type Err = InvalidStacks;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let lines = s
            .lines()
            .take_while(|s| &s[1..2] != "1")
            .map(|s| s.chars().skip(1).step_by(4).collect())
            .collect::<Vec<Vec<_>>>();
        let mut stacks = Vec::new();
        stacks.resize(lines[0].len(), Vec::new());
        for line in &lines {
            for (i, crte) in line.iter().enumerate() {
                if crte == &' ' {
                    continue;
                }
                stacks[i].push(*crte);
            }
        }
        for stack in &mut stacks {
            stack.reverse();
        }
        Ok(Self { stacks })
    }
}

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

#[derive(Debug, Copy, Clone)]
pub struct InvalidInstruction;

impl FromStr for Instruction {
    type Err = InvalidInstruction;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let s = s.split(' ').collect::<Vec<_>>();
        if s.len() != 6 {
            return Err(InvalidInstruction);
        }
        Ok(Self {
            count: s[1].parse().map_err(|_| InvalidInstruction)?,
            from: s[3].parse().map_err(|_| InvalidInstruction)?,
            to: s[5].parse().map_err(|_| InvalidInstruction)?,
        })
    }
}

#[aoc_generator(day5)]
pub fn input_generator(input: &str) -> (Stacks, Vec<Instruction>) {
    let (start, instructions) = input.split_once(&"\n\n").unwrap();
    (
        start.parse().unwrap(),
        instructions
            .lines()
            .map(|ins| ins.parse().unwrap())
            .collect(),
    )
}

#[aoc(day5, part1)]
pub fn solve_part1(input: &(Stacks, Vec<Instruction>)) -> String {
    let mut stacks = input.0.clone();
    for ins in &input.1 {
        stacks.apply_instruction(ins);
    }
    stacks.top_crates()
}

#[aoc(day5, part2)]
pub fn solve_part2(input: &(Stacks, Vec<Instruction>)) -> String {
    let mut stacks = input.0.clone();
    for ins in &input.1 {
        stacks.apply_instruction_in_order(ins);
    }
    stacks.top_crates()
}