blackout/src/visibility.rs

use std::{
    cell::RefCell,
    collections::{HashMap, HashSet},
    marker::PhantomData,
};

use bevy::{core::FixedTimestep, prelude::*};

use coord_2d::{Coord, Size};
use shadowcast::{vision_distance, Context, InputGrid};

use crate::{
    bevy_rapier2d::prelude::*,
    commands::RunIfExistsExt,
    core::{Angle, Player, PointLike},
    log::Log,
    map::{Map, MapConfig, MapObstruction},
};

#[derive(Component, Clone, Copy, Debug, Default, Reflect)]
#[reflect(Component)]
pub struct DontLogWhenVisible;

#[derive(Component, Clone, Debug, Default, Deref, DerefMut, Reflect)]
#[reflect(Component)]
struct LastCoordinates((i32, i32));

#[derive(Component, Clone, Debug, Default, Deref, DerefMut, Reflect)]
#[reflect(Component)]
pub struct RevealedTiles(pub Vec<bool>);

#[derive(Component, Clone, Debug, Eq, PartialEq)]
pub struct Viewshed {
    pub range: u32,
    pub visible_points: HashSet<(i32, i32)>,
}

impl Default for Viewshed {
    fn default() -> Self {
        Self {
            range: 15,
            visible_points: HashSet::new(),
        }
    }
}

impl Viewshed {
    pub fn is_point_visible(&self, point: &dyn PointLike) -> bool {
        self.visible_points.contains(&point.into())
    }

    fn update<D: 'static + Clone + Default + Send + Sync>(
        &mut self,
        viewer_entity: &Entity,
        visible_entities: &mut VisibleEntities,
        start: &dyn PointLike,
        rapier_context: &RapierContext,
        map: &Map<D>,
        visible_query: &Query<&Visible>,
        obstructions_query: &Query<&MapObstruction>,
        events: &mut EventWriter<VisibilityChanged>,
        cache: &mut HashMap<Coord, (u8, HashSet<Entity>)>,
    ) {
        let mut context: Context<u8> = Context::default();
        let vision_distance = vision_distance::Circle::new(self.range);
        let shape = Collider::cuboid(0.5 - f32::EPSILON, 0.5 - f32::EPSILON);
        let mut new_visible_entities = HashSet::new();
        let size = (map.width as u32, map.height as u32);
        let visibility_grid = VisibilityGrid(
            size,
            RefCell::new(Box::new(|coord: Coord| {
                let shape_pos = Vec2::new(coord.x as f32 + 0.5, coord.y as f32 + 0.5);
                if start.distance(&shape_pos) > self.range as f32 + 1. {
                    return u8::MAX;
                }
                if let Some((opacity, entities)) = cache.get(&coord) {
                    for e in entities {
                        new_visible_entities.insert(*e);
                    }
                    if entities.contains(viewer_entity) {
                        return 0;
                    } else {
                        return *opacity;
                    }
                }
                let mut opacity = 0;
                let mut entities = HashSet::new();
                rapier_context.intersections_with_shape(
                    shape_pos,
                    0.,
                    &shape,
                    InteractionGroups::all(),
                    Some(&|v| visible_query.get(v).is_ok()),
                    |entity| {
                        let obstruction = obstructions_query.get(entity).is_ok();
                        if obstruction {
                            entities.clear();
                        }
                        entities.insert(entity);
                        if let Ok(visible) = visible_query.get(entity) {
                            opacity = opacity.max(**visible);
                        }
                        !obstruction
                    },
                );
                for e in &entities {
                    new_visible_entities.insert(*e);
                }
                cache.insert(coord, (opacity, entities.clone()));
                if entities.contains(viewer_entity) {
                    0
                } else {
                    opacity
                }
            })),
        );
        let mut new_visible = HashSet::new();
        context.for_each_visible(
            Coord::new(start.x_i32(), start.y_i32()),
            &visibility_grid,
            &size,
            vision_distance,
            u8::MAX,
            |coord, _directions, _visibility| {
                new_visible.insert((coord.x, coord.y));
            },
        );
        if self.visible_points != new_visible {
            self.visible_points = new_visible;
        }
        if new_visible_entities != **visible_entities {
            for lost in visible_entities.difference(&new_visible_entities) {
                events.send(VisibilityChanged::Lost {
                    viewer: *viewer_entity,
                    viewed: *lost,
                });
            }
            for gained in new_visible_entities.difference(visible_entities) {
                events.send(VisibilityChanged::Gained {
                    viewer: *viewer_entity,
                    viewed: *gained,
                });
            }
            **visible_entities = new_visible_entities;
        }
    }
}

#[derive(Component, Clone, Copy, Debug, Deref, DerefMut, Reflect)]
#[reflect(Component)]
pub struct Visible(pub u8);

impl Default for Visible {
    fn default() -> Self {
        Self::opaque()
    }
}

impl Visible {
    pub fn transparent() -> Self {
        Self(0)
    }

