#![feature(try_blocks)]
// Bevy basically forces "complex types" with Querys
#![allow(clippy::type_complexity)]

use games::*;

#[cfg(test)]
mod test;

const TETRIS: RenderLayers = RenderLayers::layer(1);
const BATTLER: RenderLayers = RenderLayers::layer(2);

// TODO: When piece is near wall and rotates, move it over if it fits
// TODO: Make falling based on a timer resource ticking
//       This allows us to tune the falling rate over time

fn main() {
    App::new()
        .add_plugins(BaseGamePlugin {
            name: "falling-block-adventure".into(),
            target_resolution: (640, 480).into(),
            game_type: GameType::Two,
            ..default()
        })
        .init_state::<GameState>()
        .init_resource::<ShapesBuffer>()
        .init_resource::<ShapeStore>()
        .init_resource::<Score>()
        .init_resource::<OutputImages>()
        .init_resource::<LevelGoal>()
        .add_systems(
            Startup,
            (
                init_tetris,
                init_debug_ui,
                init_cameras,
                init_ui.after(init_cameras),
                init_battler,
            ),
        )
        .add_systems(
            Update,
            game_state_machine.run_if(state_changed::<GameState>)
        )
        // Input and basic systems
        .add_systems(
            Update,
            (
                kb_input.run_if(on_message::<KeyboardInput>),
                toggle_state_visibility::<GameState>.run_if(state_changed::<GameState>),
            ),
        )
        .add_systems(
            Update,
            (
                update_next_shapes
                    .run_if(resource_changed::<ShapesBuffer>.or(resource_added::<ShapesBuffer>)),
                add_piece
                    .run_if(not(any_with_component::<Shape>))
                    .run_if(in_state(GameState::Play))
                    .after(update_next_shapes),
                falling
                    .run_if(in_state(GameState::Play))
                    .run_if(on_timer(Duration::from_secs(1))),
                update_position.run_if(any_component_changed::<GridPosition>),
                update_shape_blocks
                    .run_if(any_component_changed::<ShapeLayout>.or(any_component_changed::<GridPosition>))
                    .after(update_position),
                deactivate_shape
                    .run_if(any_component_removed::<Shape>)
                    .after(update_shape_blocks),
                // Clearing lines systems
                clear_line.run_if(any_component_changed::<LineBlocks>),
                adjust_block_lines
                    .run_if(any_component_changed::<Line>)
                    .after(clear_line),
                assert_grid_position_uniqueness.run_if(any_component_changed::<GridPosition>),
                sync_health
                    .run_if(any_component_changed::<Health>.or(any_component_added::<Health>)),
                damage_on_place_shape.run_if(any_component_removed::<Shape>),
                damage_on_clear_line.run_if(any_component_removed::<LineBlock>),
                damage_over_time.run_if(on_timer(Duration::from_secs(5))),
                check_level_goal.run_if(resource_changed::<Score>),
            ),
        )
        // UI systems
        .add_systems(
            Update,
            (
                sync_resource_to_ui::<ShapesBuffer>.run_if(resource_changed::<ShapesBuffer>),
                sync_resource_to_ui::<ShapeStore>.run_if(resource_changed::<ShapeStore>),
                sync_resource_to_ui::<Score>.run_if(resource_changed::<Score>),
                sync_resource_to_ui::<LevelGoal>.run_if(resource_changed::<LevelGoal>),
                sync_singleton_to_ui::<Shape>.run_if(any_component_changed::<Shape>),
            ),
        )
        .add_observer(deal_damage)
        .add_observer(on_add_shape_layout)
        .add_observer(on_add_health)
        .run();
}

const SCALE: f32 = 30.0;

// Declare the size of the play area
const X_MAX: usize = 10;
const Y_MAX: usize = 20;

// The blocks making up this shape
#[derive(Component)]
#[relationship_target(relationship = ShapeBlock)]
struct ShapeBlocks(Vec<Entity>);

/// A part of a piece, i.e., a single square of a piece
#[derive(Component, Debug)]
#[relationship(relationship_target = ShapeBlocks)]
struct ShapeBlock {
    #[relationship]
    shape: Entity,
}

// The blocks making up this shape
#[derive(Component, Default, Debug)]
#[relationship_target(relationship = LineBlock)]
struct LineBlocks(Vec<Entity>);

/// A part of a piece, i.e., a single square of a piece
#[derive(Component, Debug)]
#[require(Transform, Visibility)]
#[relationship(relationship_target = LineBlocks)]
struct LineBlock {
    #[relationship]
    line: Entity,
}

// A line holds up to 10 blocks before being cleared
#[derive(Component, Debug, Clone, Copy)]
struct Line(usize);

// Just marks a block either of a shape or line
#[derive(Component, Debug)]
struct Block;

#[derive(Component, Debug, Clone, Copy, PartialEq)]
#[require(Transform, Visibility)]
struct GridPosition {
    x: usize,
    y: usize,
}

impl GridPosition {
    fn with_offset(self, other_x: isize, other_y: isize) -> Result<GridPosition, OutOfBoundsError> {
        let x = self.x as isize + other_x;
        let y = self.y as isize + other_y;

        if x >= X_MAX as isize {
            Err(OutOfBoundsError::Left)
        } else if x < 0 {
            Err(OutOfBoundsError::Right)
        } else if y < 0 {
            Err(OutOfBoundsError::Down)
        } else {
            Ok(GridPosition {
                x: x as usize,
                y: y as usize,
            })
        }
    }
}

