use aoc_runner_derive::{aoc, aoc_generator};

type Input = ((usize, usize), Vec<Vec<u8>>, Vec<Vec<[u8; 5]>>);

const UP: u8 = 0;
const DOWN: u8 = 1;
const LEFT: u8 = 2;
const RIGHT: u8 = 4;
const EMPTY: u8 = 8;
const WALL: u8 = 16;
const PASSED: u8 = 32;

const DIR: [(isize, isize); 5] = [(0, -1), (0, 1), (-1, 0), (0, 0), (1, 0)];
const NEXT_DIR: [usize; 5] = [RIGHT as usize, LEFT as usize, UP as usize, 0, DOWN as usize];

#[aoc_generator(day6)]
fn parse(input: &str) -> Input {
    let lines = input.split("\n");
    let mut pos = (0, 0);
    let mut stones = Vec::with_capacity(input.len());

    let map: Vec<Vec<u8>> = lines
        .enumerate()
        .map(|(y, line)| {
            line.chars()
                .enumerate()
                .map(|(x, c)| match c {
                    '.' => EMPTY,
                    '#' => {
                        stones.push((x, y));
                        WALL
                    }
                    '^' => {
                        pos = (x, y);
                        EMPTY
                    }
                    _ => todo!(),
                })
                .collect()
        })
        .collect();

    let mut jumps = vec![vec![[u8::MAX; 5]; map[0].len()]; map.len()];

    for (x, y) in stones {
        for dir in NEXT_DIR {
            let mut d = 1;
            loop {
                let pos = (x as isize - d * DIR[dir].0, y as isize - d * DIR[dir].1);
                if pos.0 < 0
                    || pos.0 >= map[0].len() as isize
                    || pos.1 < 0
                    || pos.1 >= map.len() as isize
                    || map[pos.1 as usize][pos.0 as usize] == WALL
                {
                    break;
                }

                jumps[pos.1 as usize][pos.0 as usize][dir] = (d - 1) as u8;

                d += 1;
            }
        }
    }

    (pos, map, jumps)
}

#[aoc(day6, part1)]
fn part1(input: &Input) -> u64 {
    let h = input.1.len();
    let w = input.1[0].len();

    let mut map = input.1.clone();

    let mut pos = ((input.0).0 as isize, (input.0).1 as isize);
    let mut dir = UP as usize;
    let mut count = 0;

    loop {
        if map[pos.1 as usize][pos.0 as usize] == EMPTY {
            count += 1;
            map[pos.1 as usize][pos.0 as usize] = PASSED;
        }

        if pos.0 + DIR[dir].0 < 0
            || pos.0 + DIR[dir].0 >= w as isize
            || pos.1 + DIR[dir].1 < 0
            || pos.1 + DIR[dir].1 >= h as isize
        {
            break;
        }

        while map[(pos.1 + DIR[dir].1) as usize][(pos.0 + DIR[dir].0) as usize] == WALL {
            dir = NEXT_DIR[dir];
        }

        let next_pos = (pos.0 + DIR[dir].0, pos.1 + DIR[dir].1);
        pos = next_pos;
    }

    count
}

#[aoc(day6, part2)]
fn part2(input: &Input) -> u64 {
    let h = input.1.len();
    let w = input.1[0].len();

    let initial = ((input.0).0 as isize, (input.0).1 as isize);

    let mut map = input.1.clone();
    let jumps = &input.2;

    let mut pos = initial;
    let mut dir = UP as usize;
    let mut count = 0;

    loop {
        if pos.0 + DIR[dir].0 < 0
            || pos.0 + DIR[dir].0 >= w as isize
            || pos.1 + DIR[dir].1 < 0
            || pos.1 + DIR[dir].1 >= h as isize
        {
            break;
        }

        while map[(pos.1 + DIR[dir].1) as usize][(pos.0 + DIR[dir].0) as usize] == WALL {
            dir = NEXT_DIR[dir];
        }

        let next_pos = (pos.0 + DIR[dir].0, pos.1 + DIR[dir].1);

        if next_pos != initial && map[next_pos.1 as usize][next_pos.0 as usize] == EMPTY {
            map[next_pos.1 as usize][next_pos.0 as usize] = WALL;
            if check_loop(pos, next_pos, dir, &mut map, jumps) {
                count += 1;
            }
            map[next_pos.1 as usize][next_pos.0 as usize] = PASSED;
        }

        pos = next_pos;
    }

    count
}

fn check_loop(
    mut pos: (isize, isize),
    stone: (isize, isize),
    mut dir: usize,
    map: &mut [Vec<u8>],
    jumps: &[Vec<[u8; 5]>],
) -> bool {
    let h = map.len();
    let w = map[0].len();

    let mut changed = Vec::with_capacity(w * h * 4);

    loop {
        if pos.0 + DIR[dir].0 < 0
            || pos.0 + DIR[dir].0 >= w as isize
            || pos.1 + DIR[dir].1 < 0
            || pos.1 + DIR[dir].1 >= h as isize
        {
            break;
        }

        let mut turn = false;
        while map[(pos.1 + DIR[dir].1) as usize][(pos.0 + DIR[dir].0) as usize] == WALL {
            dir = NEXT_DIR[dir];
            turn = true;
        }

        let next_pos = if (pos.0 == stone.0) || (pos.1 == stone.1) {
            (pos.0 + DIR[dir].0, pos.1 + DIR[dir].1)
        } else {
            let d = jumps[pos.1 as usize][pos.0 as usize][dir] as isize;
            if d == u8::MAX as isize {
                changed.into_iter().for_each(|pos: (isize, isize)| {
                    map[pos.1 as usize][pos.0 as usize] &= !(UP | DOWN | LEFT | RIGHT);
                });

                return false;
            }
            (pos.0 + d * DIR[dir].0, pos.1 + d * DIR[dir].1)
        };

        if turn {
            if map[next_pos.1 as usize][next_pos.0 as usize] & dir as u8 != 0 {
                changed.into_iter().for_each(|pos: (isize, isize)| {
                    map[pos.1 as usize][pos.0 as usize] &= !(UP | DOWN | LEFT | RIGHT);
                });

                return true;
            } else {
                map[next_pos.1 as usize][next_pos.0 as usize] |= dir as u8;
                changed.push(next_pos);
            }
        }

        pos = next_pos;
    }

    changed.into_iter().for_each(|pos: (isize, isize)| {
        map[pos.1 as usize][pos.0 as usize] &= !(UP | DOWN | LEFT | RIGHT);
    });

    false
}

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

    #[test]
    fn part1_example() {
        assert_eq!(part1(&parse("....#.....\n.........#\n..........\n..#.......\n.......#..\n..........\n.#..^.....\n........#.\n#.........\n......#...")), 41);
        assert_eq!(part1(&parse(include_str!("../input/2024/day6.txt"))), 4977);
    }

    #[test]
    fn part2_example() {
        assert_eq!(part2(&parse("....#.....\n.........#\n..........\n..#.......\n.......#..\n..........\n.#..^.....\n........#.\n#.........\n......#...")), 6);
        assert_eq!(part2(&parse(include_str!("../input/2024/day6.txt"))), 1729);
    }
}