use aoc_runner_derive::{aoc, aoc_generator};

#[derive(Debug, PartialEq, Eq)]
struct Path {
    steps: usize,
    path: [Option<DirButton>; 6],
}

impl Path {
    const fn from(buttons: &[DirButton]) -> Option<Self> {
        let steps = buttons.len();
        let mut path = [None, None, None, None, None, None];
        let mut i = 0;
        while i < steps {
            path[i] = Some(match buttons[i] {
                DirButton::Up => DirButton::Up,
                DirButton::Down => DirButton::Down,
                DirButton::Left => DirButton::Left,
                DirButton::Right => DirButton::Right,
                DirButton::A => DirButton::A,
            });
            i += 1;
        }
        let steps = steps + 1;
        path[i] = Some(DirButton::A);
        Some(Self { steps, path })
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum DirButton {
    Up,
    Down,
    Left,
    Right,
    A,
}

const DIRBUTTON_PATHS: [[[Option<Path>; 2]; 5]; 5] = [
    // Up
    [
        // Up
        [Path::from(&[]), None],
        // Down
        [Path::from(&[DirButton::Down]), None],
        // Left
        [Path::from(&[DirButton::Down, DirButton::Left]), None],
        // Right
        [
            Path::from(&[DirButton::Down, DirButton::Right]),
            Path::from(&[DirButton::Right, DirButton::Down]),
        ],
        // A
        [Path::from(&[DirButton::Right]), None],
    ],
    // Down
    [
        // Up
        [Path::from(&[DirButton::Up]), None],
        // Down
        [Path::from(&[]), None],
        // Left
        [Path::from(&[DirButton::Left]), None],
        // Right
        [Path::from(&[DirButton::Right]), None],
        // A
        [
            Path::from(&[DirButton::Up, DirButton::Right]),
            Path::from(&[DirButton::Right, DirButton::Up]),
        ],
    ],
    // Left
    [
        // Up
        [Path::from(&[DirButton::Right, DirButton::Up]), None],
        // Down
        [Path::from(&[DirButton::Right]), None],
        // Left
        [Path::from(&[]), None],
        // Right
        [Path::from(&[DirButton::Right, DirButton::Right]), None],
        // A
        [
            Path::from(&[DirButton::Right, DirButton::Right, DirButton::Up]),
            None,
        ],
    ],
    // Right
    [
        // Up
        [
            Path::from(&[DirButton::Up, DirButton::Left]),
            Path::from(&[DirButton::Left, DirButton::Up]),
        ],
        // Down
        [Path::from(&[DirButton::Left]), None],
        // Left
        [Path::from(&[DirButton::Left, DirButton::Left]), None],
        // Right
        [Path::from(&[]), None],
        // A
        [Path::from(&[DirButton::Up]), None],
    ],
    // A
    [
        // Up
        [Path::from(&[DirButton::Left]), None],
        // Down
        [
            Path::from(&[DirButton::Down, DirButton::Left]),
            Path::from(&[DirButton::Left, DirButton::Down]),
        ],
        // Left
        [
            Path::from(&[DirButton::Down, DirButton::Left, DirButton::Left]),
            None,
        ],
        // Right
        [Path::from(&[DirButton::Down]), None],
        // A
        [Path::from(&[]), None],
    ],
];

impl DirButton {
    fn index(&self) -> usize {
        match self {
            DirButton::Up => 0,
            DirButton::Down => 1,
            DirButton::Left => 2,
            DirButton::Right => 3,
            DirButton::A => 4,
        }
    }

    fn to(&self, other: &DirButton) -> &[Option<Path>; 2] {
        &DIRBUTTON_PATHS[self.index()][other.index()]
    }
}

#[derive(Debug, PartialEq, Eq)]
enum NumButton {
    One,
    Two,
    Three,
    Four,
    Five,
    Six,
    Seven,
    Eight,
    Nine,
    Zero,
    A,
}

impl NumButton {
    fn pos(&self) -> (usize, usize) {
        match self {
            NumButton::One => (0, 2),
            NumButton::Two => (1, 2),
            NumButton::Three => (2, 2),
            NumButton::Four => (0, 1),
            NumButton::Five => (1, 1),
            NumButton::Six => (2, 1),
            NumButton::Seven => (0, 0),
            NumButton::Eight => (1, 0),
            NumButton::Nine => (2, 0),
            NumButton::Zero => (1, 3),
            NumButton::A => (2, 3),
        }
    }

    fn to(&self, other: &NumButton) -> [Option<Path>; 2] {
        let start = self.pos();
        let end = other.pos();

        let bx = if start.0 < end.0 {
            DirButton::Right
        } else {
            DirButton::Left
        };
        let dx = start.0.abs_diff(end.0);

        let by = if start.1 < end.1 {
            DirButton::Down
        } else {
            DirButton::Up
        };
        let dy = start.1.abs_diff(end.1);

        if start.0 == 0 && end.1 == 3 {
            let mut x = [None, None, None, None, None, None];
            (0..dx).for_each(|i| {
                x[i] = Some(bx);
            });
            (0..dy).for_each(|i| {
                x[i + dx] = Some(by);
            });
            x[dx + dy] = Some(DirButton::A);
            [
                Some(Path {
                    steps: dx + dy + 1,
                    path: x,
                }),
                None,
            ]
        } else if end.0 == 0 && start.1 == 3 {
            let mut y = [None, None, None, None, None, None];
            (0..dx).for_each(|i| {
                y[i + dy] = Some(bx);
            });
            (0..dy).for_each(|i| {
                y[i] = Some(by);
            });
            y[dx + dy] = Some(DirButton::A);
            [
                Some(Path {
                    steps: dx + dy + 1,
                    path: y,
                }),
                None,
            ]
        } else {
            let mut x = [None, None, None, None, None, None];
            let mut y = [None, None, None, None, None, None];
            (0..dx).for_each(|i| {
                x[i] = Some(bx);
                y[i + dy] = Some(bx);
            });
            (0..dy).for_each(|i| {
                x[i + dx] = Some(by);
                y[i] = Some(by);
            });
            x[dx + dy] = Some(DirButton::A);
            y[dx + dy] = Some(DirButton::A);
            [
                Some(Path {
                    steps: dx + dy + 1,
                    path: x,
                }),
                Some(Path {
                    steps: dx + dy + 1,
                    path: y,
                }),
            ]
        }
    }
}

type Input = Vec<(usize, Vec<NumButton>)>;

#[aoc_generator(day21)]
fn parse(input: &str) -> Input {
    input
        .lines()
        .map(|line| {
            (
                line[0..3].parse().unwrap(),
                line.chars()
                    .map(|c| match c {
                        '1' => NumButton::One,
                        '2' => NumButton::Two,
                        '3' => NumButton::Three,
                        '4' => NumButton::Four,
                        '5' => NumButton::Five,
                        '6' => NumButton::Six,
                        '7' => NumButton::Seven,
                        '8' => NumButton::Eight,
                        '9' => NumButton::Nine,
                        '0' => NumButton::Zero,
                        'A' => NumButton::A,
                        _ => unreachable!(),
                    })
                    .collect(),
            )
        })
        .collect()
}

type DistanceMatrix = [[usize; 5]; 5];

fn precompute_steps(steps: &mut [DistanceMatrix]) {
    if steps.len() == 1 {
        return;
    }

    precompute_steps(&mut steps[1..]);

    let buttons = [
        DirButton::Up,
        DirButton::Down,
        DirButton::Left,
        DirButton::Right,
        DirButton::A,
    ];

    buttons.iter().for_each(|start| {
        buttons.iter().for_each(|end| {
            steps[0][start.index()][end.index()] = start
                .to(end)
                .iter()
                .flatten()
                .map(|path| {
                    path.path
                        .iter()
                        .flatten()
                        .fold((0, &DirButton::A), |acc, next| {
                            (acc.0 + steps[1][acc.1.index()][next.index()], next)
                        })
                        .0
                })
                .fold(usize::MAX, |acc, next| acc.min(next));
        });
    });
}

fn shortest_dir_path(path: Path, steps: &mut [DistanceMatrix]) -> usize {
    path.path
        .iter()
        .flatten()
        .fold((0, &DirButton::A), |acc, next| {
            let sum = acc.0;
            let pos = acc.1;

            let extension = steps[0][pos.index()][next.index()];

            (sum + extension, next)
        })
        .0
}

fn shortest_path(path: &[NumButton], steps: &mut [DistanceMatrix]) -> usize {
    path.iter()
        .fold((0, &NumButton::A), |acc, end| {
            let sum = acc.0;
            let start = acc.1;

            let extension = start
                .to(end)
                .into_iter()
                .flatten()
                .map(|path| shortest_dir_path(path, steps))
                .min()
                .unwrap();
            (sum + extension, end)
        })
        .0
}

#[aoc(day21, part1)]
fn part1(input: &Input) -> usize {
    let mut steps = vec![[[1; 5]; 5]; 3];

    precompute_steps(&mut steps);

    input
        .iter()
        .map(|(num, path)| num * shortest_path(path, &mut steps))
        .sum()
}

#[aoc(day21, part2)]
fn part2(input: &Input) -> usize {
    let mut steps = vec![[[1; 5]; 5]; 26];

    precompute_steps(&mut steps);

    input
        .iter()
        .map(|(num, path)| num * shortest_path(path, &mut steps))
        .sum()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn part1_example() {
        assert_eq!(part1(&parse("179A")), 68 * 179);
        assert_eq!(part1(&parse("379A")), 64 * 379);
        assert_eq!(part1(&parse("456A")), 64 * 456);
        assert_eq!(part1(&parse("980A")), 60 * 980);
        assert_eq!(part1(&parse("029A")), 68 * 29);
    }
}