impl Display for GridPosition {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "({},{})", self.x, self.y)
    }
}

#[derive(Error, Debug, PartialEq)]
enum OutOfBoundsError {
    #[error("Coordinates are out of bounds: Left")]
    Left,
    #[error("Coordinates are out of bounds: Right")]
    Right,
    #[error("Coordinates are out of bounds: Down")]
    Down,
}

impl Default for GridPosition {
    fn default() -> Self {
        GridPosition { x: 5, y: Y_MAX }
    }
}

impl From<&GridPosition> for Vec3 {
    fn from(GridPosition { x, y }: &GridPosition) -> Vec3 {
        // Grid Positions start in the bottom left of the area
        // So (0, 0) is the bottom left, (0, 9) is the bottom right, etc

        let x_0 = -SCALE * 5.0 + (0.5 * SCALE);
        let x = x_0 + ((*x as f32) * SCALE);

        let y_0 = -SCALE * 10.0 + (0.5 * SCALE);
        let y = y_0 + ((*y as f32) * SCALE);
        Vec3::new(x, y, 0.0)
    }
}

impl From<(usize, usize)> for GridPosition {
    fn from((x, y): (usize, usize)) -> GridPosition {
        GridPosition { x, y }
    }
}

impl std::ops::Add for GridPosition {
    type Output = Self;

    fn add(self, GridPosition { x: x2, y: y2 }: Self) -> Self {
        GridPosition {
            x: self.x + x2,
            y: self.y + y2,
        }
    }
}

impl std::ops::AddAssign<&GridPosition> for GridPosition {
    fn add_assign(&mut self, rhs: &GridPosition) {
        *self = *self + *rhs;
    }
}

#[derive(States, Clone, Eq, PartialEq, Debug, Hash, Default, Component)]
enum GameState {
    // Sets up a new game
    #[default]
    NewGame,
    // Actively playing the game
    Play,
    // Temporarily paused gameplay
    Pause,
    // Reached level goal
    LevelComplete,
    // Sets up next level
    NextLevel,
    // Failed level
    GameOver,
}

// The core of the Game state machine, coordinates other stuff happening
fn game_state_machine(
    curr: Res<State<GameState>>,
    mut next: ResMut<NextState<GameState>>,
    mut goal: ResMut<LevelGoal>,
    mut score: ResMut<Score>,
    items: Query<Entity, (With<GridPosition>, Or<(With<Block>, With<Shape>)>)>,
    mut commands: Commands,
) {
    match curr.get() {
        GameState::NewGame => {
            // Set new game goal
            *goal = LevelGoal::Score(3);

            // Reset score
            *score = Score(0);

            // Clear board
            items.iter().for_each(|i| {
                commands.entity(i).despawn();
            });

            // Set goal based on current goal
            next.set(GameState::Play);
        },
        GameState::Play => (), // handled by user input
        GameState::Pause => (), // handled by user input
        GameState::LevelComplete => (), // hanled by user input
        GameState::NextLevel => {
            // Increase goal
            *goal = match *goal {
                LevelGoal::Score(n) => LevelGoal::Score(n + 5),
            };

            // Reset score
            *score = Score(0);

            // Clear board
            items.iter().for_each(|i| {
                commands.entity(i).despawn();
            });

            // Progress to play
            next.set(GameState::Play)
        }
        GameState::GameOver => (),// handled by user input
    }
}

#[derive(Resource)]
enum LevelGoal {
    Score(usize),
}

impl Default for LevelGoal {
    fn default() -> Self {
        Self::Score(0)
    }
}

impl Display for LevelGoal {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            LevelGoal::Score(n) => write!(f, "{n}")
        }
    }
}

#[derive(Resource, Debug)]
struct Visuals {
    mesh: Handle<Mesh>,
    material_o: Handle<ColorMaterial>,
    material_t: Handle<ColorMaterial>,
    material_l: Handle<ColorMaterial>,
    material_j: Handle<ColorMaterial>,
    material_s: Handle<ColorMaterial>,
    material_z: Handle<ColorMaterial>,
    material_i: Handle<ColorMaterial>,
}

#[derive(Resource, Debug, Default)]
struct Score(usize);

impl Display for Score {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.0)
    }
}

#[derive(Component, Debug)]
struct Health(f32);

/// ShapesBuffer resource stores non-active shapes
#[derive(Resource, Debug, Default)]
struct ShapesBuffer {
    /// Next stores a vector of 2N shapes that will come up in play
    next: VecDeque<Shape>,
}

impl Display for ShapesBuffer {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        if let Some(shape) = self.next.front() {
            write!(f, "{shape}")
        } else {
            write!(f, "ERR")
        }
    }
}

#[derive(Resource, Debug, Default)]
struct ShapeStore(Option<Shape>);

impl Display for ShapeStore {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match &self.0 {
            Some(s) => write!(f, "{}", ShapeLayout::from_shape(s).as_ascii()),
            None => write!(f, "---"),
        }
    }
}

