use std::str::FromStr;

use aoc_runner_derive::{aoc, aoc_generator};
use yap::{IntoTokens, Tokens};

use crate::parsing::parse_number;

#[aoc_generator(day2)]
pub fn parse_movements(input: &str) -> Vec<Movement> {
    input
        .split('\n')
        .map(|input| input.parse().unwrap())
        .collect()
}

#[aoc(day2, part1)]
pub fn part1(input: &[Movement]) -> isize {
    NaiveState::default().movements(input)
}

#[aoc(day2, part2)]
pub fn part2(input: &[Movement]) -> isize {
    CorrectState::default().movements(input)
}

pub trait State {
    fn movement(&mut self, movement: &Movement);
    fn depth(&self) -> isize;
    fn horizontal_pos(&self) -> isize;
    fn result(&self) -> isize {
        self.depth() * self.horizontal_pos()
    }
    fn movements(&mut self, moves: &[Movement]) -> isize {
        for movement in moves {
            self.movement(movement);
        }
        self.result()
    }
}

#[derive(Default)]
pub struct CorrectState {
    horizontal_pos: isize,
    depth: isize,
    aim: isize,
}

impl State for CorrectState {
    fn movement(&mut self, movement: &Movement) {
        match movement {
            Movement::Forward(units) => {
                self.horizontal_pos += units;
                self.depth += units * self.aim;
            }
            Movement::Down(units) => self.aim += units,
            Movement::Up(units) => self.aim -= units,
        }
    }

    fn depth(&self) -> isize {
        self.depth
    }

    fn horizontal_pos(&self) -> isize {
        self.horizontal_pos
    }
}

#[derive(Default)]
pub struct NaiveState {
    horizontal_pos: isize,
    depth: isize,
}

impl State for NaiveState {
    fn movement(&mut self, movement: &Movement) {
        match movement {
            Movement::Forward(units) => self.horizontal_pos += units,
            Movement::Down(units) => self.depth += units,
            Movement::Up(units) => self.depth -= units,
        }
    }

    fn depth(&self) -> isize {
        self.depth
    }

    fn horizontal_pos(&self) -> isize {
        self.horizontal_pos
    }
}

pub enum Movement {
    Forward(isize),
    Down(isize),
    Up(isize),
}

#[allow(clippy::unnecessary_mut_passed)]
impl FromStr for Movement {
    type Err = ();
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let mut tokens = s.into_tokens();
        yap::one_of!(ts from &mut tokens;
        ts.tokens("forward ".chars()).then(|| parse_number(ts).map(Movement::Forward)).flatten(),
        ts.tokens("down ".chars()).then(|| parse_number(ts).map(Movement::Down)).flatten(),
        ts.tokens("up ".chars()).then(|| parse_number(ts).map(Movement::Up)).flatten(),
        )
        .ok_or(())
    }
}

#[cfg(test)]
mod tests {
    const EXAMPLE_INPUT: &str = "forward 5
down 5
forward 8
up 3
down 8
forward 2";

    #[test]
    fn part1_example() {
        let input = super::parse_movements(EXAMPLE_INPUT);
        let result = super::part1(&input);
        assert_eq!(result, 150);
    }

    #[test]
    fn part2_example() {
        let input = super::parse_movements(EXAMPLE_INPUT);
        let result = super::part2(&input);
        assert_eq!(result, 900);
    }
}