use aoc_runner_derive::{aoc, aoc_generator};
use std::{error::Error, str::FromStr};

#[derive(Debug, Copy, Clone)]
enum Direction {
    Forward,
    Down,
    Up,
}
#[derive(Debug, Copy, Clone)]
pub struct Instruction {
    direction: Direction,
    amount: usize,
}

#[derive(Debug, Copy, Clone, Default)]
pub struct Submarine {
    depth: i64,
    pos: i64,
    aim: i64,
}

impl Submarine {
    fn new() -> Self {
        Default::default()
    }
    fn run_p1(&mut self, inst: &Instruction) {
        let d = inst.direction;
        let a = inst.amount as i64;
        match d {
            Direction::Forward => {
                self.pos += a;
            }
            Direction::Down => {
                self.depth += a;
            }
            Direction::Up => {
                self.depth -= a;
            }
        };
    }
    fn run_p2(&mut self, inst: &Instruction) {
        let d = inst.direction;
        let a = inst.amount as i64;
        match d {
            Direction::Forward => {
                self.pos += a;
                self.depth += self.aim * a;
            }
            Direction::Down => {
                self.aim += a;
            }
            Direction::Up => {
                self.aim -= a;
            }
        };
    }
}

impl FromStr for Instruction {
    type Err = Box<dyn Error>;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let mut s = s.split_ascii_whitespace();
        let inst = s.next().unwrap();
        let amount: usize = s.next().unwrap().parse().unwrap();
        let ret = match inst {
            "forward" => Self {
                direction: Direction::Forward,
                amount,
            },
            "up" => Self {
                direction: Direction::Up,
                amount,
            },
            "down" => Self {
                direction: Direction::Down,
                amount,
            },
            other => {
                panic!("Invalid instruction: {}", other)
            }
        };
        Ok(ret)
    }
}

pub type Data = Vec<Instruction>;

#[aoc_generator(day02)]
pub fn input_generator(input: &str) -> Data {
    input
        .lines()
        .map(Instruction::from_str)
        .collect::<Result<Vec<_>, _>>()
        .unwrap()
}

#[aoc(day02, part1)]
pub fn solve_part1(input: &Data) -> usize {
    let mut sm = Submarine::new();
    for inst in input {
        sm.run_p1(inst);
    }
    (sm.pos * sm.depth) as usize
}

#[aoc(day02, part2)]
pub fn solve_part2(input: &Data) -> usize {
    let mut sm = Submarine::new();
    for inst in input {
        sm.run_p2(inst);
    }
    (sm.pos * sm.depth) as usize
}