use std::iter::once;

use aoc_runner_derive::{aoc, aoc_generator};
use hashbrown::{hash_map::Entry, HashMap};

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

impl Path {
    fn new() -> Self {
        Self {
            steps: 0,
            path: [None, None, None, None, None, None],
        }
    }

    fn push(&mut self, button: Button) {
        self.path[self.steps] = Some(button);
        self.steps += 1;
    }

    fn append(&mut self, other: &Self) {
        for i in 0..other.steps {
            self.push(other.path[i].unwrap());
        }
    }
}

impl IntoIterator for Path {
    type Item = Button;
    type IntoIter = std::iter::Flatten<std::array::IntoIter<Option<Button>, 6>>;

    fn into_iter(self) -> Self::IntoIter {
        self.path.into_iter().flatten()
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum Button {
    Directional(DirectionalButton),
    Numeric(NumericButton),
}

impl Button {
    fn pos(&self) -> (isize, isize) {
        match self {
            Button::Directional(button) => button.pos(),
            Button::Numeric(button) => button.pos(),
        }
    }

    fn to(&self, other: &Button) -> (Path, Option<Path>) {
        match (self, other) {
            (Button::Directional(a), Button::Directional(b)) => a.to(b),
            (Button::Numeric(a), Button::Numeric(b)) => a.to(b),
            _ => unimplemented!(),
        }
    }
}

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

impl DirectionalButton {
    fn pos(&self) -> (isize, isize) {
        match self {
            DirectionalButton::Up => (1, 0),
            DirectionalButton::Down => (1, 1),
            DirectionalButton::Left => (0, 1),
            DirectionalButton::Right => (2, 1),
            DirectionalButton::A => (2, 0),
        }
    }

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

        let mut x = Path::new();
        let mut y = Path::new();

        if start.0 < end.0 {
            (start.0..end.0).for_each(|_| x.push(Button::Directional(DirectionalButton::Right)));
        } else {
            (end.0..start.0).for_each(|_| x.push(Button::Directional(DirectionalButton::Left)));
        };
        if start.1 < end.1 {
            (start.1..end.1).for_each(|_| y.push(Button::Directional(DirectionalButton::Down)));
        } else {
            (end.1..start.1).for_each(|_| y.push(Button::Directional(DirectionalButton::Up)));
        };

        if start.0 == 0 && end.1 == 0 {
            x.append(&y);
            x.push(Button::Directional(DirectionalButton::A));
            (x, None)
        } else if end.0 == 0 && start.1 == 0 {
            y.append(&x);
            y.push(Button::Directional(DirectionalButton::A));
            (y, None)
        } else {
            x.append(&y);
            x.push(Button::Directional(DirectionalButton::A));
            y.append(&x);
            y.push(Button::Directional(DirectionalButton::A));
            (x, Some(y))
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum NumericButton {
    Number(u8),
    A,
}

impl NumericButton {
    fn pos(&self) -> (isize, isize) {
        match self {
            NumericButton::Number(1) => (0, 2),
            NumericButton::Number(2) => (1, 2),
            NumericButton::Number(3) => (2, 2),
            NumericButton::Number(4) => (0, 1),
            NumericButton::Number(5) => (1, 1),
            NumericButton::Number(6) => (2, 1),
            NumericButton::Number(7) => (0, 0),
            NumericButton::Number(8) => (1, 0),
            NumericButton::Number(9) => (2, 0),
            NumericButton::Number(0) => (1, 3),
            NumericButton::A => (2, 3),
            _ => unreachable!(),
        }
    }

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

        let mut x = Path::new();
        let mut y = Path::new();

        if start.0 < end.0 {
            (start.0..end.0).for_each(|_| x.push(Button::Directional(DirectionalButton::Right)));
        } else {
            (end.0..start.0).for_each(|_| x.push(Button::Directional(DirectionalButton::Left)));
        };
        if start.1 < end.1 {
            (start.1..end.1).for_each(|_| y.push(Button::Directional(DirectionalButton::Down)));
        } else {
            (end.1..start.1).for_each(|_| y.push(Button::Directional(DirectionalButton::Up)));
        };

        if start.0 == 0 && end.1 == 3 {
            x.append(&y);
            x.push(Button::Directional(DirectionalButton::A));
            (x, None)
        } else if end.0 == 0 && start.1 == 3 {
            y.append(&x);
            y.push(Button::Directional(DirectionalButton::A));
            (y, None)
        } else {
            let x_copy = x;
            x.append(&y);
            x.push(Button::Directional(DirectionalButton::A));
            y.append(&x_copy);
            y.push(Button::Directional(DirectionalButton::A));
            (x, Some(y))
        }
    }
}

type Input = Vec<Vec<Button>>;

#[aoc_generator(day21)]
fn parse(input: &str) -> Input {
    input
        .lines()
        .map(|line| {
            line.chars()
                .map(|c| match c {
                    '0' => NumericButton::Number(0),
                    '1' => NumericButton::Number(1),
                    '2' => NumericButton::Number(2),
                    '3' => NumericButton::Number(3),
                    '4' => NumericButton::Number(4),
                    '5' => NumericButton::Number(5),
                    '6' => NumericButton::Number(6),
                    '7' => NumericButton::Number(7),
                    '8' => NumericButton::Number(8),
                    '9' => NumericButton::Number(9),
                    'A' => NumericButton::A,
                    _ => unreachable!(),
                })
                .map(Button::Numeric)
                .collect()
        })
        .collect()
}

fn shortest_path(num_robots: usize, path: Path) -> usize {
    type Memos = Vec<HashMap<Path, usize>>;

    let mut memos = vec![HashMap::new(); num_robots];

    fn expand(robots: usize, path: Path, memos: &mut Memos) -> usize {
        if robots == 0 {
            return path.steps;
        }

        let pos = match path.path[0].unwrap() {
            Button::Numeric(_) => Button::Numeric(NumericButton::A),
            Button::Directional(_) => Button::Directional(DirectionalButton::A),
        };

        path.into_iter()
            .fold((pos, 0), |acc, next| {
                let pos = acc.0;
                let steps = acc.1;

                let (first, second) = pos.to(&next);
                let expansion = once(Some(first))
                    .chain(once(second))
                    .flatten()
                    .min_by_key(|path| {
                        if memos[robots - 1].contains_key(path) {
                            memos[robots - 1][path]
                        } else {
                            let expanded = expand(robots - 1, *path, memos);
                            memos[robots - 1].insert(*path, expanded);
                            expanded
                        }
                    })
                    .map(|path| expand(robots - 1, path, memos))
                    .unwrap();

                (next, steps + expansion)
            })
            .1
    }

    expand(num_robots, path, &mut memos)
}

#[aoc(day21, part1)]
fn part1(input: &Input) -> usize {
    input
        .iter()
        .cloned()
        .map(|instruction| {
            let numeric = instruction
                .iter()
                .filter_map(|button| match button {
                    Button::Numeric(NumericButton::Number(numeric)) => Some(numeric),
                    _ => None,
                })
                .fold(String::new(), |mut acc, num| {
                    acc.push_str(&num.to_string());
                    acc
                })
                .parse::<usize>()
                .unwrap();

            let mut path = Path::new();
            instruction.iter().for_each(|step| path.push(*step));

            let path = shortest_path(3, path);

            numeric * path
        })
        .sum()
}

#[aoc(day21, part2)]
fn part2(input: &Input) -> usize {
    input
        .iter()
        .cloned()
        .map(|instruction| {
            let numeric = instruction
                .iter()
                .filter_map(|button| match button {
                    Button::Numeric(NumericButton::Number(numeric)) => Some(numeric),
                    _ => None,
                })
                .fold(String::new(), |mut acc, num| {
                    acc.push_str(&num.to_string());
                    acc
                })
                .parse::<usize>()
                .unwrap();

            let mut path = Path::new();
            instruction.iter().for_each(|step| path.push(*step));

            let path = shortest_path(8, path);

            numeric * path
        })
        .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);
    }

    // #[test]
    // fn part2_example() {
    //     assert_eq!(part2(&parse("<EXAMPLE>")), 0);
    // }
}

// 89666 too low
// 97150 wrong