#[derive(Component)]
struct GridBackground;

fn init_tetris(
    mut commands: Commands,
    mut meshes: ResMut<Assets<Mesh>>,
    mut materials: ResMut<Assets<ColorMaterial>>,
) {
    let mesh = meshes.add(Rectangle::new(SCALE, SCALE));
    commands.insert_resource(Visuals {
        mesh: mesh.clone(),
        material_o: materials.add(ColorMaterial {
            color: YELLOW.into(),
            ..default()
        }),
        material_t: materials.add(ColorMaterial {
            color: PURPLE.into(),
            ..default()
        }),
        material_l: materials.add(ColorMaterial {
            color: ORANGE.into(),
            ..default()
        }),
        material_j: materials.add(ColorMaterial {
            color: BLUE.into(),
            ..default()
        }),
        material_s: materials.add(ColorMaterial {
            color: LIME.into(),
            ..default()
        }),
        material_z: materials.add(ColorMaterial {
            color: RED.into(),
            ..default()
        }),
        material_i: materials.add(ColorMaterial {
            color: AQUA.into(),
            ..default()
        }),
    });

    (0..Y_MAX).for_each(|y| {
        // Spawn tiles for this line
        (0..X_MAX).for_each(|x| {
            commands.spawn((
                Mesh2d(mesh.clone()),
                MeshMaterial2d(materials.add(ColorMaterial {
                    color: BLACK.into(),
                    ..default()
                })),
                GridPosition { x, y },
                Transform::from_xyz(0.0, 0.0, -1.0),
                GridBackground,
                TETRIS,
            ));
        });

        // Spawn line for holding blocks
        commands.spawn((Line(y), LineBlocks::default(), TETRIS));
    });
}

#[derive(Component, Debug)]
struct Protagonist;

#[derive(Component, Debug)]
struct Enemy;

fn init_battler(
    mut commands: Commands,
    mut meshes: ResMut<Assets<Mesh>>,
    mut materials: ResMut<Assets<ColorMaterial>>,
) {
    {
        let mat = materials.add(ColorMaterial {
            color: BLUE.into(),
            ..default()
        });
        let mesh = meshes.add(Ellipse::new(SCALE, SCALE * 2.0));
        let t = Transform::from_xyz(-3.0 * SCALE, 0.0, 0.0);
        commands.spawn((
            Mesh2d(mesh),
            MeshMaterial2d(mat),
            t,
            BATTLER,
            Health(100.0),
            Name::new("Protagonist"),
            Protagonist,
        ));
    }

    {
        let mat = materials.add(ColorMaterial {
            color: RED.into(),
            ..default()
        });
        let mesh = meshes.add(Ellipse::new(SCALE, SCALE));
        let t = Transform::from_xyz(3.0 * SCALE, 0.0, 0.0);
        commands.spawn((
            Mesh2d(mesh),
            MeshMaterial2d(mat),
            t,
            BATTLER,
            Health(100.0),
            Name::new("ENEMY"),
            Enemy,
        ));
    }
}

#[derive(Resource, Default)]
struct OutputImages {
    tetris: Handle<Image>,
    battler: Handle<Image>,
}

#[derive(Component)]
struct TetrisCamera;

#[derive(Component)]
struct BattlerCamera;

fn init_cameras(
    mut commands: Commands,
    mut images: ResMut<Assets<Image>>,
    mut output_images: ResMut<OutputImages>,
) {
    fn render_target_image((width, height): (u32, u32)) -> Image {
        let size = Extent3d {
            width,
            height,
            ..default()
        };
        let dimension = TextureDimension::D2;
        let format = TextureFormat::Bgra8UnormSrgb;
        let asset_usage = RenderAssetUsages::all();
        let mut image = Image::new_uninit(size, dimension, format, asset_usage);

        image.texture_descriptor.usage = TextureUsages::TEXTURE_BINDING
            | TextureUsages::COPY_DST
            | TextureUsages::RENDER_ATTACHMENT;

        image
    }

    {
        let (width, height) = (SCALE as u32 * X_MAX as u32, SCALE as u32 * Y_MAX as u32);
        let img = render_target_image((width, height));
        let target = {
            let handle = images.add(img);
            output_images.tetris = handle.clone();
            let scale_factor = FloatOrd(1.0);

            RenderTarget::Image(ImageRenderTarget {
                handle,
                scale_factor,
            })
        };
        commands.spawn((
            Camera2d,
            Camera {
                order: 1,
                target,
                clear_color: ClearColorConfig::Custom(WHITE.into()),
                ..default()
            },
            TETRIS,
            TetrisCamera,
        ));
    }

    {
        let (width, height) = (SCALE as u32 * Y_MAX as u32, SCALE as u32 * X_MAX as u32);
        let img = render_target_image((width, height));
        let target = {
            let handle = images.add(img);
            output_images.battler = handle.clone();
            let scale_factor = FloatOrd(1.0);

            RenderTarget::Image(ImageRenderTarget {
                handle,
                scale_factor,
            })
        };
        commands.spawn((
            Camera2d,
            Camera {
                order: 2,
                target,
                clear_color: ClearColorConfig::Custom(WHITE.into()),
                ..default()
            },
            BATTLER,
            BattlerCamera,
        ));
    }
}

