use std::collections::HashMap;

use aoc_2022::*;

type Input = Vec<char>;
type Output = usize;

type Position = Vector2D<isize>;

struct InfiniteIndex {
    size: usize,
    i: usize,
}

impl InfiniteIndex {
    fn new(size: usize) -> InfiniteIndex {
        InfiniteIndex { size, i: size - 1 }
    }
}

impl Iterator for InfiniteIndex {
    type Item = usize;

    fn next(&mut self) -> Option<Self::Item> {
        self.i = (self.i + 1) % self.size;
        Some(self.i)
    }
}

struct Tunnel {
    buffer: usize,
    jets: Vec<char>,
    rocks: Vec<Vec<char>>,
    flushed: usize,
    next_jet: InfiniteIndex,
    next_shape: InfiniteIndex,
}

lazy_static! {
    static ref SHAPES: Vec<Vec<Vec<char>>> = vec![
        vec!["####"],
        vec![" # ", "###", " # "],
        vec!["  #", "  #", "###"],
        vec!["#", "#", "#", "#"],
        vec!["##", "##"],
    ]
    .iter()
    .map(|shape| shape.iter().map(|l| l.chars().collect()).collect())
    .collect_vec();
}

fn below(position: Vector2D<isize>) -> Vector2D<isize> {
    position + Vector2D::new(0, 1)
}

fn left(position: Vector2D<isize>) -> Vector2D<isize> {
    position + Vector2D::new(-1, 0)
}

fn right(position: Vector2D<isize>) -> Vector2D<isize> {
    position + Vector2D::new(1, 0)
}

fn occupied(shape: &[Vec<char>]) -> impl Iterator<Item = Position> + '_ {
    shape.iter().enumerate().flat_map(|(y, row)| {
        row.iter().enumerate().filter_map(move |(x, c)| {
            if c == &'#' {
                Some(Vector2D::new(x as isize, y as isize))
            } else {
                None
            }
        })
    })
}

impl Tunnel {
    fn new(buffer: usize, jets: Vec<char>) -> Tunnel {
        let jet_count = jets.len();

        Tunnel {
            buffer,
            jets,
            rocks: vec!["#########".chars().collect_vec()],
            flushed: 0,
            next_jet: InfiniteIndex::new(jet_count),
            next_shape: InfiniteIndex::new(SHAPES.len()),
        }
    }

    fn hit(&self, shape: &[Vec<char>], position: Position) -> bool {
        occupied(shape).any(|rock| self.rocks[position + rock] != ' ')
    }

    fn emplace(&mut self, shape: &[Vec<char>], position: Position) {
        for rock in occupied(shape) {
            self.rocks[position + rock] = '#';
        }

        while !self.rocks.first().unwrap().contains(&'#') {
            self.rocks.remove(0);
        }

        while self.rocks.len() > self.buffer {
            self.rocks.pop();
            self.flushed += 1;
        }
    }

    fn add_one(&mut self) {
        let shape = &SHAPES[self.next_shape.next().unwrap()];

        for _ in 0..shape.len() + 3 {
            self.rocks.insert(0, "|       |".chars().collect_vec());
        }

        let mut position = Vector2D::new(3, 0);
        loop {
            let jet = self.jets[self.next_jet.next().unwrap()];

            if jet == '>' && !self.hit(shape, right(position)) {
                position = right(position);
            } else if jet == '<' && !self.hit(shape, left(position)) {
                position = left(position);
            }

            if self.hit(shape, below(position)) {
                break;
            }

            position = below(position);
        }

        self.emplace(shape, position);
    }

    fn add(&mut self, mut count: usize) -> &Tunnel {
        let mut seen: HashMap<String, (usize, usize)> = HashMap::new();

        while count > 0 {
            let hash = self.rocks.iter().map(|l| l.iter().join("")).join("");

            match seen.get(&hash) {
                Some((to_add, height)) => {
                    let p_height = self.height() - height;
                    let p_length = to_add - count;

                    self.flushed += (count / p_length) * p_height;
                    count %= p_length;
                    break;
                }
                None => {
                    seen.insert(hash, (count, self.height()));
                    self.add_one();
                    count -= 1;
                }
            }
        }

        for _ in 0..count {
            self.add_one();
        }

        self
    }

    fn height(&self) -> usize {
        self.rocks.len() + self.flushed - 1
    }
}

struct Day17;
impl Solution<Input, Output> for Day17 {
    fn parse_input<P: AsRef<Path>>(pathname: P) -> Input {
        file_to_string(pathname).trim_end().chars().collect()
    }

    fn part_1(input: &Input) -> Output {
        Tunnel::new(100, input.clone()).add(2022).height()
    }

    fn part_2(input: &Input) -> Output {
        Tunnel::new(100, input.clone()).add(1000000000000).height()
    }
}

fn main() -> Result<()> {
    // Day17::run("sample")
    Day17::main()
}

test_sample!(day_17, Day17, 3068, 1514285714288);

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

    #[test]
    fn test_infinite_index() {
        let index = InfiniteIndex::new(3);

        assert_eq!(
            index.take(10).collect::<Vec<usize>>(),
            vec![0, 1, 2, 0, 1, 2, 0, 1, 2, 0]
        );
    }
}