use std::cmp::max;
use std::collections::HashSet;

use nom::{branch::alt, character::complete::char, multi::many1, IResult, Parser};

#[derive(Debug, Clone, Copy, PartialEq)]
enum Dir {
    Left,
    Right,
}
use Dir::*;

struct Jets {
    dirs: Vec<Dir>,
    current_jet_index: usize,
}

impl Jets {
    fn next(&mut self) -> Dir {
        let dir = self.dirs[self.current_jet_index];
        self.current_jet_index = (self.current_jet_index + 1) % self.dirs.len();
        dir
    }
}

#[derive(Debug, Clone, Copy)]
enum Shape {
    HorLine,
    Plus,
    LShape,
    VerLine,
    Block,
}
use Shape::*;

impl Shape {
    fn next(&self) -> Shape {
        match self {
            HorLine => Plus,
            Plus => LShape,
            LShape => VerLine,
            VerLine => Block,
            Block => HorLine,
        }
    }

    fn height(&self) -> u64 {
        self.left().len() as u64
    }

    fn width(&self) -> u64 {
        self.under().len() as u64
    }

    fn under(&self) -> Vec<u64> {
        match self {
            HorLine => vec![0, 0, 0, 0],
            Plus => vec![1, 0, 1],
            LShape => vec![0, 0, 0],
            VerLine => vec![0],
            Block => vec![0, 0],
        }
    }

    fn left(&self) -> Vec<u64> {
        match self {
            HorLine => vec![0],
            Plus => vec![1, 0, 1],
            LShape => vec![0, 2, 2],
            VerLine => vec![0, 0, 0, 0],
            Block => vec![0, 0],
        }
    }

    fn right(&self) -> Vec<u64> {
        match self {
            HorLine => vec![0],
            Plus => vec![1, 0, 1],
            LShape => vec![0, 0, 0],
            VerLine => vec![0, 0, 0, 0],
            Block => vec![0, 0],
        }
    }

    fn shape(&self) -> Vec<(u64, u64)> {
        match self {
            HorLine => vec![(0, 0), (1, 0), (2, 0), (3, 0)],
            Plus => vec![(1, 0), (0, 1), (1, 1), (2, 1), (1, 2)],
            LShape => vec![(0, 0), (1, 0), (2, 0), (2, 1), (2, 2)],
            VerLine => vec![(0, 0), (0, 1), (0, 2), (0, 3)],
            Block => vec![(0, 0), (0, 1), (1, 0), (1, 1)],
        }
    }
}

fn parse_input(input: &str) -> IResult<&str, Jets> {
    let (input, dirs) = many1(alt((char('<').map(|_| Left), char('>').map(|_| Right))))(input)?;

    let jets = Jets {
        dirs,
        current_jet_index: 0,
    };

    Ok((input, jets))
}

fn draw_cave(cave_map: &HashSet<(u64, u64)>, falling_shape: &Shape, shape_coord: &(u64, u64)) {
    let shape_coords: Vec<(u64, u64)> = falling_shape
        .shape()
        .iter()
        .map(|&(x, y)| (shape_coord.0 + x, shape_coord.1 + y))
        .collect();

    let max_height = shape_coord.1 + falling_shape.height() - 1;

    for height in 0..max_height {
        let y = max_height - height;

        print!("|");

        for x in 1..=7 {
            if cave_map.contains(&(x, y)) {
                print!("#");
            } else if shape_coords.contains(&(x, y)) {
                print!("@");
            } else {
                print!(".");
            }
        }

        println!("|");
    }
    println!("+-------+");
}

fn fall(
    cave_map: &mut HashSet<(u64, u64)>,
    jetstream: &mut Jets,
    falling_shape: &Shape,
    highest_unit: u64,
) -> u64 {
    let mut rock_coord = (3, highest_unit + 4);
    let shape_under = falling_shape.under();
    let shape_right = falling_shape.right();
    let shape_left = falling_shape.left();
    let shape_width = falling_shape.width();

    loop {
        // First, move the shape to the left or right
        let jet_dir = jetstream.next();
        // match jet_dir {
        //     Right => println!("Jet: Right"),
        //     Left => println!("Jet: Left"),
        // }

        if jet_dir == Left {
            if rock_coord.0 > 1 {
                let move_result = shape_left.iter().enumerate().any(|(index, &indent)| {
                    cave_map.contains(&(rock_coord.0 + indent - 1, rock_coord.1 + index as u64))
                });
                if !move_result {
                    rock_coord = (rock_coord.0 - 1, rock_coord.1);
                }
            }
        } else {
            if rock_coord.0 + shape_width < 8 {
                let move_result = shape_right.iter().enumerate().any(|(index, &indent)| {
                    cave_map.contains(&(
                        rock_coord.0 + shape_width - indent,
                        rock_coord.1 + index as u64,
                    ))
                });
                if !move_result {
                    rock_coord = (rock_coord.0 + 1, rock_coord.1);
                }
            }
        }

        // draw_cave(&cave_map, &falling_shape, &rock_coord);

        // Then, move the shape down
        if rock_coord.1 > 1
            && !shape_under.iter().enumerate().any(|(index, &indent)| {
                cave_map.contains(&(rock_coord.0 + index as u64, rock_coord.1 + indent - 1))
            })
        {
            rock_coord = (rock_coord.0, rock_coord.1 - 1);
        } else {
            for (x_indent, y_indent) in falling_shape.shape() {
                cave_map.insert((rock_coord.0 + x_indent, rock_coord.1 + y_indent));
            }
            break;
        }
    }

    // draw_cave(&cave_map, &falling_shape, &rock_coord);

    max(rock_coord.1 + falling_shape.height() - 1, highest_unit)
}

fn main() {
    let input_text = include_str!("../test.txt");
    let (_rest, mut jetstream) = parse_input(input_text).unwrap();

    let mut number_of_rocks: u64 = 1000000000000;
    let mut cave_map: HashSet<(u64, u64)> = HashSet::new();
    let mut highest_unit = 0;
    let mut falling_rock = HorLine;

    while number_of_rocks > 0 {
        highest_unit = fall(&mut cave_map, &mut jetstream, &falling_rock, highest_unit);

        falling_rock = falling_rock.next();
        number_of_rocks -= 1;
    }

    println!("Height of the rock pile: {highest_unit}");
}