fn init_ui(mut commands: Commands, output_images: Res<OutputImages>, images: Res<Assets<Image>>) {
    commands
        .spawn((Node {
            align_self: AlignSelf::Center,
            justify_self: JustifySelf::End,
            flex_direction: FlexDirection::Column,
            ..default()
        },))
        .with_children(|parent| {
            parent
                .spawn((Node {
                    flex_direction: FlexDirection::Column,
                    align_items: AlignItems::Center,
                    ..default()
                },))
                .with_children(|parent| {
                    parent.spawn(Text::new("Next:"));
                    parent.spawn((Text::new("???"), SyncResource::<ShapesBuffer>::default()));
                });
            parent
                .spawn((Node {
                    flex_direction: FlexDirection::Column,
                    align_items: AlignItems::Center,
                    ..default()
                },))
                .with_children(|parent| {
                    parent.spawn(Text::new("Swap:"));
                    parent.spawn((Text::new("???"), SyncResource::<ShapeStore>::default()));
                });
            parent
                .spawn((Node {
                    flex_direction: FlexDirection::Column,
                    align_items: AlignItems::Center,
                    ..default()
                },))
                .with_children(|parent| {
                    parent.spawn(Text::new("Score:"));
                    parent.spawn((Text::new("???"), SyncResource::<Score>::default()));
                });
            parent
                .spawn((Node {
                    flex_direction: FlexDirection::Column,
                    align_items: AlignItems::Center,
                    ..default()
                },))
                .with_children(|parent| {
                    parent.spawn(Text::new("Goal:"));
                    parent.spawn((Text::new("???"), SyncResource::<LevelGoal>::default()));
                });
        });

    commands
        .spawn((
            Node {
                align_self: AlignSelf::Center,
                justify_self: JustifySelf::Center,
                ..default()
            },
            ZIndex(1),
            GameState::Pause,
            children![(Text::new("paused"), TextColor(WHITE.into()), GameState::Pause)],
        ));

    commands
        .spawn((
            Node {
                align_self: AlignSelf::Center,
                justify_self: JustifySelf::Center,
                ..default()
            },
            ZIndex(1),
            GameState::LevelComplete,
            children![(Text::new("you did it!"), TextColor(WHITE.into()), GameState::LevelComplete)],
        ));

    commands
        .spawn((
            Node {
                align_self: AlignSelf::Center,
                justify_self: JustifySelf::Center,
                ..default()
            },
            ZIndex(1),
            GameState::GameOver,
            children![(Text::new("aww, you suck. i'm sorry."), TextColor(WHITE.into()), GameState::GameOver)],
        ));

    commands
        .spawn((
            Node {
                align_self: AlignSelf::End,
                justify_self: JustifySelf::Center,
                border: UiRect::all(Val::Px(5.0)),
                ..default()
            },
            BorderColor::all(WHITE),
        ))
        .with_children(|parent| {
            let img = images.get(&output_images.tetris).unwrap();
            parent.spawn((
                Node {
                    width: Val::Px(img.size_f32().x * 0.75),
                    height: Val::Px(img.size_f32().y * 0.75),
                    ..default()
                },
                ImageNode {
                    image: output_images.tetris.clone(),
                    image_mode: NodeImageMode::Stretch,
                    ..default()
                },
            ));
        });

    commands
        .spawn((
            Node {
                align_self: AlignSelf::Start,
                justify_self: JustifySelf::Center,
                border: UiRect::all(Val::Px(5.0)),
                ..default()
            },
            BorderColor::all(WHITE),
        ))
        .with_children(|parent| {
            let img = images.get(&output_images.battler).unwrap();
            parent.spawn((
                Node {
                    width: Val::Px(img.size_f32().x * 0.75),
                    height: Val::Px(img.size_f32().y * 0.75),
                    ..default()
                },
                ImageNode {
                    image: output_images.battler.clone(),
                    image_mode: NodeImageMode::Stretch,
                    ..default()
                },
            ));
        });
}

fn init_debug_ui(mut commands: Commands) {
    commands
        .spawn((
            Node {
                top: Val::Px(0.0),
                left: Val::Px(0.0),
                ..default()
            },
            DebuggingState::On,
        ))
        .with_children(|parent| {
            parent.spawn((
                Node::default(),
                children![(Text::new("SHAPE"), SyncSingleton::<Shape>::default()),],
            ));
        });
}

/// Enum describing a shape option
#[derive(Component, Debug, Clone)]
enum Shape {
    O,
    T,
    L,
    J,
    S,
    Z,
    I,
}

impl Default for Shape {
    fn default() -> Self {
        Self::T
    }
}

impl Display for Shape {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{self:?}")
    }
}

impl Shape {
    fn reposition(
        (x_offset, y_offset): (isize, isize),
        center: &GridPosition,
    ) -> Result<GridPosition, OutOfBoundsError> {
        center.with_offset(x_offset, y_offset)
    }
}

#[derive(Component, Debug, Clone, PartialEq, Eq)]
struct ShapeLayout {
    inner: Vec<Vec<u8>>,
}

impl From<Vec<Vec<u8>>> for ShapeLayout {
    fn from(inner: Vec<Vec<u8>>) -> Self {
        Self { inner }
    }
}

