use std::collections::HashSet;

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

type Coord = (i32, i32);

#[aoc_generator(day9)]
fn parse_input(input: &str) -> Vec<Coord> {
    let mut out = Vec::new();
    let (mut x, mut y) = (0, 0);
    out.push((x, y));
    for line in input.lines() {
        let (dir, num) = line.split_once(' ').unwrap();
        let num: i32 = num.parse().unwrap();
        for _ in 0..num {
            match dir {
                "U" => y += 1,
                "D" => y -= 1,
                "L" => x -= 1,
                "R" => x += 1,
                _ => unreachable!(),
            }
            out.push((x, y));
        }
    }
    out
}

#[aoc_run(day9, part1)]
fn part1(input: &[Coord]) -> usize {
    let mut tail = (0, 0);
    let mut set = HashSet::new();
    set.insert(tail);
    for mv in input {
        tail = new_tail(mv, &tail);
        set.insert(tail);
    }
    set.len()
}

#[aoc_run(day9, part2)]
fn part2(input: &[Coord]) -> usize {
    let mut knots = Vec::from_iter((0..9).map(|_| (0, 0)));
    let mut set = HashSet::new();
    set.insert((0, 0));
    for mv in input {
        let mut head = *mv;
        for knot in knots.iter_mut() {
            *knot = new_tail(&head, knot);
            head = *knot;
        }
        set.insert(*knots.last().unwrap());
    }
    set.len()
}

fn new_tail(head: &Coord, tail: &Coord) -> Coord {
    let dy = head.1 - tail.1;
    let dx = head.0 - tail.0;
    match (dx, dy) {
        (m, n) if m.abs() < 2 && n.abs() < 2 => *tail,
        _ => (tail.0 + dx.signum(), tail.1 + dy.signum()),
    }
}