use std::{
    cmp::Ordering,
    collections::HashSet,
    ops::{Add, Sub},
};

use aoc_runner_derive::{aoc, aoc_generator};

#[derive(Clone, Copy)]
enum Direction {
    Up,
    Down,
    Left,
    Right,
    UpLeft,
    UpRight,
    DownLeft,
    DownRight,
}

impl Direction {
    fn to_vec2(self) -> Vec2 {
        match self {
            Direction::Up => (0, 1),
            Direction::Down => (0, -1),
            Direction::Left => (-1, 0),
            Direction::Right => (1, 0),
            Direction::UpLeft => (-1, 1),
            Direction::UpRight => (1, 1),
            Direction::DownLeft => (-1, -1),
            Direction::DownRight => (1, -1),
        }
        .into()
    }
}

#[derive(Clone, Copy, PartialEq, Eq, Hash)]
struct Vec2 {
    x: i32,
    y: i32,
}

impl From<(i32, i32)> for Vec2 {
    fn from((x, y): (i32, i32)) -> Self {
        Self { x, y }
    }
}

impl Add for Vec2 {
    type Output = Self;

    fn add(self, rhs: Self) -> Self::Output {
        Self {
            x: self.x + rhs.x,
            y: self.y + rhs.y,
        }
    }
}

impl Sub for Vec2 {
    type Output = Self;

    fn sub(self, rhs: Self) -> Self::Output {
        Self {
            x: self.x - rhs.x,
            y: self.y - rhs.y,
        }
    }
}

#[aoc_generator(day9)]
fn generator(input: &str) -> Vec<(Direction, u32)> {
    input
        .lines()
        .map(|l| {
            (
                match l.chars().next().unwrap() {
                    'U' => Direction::Up,
                    'D' => Direction::Down,
                    'L' => Direction::Left,
                    'R' => Direction::Right,
                    _ => panic!(),
                },
                l[2..].parse().unwrap(),
            )
        })
        .collect()
}

fn do_pull(head: Vec2, tail: Vec2) -> Vec2 {
    let head_relative_to_tail = head - tail;

    if head_relative_to_tail.x.abs() <= 1 && head_relative_to_tail.y.abs() <= 1 {
        return tail;
    }

    use Ordering::*;

    let dir_to_move_tail = match (
        head_relative_to_tail.x.cmp(&0),
        head_relative_to_tail.y.cmp(&0),
    ) {
        (Less, Less) => Direction::DownLeft,
        (Less, Equal) => Direction::Left,
        (Less, Greater) => Direction::UpLeft,
        (Equal, Less) => Direction::Down,
        (Equal, Equal) => unreachable!(),
        (Equal, Greater) => Direction::Up,
        (Greater, Less) => Direction::DownRight,
        (Greater, Equal) => Direction::Right,
        (Greater, Greater) => Direction::UpRight,
    };

    tail + dir_to_move_tail.to_vec2()
}

struct Rope<const N: usize> {
    inner: [Vec2; N],
}

impl<const N: usize> Rope<N> {
    fn new() -> Self {
        Self {
            inner: [(0, 0).into(); N],
        }
    }

    fn do_move(&mut self, dir: Direction) {
        self.inner[0] = self.inner[0] + dir.to_vec2();
        for i in 0..N - 1 {
            self.inner[i + 1] = do_pull(self.inner[i], self.inner[i + 1]);
        }
    }

    fn tail(&self) -> Vec2 {
        *self.inner.last().unwrap()
    }
}

fn run_simulation<const N: usize>(input: &[(Direction, u32)]) -> usize {
    let mut rope = Rope::<N>::new();
    let mut tail_locations = HashSet::new();
    tail_locations.insert(rope.tail());

    for instruction in input {
        for _ in 0..instruction.1 {
            rope.do_move(instruction.0);
            tail_locations.insert(rope.tail());
        }
    }

    tail_locations.len()
}

#[aoc(day9, part1)]
fn part1_indexing(input: &[(Direction, u32)]) -> usize {
    run_simulation::<2>(input)
}

#[aoc(day9, part2)]
fn part2(input: &[(Direction, u32)]) -> usize {
    run_simulation::<10>(input)
}