impl ShapeLayout {
    fn from_shape(s: &Shape) -> Self {
        match s {
            Shape::O => Self::new_o(),
            Shape::T => Self::new_t(),
            Shape::L => Self::new_l(),
            Shape::J => Self::new_j(),
            Shape::S => Self::new_s(),
            Shape::Z => Self::new_z(),
            Shape::I => Self::new_i(),
        }
    }

    fn new_o() -> Self {
        vec![vec![1, 1], vec![1, 1]].into()
    }

    fn new_t() -> Self {
        vec![
            vec![0, 1, 0],
            vec![1, 1, 1],
            vec![0, 0, 0]
        ].into()
    }

    fn new_l() -> Self {
        vec![vec![1, 0], vec![1, 0], vec![1, 1]].into()
    }

    fn new_j() -> Self {
        vec![vec![0, 1], vec![0, 1], vec![1, 1]].into()
    }

    fn new_s() -> Self {
        vec![vec![0, 1, 1], vec![1, 1, 0]].into()
    }

    fn new_z() -> Self {
        vec![vec![1, 1, 0], vec![0, 1, 1]].into()
    }

    fn new_i() -> Self {
        vec![vec![1], vec![1], vec![1], vec![1]].into()
    }

    fn rotated(&self) -> Self {
        let mut inner = vec![];
        for _ in 0..self.inner[0].len() {
            inner.push(vec![]);
        }
        for y in self.inner.iter() {
            for (j, x) in y.iter().enumerate() {
                inner[j].insert(0, *x);
            }
        }
        ShapeLayout { inner }
    }

    fn center(&self) -> (usize, usize) {
        let mid_x = match self.inner[0].len() % 2 {
            0 => (self.inner[0].len() - 1) / 2,
            1 => self.inner[0].len() / 2,

            _ => panic!("That's not how mod works!"),
        };

        let mid_y = match self.inner.len() % 2 {
            // If we have an even number of elements
            0 => self.inner.len().div_ceil(2),
            // If we have an odd number of elements
            1 => self.inner.len() / 2,

            _ => panic!("That's not how mod works!"),
        };

        (mid_x, mid_y)
    }

    fn coordinates_at(
        &self,
        center: &GridPosition,
    ) -> impl Iterator<Item = Result<GridPosition, OutOfBoundsError>> {
        self.coordinates()
            .map(|(x, y)| Shape::reposition((x, y), center))
    }

    fn coordinates(&self) -> impl Iterator<Item = (isize, isize)> {
        let (mid_x, mid_y) = self.center();
        // Loop over outer vec (i)
        self.inner
            .iter()
            .rev()
            .enumerate()
            .flat_map(move |(i, ys)| {
                ys.iter().enumerate().filter_map(move |(j, val)| {
                    // this_x = j - mid_x
                    let x = j as isize - mid_x as isize;
                    // this_y = i - mid_y
                    let y = i as isize - mid_y as isize;

                    (*val > 0).then_some((x, y))
                })
            })
    }

    // TODO: Make fnctional with chaining
    fn as_ascii(&self) -> String {
        let mut s = String::new();
        for y in self.inner.iter() {
            for x in y {
                match x {
                    0 => s.push('0'),
                    1 => s.push('1'),
                    _ => panic!("aur nau"),
                }
            }
            s.push('\n');
        }
        s
    }

    fn height(&self) -> usize {
        self.inner.len()
    }
}

fn update_position(
    mut changed: Query<
        (Entity, &GridPosition, &mut Transform),
        Or<(Added<GridPosition>, Changed<GridPosition>)>,
    >,
) {
    changed.iter_mut().for_each(|(e, gp, mut t)| {
        let v3: Vec3 = gp.into();
        debug!(
            "Updating {e} with grid position {:?} to coordinates {:?}",
            gp, v3
        );

        t.translation.x = v3.x;
        t.translation.y = v3.y;
    });
}

fn update_shape_blocks(
    query: Query<
        (Entity, &ShapeLayout, &GridPosition),
        Or<(
            Changed<ShapeLayout>,
            Changed<GridPosition>,
        )>,
    >,
    mut blocks: Query<&mut GridPosition, (With<ShapeBlock>, Without<ShapeLayout>)>,
) {
    query.iter().for_each(|(e, sl, center)| {
        debug!("Setting piece: {e:?} {center:?}\n{}", sl.as_ascii());

        debug_assert!(!blocks.is_empty());

        let mut p = sl.coordinates_at(center);
        blocks.iter_mut().for_each(|mut gp| {
            match p.next() {
                Some(next) => match next {
                    Ok(next_gp) => *gp = next_gp,
                    Err(_) => warn!("Next coordinate was an Err?"),
                },
                None => warn!("Next coordinate was a None?"),
            }
        });
    });
}