    pub fn opaque() -> Self {
        Self(u8::MAX)
    }
}

#[derive(Component, Clone, Debug, Default, Deref, DerefMut)]
pub struct VisibleEntities(HashSet<Entity>);

#[derive(Bundle, Default)]
pub struct ViewshedBundle {
    pub viewshed: Viewshed,
    pub visible_entities: VisibleEntities,
}

impl ViewshedBundle {
    pub fn new(range: u32) -> Self {
        Self {
            viewshed: Viewshed { range, ..default() },
            ..default()
        }
    }
}

pub enum VisibilityChanged {
    Gained { viewer: Entity, viewed: Entity },
    Lost { viewer: Entity, viewed: Entity },
}

impl PointLike for Coord {
    fn x(&self) -> f32 {
        self.x as f32
    }

    fn y(&self) -> f32 {
        self.y as f32
    }
}

fn add_visibility_indices<D: 'static + Clone + Default + Send + Sync>(
    mut commands: Commands,
    query: Query<(Entity, &Map<D>), (Added<Map<D>>, Without<RevealedTiles>)>,
    map_config: Res<MapConfig>,
) {
    for (entity, map) in query.iter() {
        let count = map.width * map.height;
        commands
            .entity(entity)
            .insert(RevealedTiles(vec![map_config.start_revealed; count]));
    }
}

fn viewshed_removed(
    query: RemovedComponents<Viewshed>,
    visible_entities: Query<&VisibleEntities>,
    mut events: EventWriter<VisibilityChanged>,
) {
    for entity in query.iter() {
        if let Ok(visible) = visible_entities.get(entity) {
            for e in visible.iter() {
                events.send(VisibilityChanged::Lost {
                    viewer: entity,
                    viewed: *e,
                })
            }
        }
    }
}

pub struct VisibilityGrid<F>(pub (u32, u32), pub RefCell<Box<F>>);

impl<F> InputGrid for VisibilityGrid<F>
where
    F: FnMut(Coord) -> u8,
{
    type Grid = (u32, u32);

    type Opacity = u8;

    fn size(&self, grid: &Self::Grid) -> Size {
        Size::new(grid.0, grid.1)
    }

    fn get_opacity(&self, _grid: &Self::Grid, coord: Coord) -> Self::Opacity {
        (self.1.borrow_mut())(coord)
    }
}

fn update_viewshed<D: 'static + Clone + Default + Send + Sync>(
    mut commands: Commands,
    config: Res<RapierConfiguration>,
    positions: Query<
        (Entity, &Transform, &Collider, Option<&LastCoordinates>),
        Or<(Changed<Transform>, Changed<Visible>)>,
    >,
    visible: Query<&Visible>,
    obstructions: Query<&MapObstruction>,
    mut viewers: Query<(Entity, &mut Viewshed, &mut VisibleEntities, &Transform)>,
    map: Query<&Map<D>>,
    rapier_context: Res<RapierContext>,
    mut changed: EventWriter<VisibilityChanged>,
) {
    if !config.query_pipeline_active {
        return;
    }
    let mut to_update = HashSet::new();
    for (entity, transform, collider, last_coordinates) in positions.iter() {
        if to_update.contains(&entity) {
            continue;
        }
        let coordinates = (transform.translation.x, transform.translation.y);
        let coordinates_i32 = coordinates.i32();
        if viewers.get_mut(entity).is_ok() {
            commands.run_if_exists(entity, move |mut entity| {
                entity.insert(LastCoordinates(coordinates_i32));
            });
            if let Some(last_coordinates) = last_coordinates {
                if **last_coordinates != coordinates.i32() {
                    to_update.insert(entity);
                }
            } else {
                to_update.insert(entity);
            }
        }
        if visible.get(entity).is_err() {
            continue;
        }
        for (viewer_entity, viewshed, visible_entities, viewer_coordinates) in viewers.iter_mut() {
            if viewer_entity != entity {
                let forward = transform.local_x();
                let rotation = forward.y.atan2(forward.x);
                let distance = collider.distance_to_point(
                    transform.translation.truncate(),
                    rotation,
                    viewer_coordinates.translation.truncate(),
                    true,
                );
                if distance <= viewshed.range as f32 || visible_entities.contains(&entity) {
                    to_update.insert(viewer_entity);
                    to_update.insert(entity);
                }
            }
        }
    }
    let mut cache = HashMap::new();
    for entity in to_update.iter() {
        if let Ok((viewer_entity, mut viewshed, mut visible_entities, viewer_coordinates)) =
            viewers.get_mut(*entity)
        {
            if let Ok(map) = map.get_single() {
                viewshed.update(
                    &viewer_entity,
                    &mut visible_entities,
                    viewer_coordinates,
                    &*rapier_context,
                    map,
                    &visible,
                    &obstructions,
                    &mut changed,
                    &mut cache,
                );
            }
        }
    }
}