fn on_add_shape_layout(
    trigger: On<Add, ShapeLayout>,
    visuals: Res<Visuals>,
    query: Query<(Entity, &Shape, &ShapeLayout, &GridPosition)>,
    mut commands: Commands,
) {
    let (e, s, sl, center) = query.get(trigger.entity).unwrap();
    let mesh = Mesh2d(visuals.mesh.clone());
    let mat = match s {
        Shape::O => &visuals.material_o,
        Shape::T => &visuals.material_t,
        Shape::L => &visuals.material_l,
        Shape::J => &visuals.material_j,
        Shape::S => &visuals.material_s,
        Shape::Z => &visuals.material_z,
        Shape::I => &visuals.material_i,
    };
    commands
        .entity(e)
        .with_related_entities::<ShapeBlock>(|parent| {
            sl.coordinates_at(center).for_each(|gp| {
                parent
                    .spawn((mesh.clone(), MeshMaterial2d(mat.clone()), gp.unwrap(), Block, TETRIS))
                    .observe(movement);
            });
        });
}

fn kb_input(
    mut events: MessageReader<KeyboardInput>,
    mut query: Query<Entity, With<Shape>>, // Single?
    curr: Res<State<GameState>>,
    mut next: ResMut<NextState<GameState>>,
    mut commands: Commands,
) {
    events.read().for_each(
        |KeyboardInput {
             key_code, state, ..
         }| {
            if let ButtonState::Pressed = state {
                query.iter_mut().for_each(|e| {
                    match curr.get() {
                        GameState::Play => match key_code {
                            // Up arrow should rotate if in falling mode
                            // Only move up if in falling::off mode
                            KeyCode::ArrowUp => {
                                commands.entity(e).trigger(|entity| Movement {
                                    entity,
                                    direction: MovementDirection::Rotate,
                                });
                            }
                            KeyCode::ArrowDown => {
                                commands.entity(e).trigger(|entity| Movement {
                                    entity,
                                    direction: MovementDirection::Down,
                                });
                            }
                            KeyCode::ArrowLeft => {
                                commands.entity(e).trigger(|entity| Movement {
                                    entity,
                                    direction: MovementDirection::Left,
                                });
                            }
                            KeyCode::ArrowRight => {
                                commands.entity(e).trigger(|entity| Movement {
                                    entity,
                                    direction: MovementDirection::Right,
                                });
                            }
                            KeyCode::Enter => {
                                commands.entity(e).trigger(|entity| Movement {
                                    entity,
                                    direction: MovementDirection::Skip,
                                });
                            }
                            KeyCode::Space => {
                                commands.entity(e).trigger(|entity| Swap { entity });
                            }
                            KeyCode::Escape => next.set(GameState::Pause),
                            _ => (),
                        },
                        GameState::Pause => match key_code {
                            KeyCode::Escape | KeyCode::Enter => next.set(GameState::Play),
                            _ => ()
                        }
                        GameState::GameOver => match key_code {
                            KeyCode::Escape => todo!("Quit the game"),
                            KeyCode::Enter => next.set(GameState::NewGame),
                            _ => ()
                        },
                        GameState::LevelComplete => match key_code {
                            KeyCode::Escape => todo!("Quit the game"),
                            KeyCode::Enter => next.set(GameState::NextLevel),
                            _ => ()
                        },
                        GameState::NewGame | GameState::NextLevel => (),
                    }
                });
            }
        },
    );
}

fn falling(mut shape: Query<Entity, With<Shape>>, mut commands: Commands) {
    shape.iter_mut().for_each(|e| {
        debug!("Making {:?} fall", e);
        commands.entity(e).trigger(|entity| Movement {
            entity,
            direction: MovementDirection::Down,
        });
    });
}

fn add_piece(mut commands: Commands, mut shapes: ResMut<ShapesBuffer>) {
    let shape = shapes.next.pop_front().unwrap();
    let shape_layout = ShapeLayout::from_shape(&shape);
    commands
        .spawn((
            GridPosition {
                y: Y_MAX - shape_layout.height(),
                ..default()
            },
            shape_layout,
            shape,
            TETRIS,
        ))
        .observe(movement)
        .observe(swap);
}

/// When a line reaches 10 blocks, clear it
fn clear_line(
    changed_lines: Query<Entity, Changed<LineBlocks>>,
    mut lines: Query<&mut Line>,
    line_blocks: Query<&LineBlocks>,
    mut score: ResMut<Score>,
    mut commands: Commands,
) {
    let mut cleared_lines: Vec<usize> = changed_lines
        .iter()
        .filter_map(|e| try { (Some(e)?, line_blocks.get(e).ok()?, lines.get(e).ok()?) })
        .filter_map(|(e, lb, Line(i))| {
            if lb.0.len() == 10 {
                commands
                    .entity(e)
                    .despawn_related::<LineBlocks>()
                    .insert(LineBlocks::default());
                score.0 += 1;
                info!("New score: {:?}", score.0);
                Some(*i)
            } else {
                None
            }
        })
        .sorted()
        .collect();

    if !cleared_lines.is_empty() {
        info!("Cleared lines: {:?}", cleared_lines);

        #[cfg(debug_assertions)]
        {
            debug_assert_eq!(lines.iter().count(), 20, "There should be 20 lines");
            // Check that all line numbers are present
            lines
                .iter()
                .map(|Line(i)| i)
                .sorted()
                .enumerate()
                .for_each(|(i, line_num)| {
                    debug_assert_eq!(i, *line_num, "Line numbers should match their sorted index");
                });
        }

        let original_cleared_lines_len = cleared_lines.len();

        // Iterate over all lines in reverse sorted order (largest to smallest)
        lines
            .iter_mut()
            .sorted_by(|i, j| i.0.cmp(&j.0))
            .rev()
            .for_each(|mut l| {
                // If the current index is in the set of cleared lines, move it to the top
                // Otherwise, move it down by the number of cleared lines
                if cleared_lines.contains(&l.0) {
                    // Move to the N-offset line number (top, top-1, etc)
                    let offset = original_cleared_lines_len - cleared_lines.len();
                    debug!("Moving line {:?}->{:?}", l.0, Y_MAX - 1 - offset);
                    l.0 = Y_MAX - 1 - offset;
                    cleared_lines.pop();
                } else {
                    debug!("Moving line {:?}->{:?}", l.0, l.0 - cleared_lines.len());
                    l.0 -= cleared_lines.len();
                }
            });
    }
}

fn adjust_block_lines(
    query: Query<(Entity, &Line), Changed<Line>>,
    parent: Query<&LineBlocks>,
    mut blocks: Query<&mut GridPosition>,
) {
    query.iter().for_each(|(e, Line(i))| {
        parent.iter_descendants(e).for_each(|block| {
            if let Ok(mut gp) = blocks.get_mut(block) {
                gp.y = *i;
            }
        });
    });
}

/// Swap the current piece out
#[derive(EntityEvent, Message, Copy, Clone, PartialEq)]
struct Swap {
    entity: Entity,
}

/// Movement events evented on the piece
#[derive(Message, EntityEvent, Copy, Clone, PartialEq)]
struct Movement {
    entity: Entity,
    direction: MovementDirection,
}

#[derive(Clone, Copy, Debug, PartialEq)]
enum MovementDirection {
    Down,
    Left,
    Right,
    Rotate,
    Skip,
}

// TODO: When out of bounds left/right, try to move piece away from wall
fn movement(
    event: On<Movement>,
    mut grid_positions: Query<
        &mut GridPosition,
        Or<(With<ShapeBlock>, With<ShapeBlocks>, Without<LineBlock>)>,
    >,
    mut shape_layouts: Query<&mut ShapeLayout>,
    inactive: Query<&GridPosition, (Without<ShapeBlock>, Without<ShapeBlocks>, With<LineBlock>)>,
    mut commands: Commands,
) {
    if let (Ok(this_shape_layout), Ok(center)) = (
        shape_layouts.get_mut(event.entity),
        grid_positions.get(event.entity),
    ) {
        let new_positions = match event.event().direction {
            MovementDirection::Down => vec![center.with_offset(0, -1)],
            MovementDirection::Left => vec![center.with_offset(-1, 0)],
            MovementDirection::Right => vec![center.with_offset(1, 0)],
            MovementDirection::Rotate => vec![Ok(*center)],
            MovementDirection::Skip => (1..=center.y)
                .map(|i| center.with_offset(0, -(i as isize)))
                .collect(),
        };
        let new_shape_layout = match event.event().direction {
            MovementDirection::Down
            | MovementDirection::Left
            | MovementDirection::Right
            | MovementDirection::Skip => this_shape_layout.clone(),
            MovementDirection::Rotate => this_shape_layout.rotated(),
        };
        debug!(
            "Proposed change: {:?}\n{}",
            new_positions,
            new_shape_layout.as_ascii()
        );
        for position in new_positions {
            match position {
                Err(OutOfBoundsError::Left) | Err(OutOfBoundsError::Right) => (), // Do nothing
                Err(OutOfBoundsError::Down) => {
                    commands.entity(event.entity).remove::<Shape>();
                }
                Ok(new_center) => {
                    let new_blocks = new_shape_layout.coordinates_at(&new_center);
                    for block_gp in new_blocks {
                        match block_gp {
                            Err(OutOfBoundsError::Left) | Err(OutOfBoundsError::Right) => {
                                return;
                            } // Do nothing
                            Err(OutOfBoundsError::Down) => {
                                commands.entity(event.entity).remove::<Shape>();
                                return;
                            }
                            Ok(gp) => {
                                for other_gp in inactive.iter() {
                                    // If there would be a collision between blocks
                                    if gp == *other_gp {
                                        // And we are moving down
                                        if event.event().direction == MovementDirection::Down {
                                            // De-activate this piece
                                            commands.entity(event.entity).remove::<Shape>();
                                        }
                                        // Regardless, cancel the move
                                        return;
                                    }
                                }
                            }
                        }
                    }
                    debug!("Checks passed for {position:?}, committing change");

                    // Update center
                    let mut gp = grid_positions.get_mut(event.entity).unwrap();
                    *gp = new_center;

                    // Update shape/rotation
                    let mut sl = shape_layouts.get_mut(event.entity).unwrap();
                    *sl = new_shape_layout.clone();
                }
            }
        }
    } else {
        warn!("Oned movement on non-shape entity");
    }
}

fn swap(
    event: On<Swap>,
    mut shapes: Query<&mut Shape>,
    mut store: ResMut<ShapeStore>,
    mut commands: Commands,
) {
    info!("Swap called");
    match store.0.as_mut() {
        None => {
            // Copy current shape into store
            // De-activate entity
            store.0 = Some(shapes.get(event.entity).unwrap().clone());
            commands
                .entity(event.entity)
                .despawn_related::<ShapeBlocks>()
                .despawn();
        }
        Some(inner) => {
            // Copy current shape into store
            // Copy old shape into entity
            let mut curr = shapes.get_mut(event.entity).unwrap();
            std::mem::swap(&mut (*inner), &mut (*curr));
        }
    }
}