fn remove_visible<D: 'static + Clone + Default + Send + Sync>(
    removed: RemovedComponents<Visible>,
    mut viewers: Query<(Entity, &mut Viewshed, &mut VisibleEntities, &Transform)>,
    map: Query<&Map<D>>,
    rapier_context: Res<RapierContext>,
    visible: Query<&Visible>,
    obstructions: Query<&MapObstruction>,
    mut changed: EventWriter<VisibilityChanged>,
) {
    for removed in removed.iter() {
        let mut cache = HashMap::new();
        for (viewer_entity, mut viewshed, mut visible_entities, start) in viewers.iter_mut() {
            if !visible_entities.contains(&removed) {
                continue;
            }
            visible_entities.remove(&removed);
            if let Ok(map) = map.get_single() {
                viewshed.update(
                    &viewer_entity,
                    &mut visible_entities,
                    start,
                    &*rapier_context,
                    map,
                    &visible,
                    &obstructions,
                    &mut changed,
                    &mut cache,
                );
            }
        }
    }
}

fn update_revealed_tiles<D: 'static + Clone + Default + Send + Sync>(
    mut map: Query<(&Map<D>, &mut RevealedTiles)>,
    viewers: Query<&Viewshed, (With<Player>, Changed<Viewshed>)>,
) {
    for viewshed in viewers.iter() {
        for (map, mut revealed_tiles) in map.iter_mut() {
            for v in viewshed.visible_points.iter() {
                let idx = v.to_index(map.width);
                if idx >= revealed_tiles.len() {
                    continue;
                }
                revealed_tiles[idx] = true;
            }
        }
    }
}

fn log_visible(
    time: Res<Time>,
    mut recently_lost: Local<HashMap<Entity, Timer>>,
    mut events: EventReader<VisibilityChanged>,
    mut log: Query<&mut Log>,
    viewers: Query<(Entity, &Transform), (With<Player>, With<Viewshed>)>,
    visible: Query<(&Name, &Transform), Without<DontLogWhenVisible>>,
) {
    for timer in recently_lost.values_mut() {
        timer.tick(time.delta());
    }
    recently_lost.retain(|_entity, timer| !timer.finished());
    for event in events.iter() {
        let viewer = match event {
            VisibilityChanged::Gained { viewer, .. } => viewer,
            VisibilityChanged::Lost { viewer, .. } => viewer,
        };
        if let Ok((viewer_entity, viewer_transform)) = viewers.get(*viewer) {
            if let VisibilityChanged::Gained { viewed, .. } = event {
                if *viewed == viewer_entity || recently_lost.contains_key(viewed) {
                    continue;
                }
                if let Ok((name, transform)) = visible.get(*viewed) {
                    let viewed_coordinates = (transform.translation.x, transform.translation.y);
                    let forward = viewer_transform.local_x();
                    let yaw = Angle::Radians(forward.y.atan2(forward.x));
                    let viewer_coordinates = (
                        viewer_transform.translation.x,
                        viewer_transform.translation.y,
                    );
                    let location =
                        viewer_coordinates.direction_and_distance(&viewed_coordinates, Some(yaw));
                    if let Ok(mut log) = log.get_single_mut() {
                        log.push(format!("{}: {location}", name));
                    }
                }
            } else if let VisibilityChanged::Lost { viewed, .. } = event {
                recently_lost.insert(*viewed, Timer::from_seconds(2., false));
            }
        }
    }
}

pub const LOG_VISIBLE_LABEL: &str = "LOG_VISIBLE";

pub struct VisibilityPlugin<D: 'static + Clone + Default + Send + Sync>(PhantomData<D>);

impl<D: 'static + Clone + Default + Send + Sync> Default for VisibilityPlugin<D> {
    fn default() -> Self {
        Self(Default::default())
    }
}

pub const UPDATE_VIEWSHED_LABEL: &str = "UPDATE_VIEWSHED";

impl<D: 'static + Clone + Default + Send + Sync> Plugin for VisibilityPlugin<D> {
    fn build(&self, app: &mut App) {
        const UPDATE_VIEWSHED_STAGE: &str = "UPDATE_VIEWSHED_STAGE";
        app.add_event::<VisibilityChanged>()
            .add_stage_after(
                CoreStage::PreUpdate,
                UPDATE_VIEWSHED_STAGE,
                SystemStage::parallel(),
            )
            .add_system_to_stage(CoreStage::PreUpdate, add_visibility_indices::<D>)
            .add_system_set_to_stage(
                UPDATE_VIEWSHED_STAGE,
                SystemSet::new()
                    .with_run_criteria(FixedTimestep::step(0.05))
                    .with_system(update_viewshed::<D>)
                    .label(UPDATE_VIEWSHED_LABEL),
            )
            .add_system_to_stage(CoreStage::PostUpdate, viewshed_removed)
            .add_system_to_stage(CoreStage::PostUpdate, remove_visible::<D>)
            .add_system_to_stage(CoreStage::PreUpdate, update_revealed_tiles::<D>)
            .add_system_to_stage(CoreStage::PreUpdate, log_visible);
    }
}