fn deactivate_shape(
    mut events: RemovedComponents<Shape>,
    grid_positions: Query<&GridPosition, With<ShapeBlock>>,
    parent: Query<&ShapeBlocks>,
    lines: Query<(Entity, &Line), With<LineBlocks>>,
    mut commands: Commands,
) {
    events.read().for_each(|target| {
        parent.iter_descendants(target).for_each(|block| {
            let GridPosition { y, .. } = grid_positions.get(block).unwrap();
            if let Some(parent_line) = lines
                .iter()
                .find_map(|(e, Line(i))| (*y == *i).then_some(e))
            {
                commands
                    .entity(parent_line)
                    .add_one_related::<LineBlock>(block);
            } else {
                error!("wtf?");
            }
        });
        commands.entity(target).despawn();
    });
}

fn update_next_shapes(mut buffer: ResMut<ShapesBuffer>) {
    // If the buffer contains less than n+1 shapes (where n is the number of possible shapes)
    // Ideally we have between 1n and 2n shapes in the `next` buffer
    while buffer.next.len() < 8 {
        // TODO: Shuffle these!
        buffer.next.extend([
            Shape::O,
            Shape::T,
            Shape::L,
            Shape::J,
            Shape::S,
            Shape::Z,
            Shape::I,
        ]);
    }
}

fn assert_grid_position_uniqueness(
    grid_positions: Query<(Entity, &GridPosition, Option<&LineBlock>, Option<&ShapeBlock>), With<Block>>,
    mut next: ResMut<NextState<GameState>>,
) {
    grid_positions.iter_combinations().for_each(|[(e1, a, lb1, sb1), (e2, b, lb2, sb2)]| {
        if a == b {
            error!("Two entities are in the same grid position: {a:?} <-> {b:?}!");
            error!("\tE1: {:?} (Parent Block: {:?}| Parent Line: {:?})!", e1, lb1, sb1);
            error!("\tE2: {:?} (Parent Block: {:?}| Parent Line: {:?})!", e2, lb2, sb2);
            next.set(GameState::Pause);
        }
        // assert_ne!(a, b, "Two entities are in the same grid position!");
    });
}

fn sync_health(
    query: Query<(Entity, &Health), Changed<Health>>,
    parent: Query<&Children>,
    mut texts: Query<&mut Text2d>,
) {
    query.iter().for_each(|(e, h)| {
        let child = parent
            .iter_descendants(e)
            .find(|child| texts.contains(*child))
            .unwrap();
        debug!("Updating health");
        let mut t = texts.get_mut(child).unwrap();
        t.0 = format!("{}", h.0);
    });
}

fn on_add_health(
    event: On<Add, Health>,
    query: Query<(&Health, &Mesh2d)>,
    meshes: Res<Assets<Mesh>>,
    mut commands: Commands,
) {
    let (h, m) = query.get(event.entity).unwrap();
    info!("Creating health display");
    commands.entity(event.entity).with_children(|parent| {
        let mesh = meshes.get(&m.0).unwrap();
        let aabb = mesh.compute_aabb().unwrap();
        let offset = Vec3::new(0.0, aabb.half_extents.y + 10.0, 0.0);
        parent.spawn((
            Text2d(format!("{}", h.0)),
            TextColor(BLACK.into()),
            Transform::from_translation(offset),
            BATTLER,
        ));
    });
}

#[derive(Message, EntityEvent)]
struct Damage {
    entity: Entity,
    quantity: f32,
}

fn deal_damage(event: On<Damage>, mut healths: Query<&mut Health>) {
    healths.get_mut(event.entity).unwrap().0 -= event.event().quantity
}

fn damage_on_place_shape(
    mut events: RemovedComponents<Shape>,
    enemies: Query<Entity, With<Enemy>>,
    mut commands: Commands,
) {
    events.read().for_each(|_| {
        enemies.iter().for_each(|e| {
            commands.entity(e).trigger(|entity| Damage {
                entity,
                quantity: 1.0,
            });
        });
    });
}

fn damage_on_clear_line(
    mut events: RemovedComponents<LineBlock>,
    enemies: Query<Entity, With<Enemy>>,
    mut commands: Commands,
) {
    events.read().for_each(|_| {
        enemies.iter().for_each(|e| {
            commands.entity(e).trigger(|entity| Damage {
                entity,
                quantity: 1.0,
            });
        });
    });
}

fn damage_over_time(protagonist: Single<Entity, With<Protagonist>>, mut commands: Commands) {
    commands.entity(*protagonist).trigger(|entity| Damage {
        entity,
        quantity: 1.0,
    });
}

fn check_level_goal(
    goal: Res<LevelGoal>,
    score: Res<Score>,
    mut next: ResMut<NextState<GameState>>,
) {
    match *goal {
        LevelGoal::Score(g) => {
            if score.0 >= g {
                next.set(GameState::LevelComplete);
            }
        }
    }
}
