blackout/src/visibility.rs

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

use avian2d::prelude::*;
use bevy::prelude::*;
use coord_2d::{Coord, Size};
use shadowcast::{vision_distance, Context, InputGrid};

use crate::{
    core::{GlobalTransformExt, Player, PointLike},
    log::Log,
    map::{Map, MapPlugin},
};

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

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

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

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

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

    fn update(
        &mut self,
        commands: &mut Commands,
        viewer_entity: &Entity,
        visible_entities: &mut VisibleEntities,
        start: &Vec2,
        opacity_map: &OpacityMap,
        sensors: &Query<&Sensor>,
        spatial_query: &SpatialQuery,
    ) {
        let mut context: Context<u8> = Context::default();
        let vision_distance = vision_distance::Circle::new(self.range);
        let size = (
            (start.x.abs() + self.range as f32 * 2.) as u32,
            (start.y.abs() + self.range as f32 * 2.) as u32,
        );
        let visibility_grid = VisibilityGrid(
            *viewer_entity,
            start.as_ivec2(),
            self.range,
            opacity_map.clone(),
        );
        let mut new_visible = HashSet::new();
        let mut new_visible_entities = HashSet::new();
        // println!("Start: {viewer_entity}");
        context.for_each_visible(
            Coord::new(start.x_i32(), start.y_i32()),
            &visibility_grid,
            &size,
            vision_distance,
            u8::MAX,
            |coord, _directions, _visibility| {
                let coord = IVec2::new(coord.x, coord.y);
                new_visible.insert(coord);
                // println!("Checking {coord:?}");
                if let Some((_, entities)) = opacity_map.get(&coord) {
                    for e in entities {
                        if entities.len() == 1 || sensors.contains(*e) {
                            // println!("Spotted {e:?}");
                            new_visible_entities.insert(*e);
                        } else {
                            let should_push = std::cell::RefCell::new(true);
                            let coord = Vec2::new(coord.x as f32 + 0.5, coord.y as f32 + 0.5);
                            let dir = Dir2::new_unchecked((coord - *start).normalize_or_zero());
                            // println!("Casting from {coord} to {dir:?}");
                            spatial_query.cast_ray_predicate(
                                *start,
                                dir,
                                dir.distance(*start),
                                true,
                                &default(),
                                &|entity| {
                                    if *e != entity && entities.contains(e) {
                                        // println!("{entities:?} contains {e}");
                                        *should_push.borrow_mut() = false;
                                    }
                                    true
                                },
                            );
                            if *should_push.borrow() {
                                new_visible_entities.insert(*e);
                            }
                        }
                    }
                }
            },
        );
        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) {
                // println!("Lost {lost}");
                commands.trigger_targets(VisibilityChanged::Lost(*lost), *viewer_entity);
            }
            for gained in new_visible_entities.difference(visible_entities) {
                // println!("Gained {gained}");
                commands.trigger_targets(VisibilityChanged::Gained(*gained), *viewer_entity);
            }
            **visible_entities = new_visible_entities;
        }
    }
}

#[derive(Component, Clone, Copy, Debug, Deref, DerefMut, Eq, PartialEq, 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(Resource, Clone, Deref, DerefMut, Default)]
pub struct OpacityMap(HashMap<IVec2, (u8, HashSet<Entity>)>);

impl OpacityMap {
    fn update(
        &mut self,
        spatial_query: &SpatialQuery,
        coordinates: HashSet<IVec2>,
        visible: &Query<(Entity, &GlobalTransform, &Collider, &Visible)>,
    ) {
        self.retain(|k, _| !coordinates.contains(k));
        let shape = Collider::rectangle(0.98, 0.98);
        for coords in &coordinates {
            let shape_pos = Vec2::new(coords.x as f32 + 0.5, coords.y as f32 + 0.5);
            let mut opacity = 0;
            let mut coord_entities = HashSet::new();
            spatial_query.shape_intersections_callback(
                &shape,
                shape_pos,
                0.,
                &default(),
                |entity| {
                    if let Ok((_, _, _, visible)) = visible.get(entity) {
                        coord_entities.insert(entity);
                        opacity = opacity.max(**visible);
                    }
                    true
                },
            );
            self.insert(*coords, (opacity, coord_entities));
        }
    }
}

fn update_opacity_map(
    mut opacity_map: ResMut<OpacityMap>,
    spatial_query: SpatialQuery,
    query: Query<
        (Entity, &GlobalTransform, &ColliderAabb, &Visible),
        Or<(
            Changed<GlobalTransform>,
            Changed<ColliderAabb>,
            Changed<Visible>,
        )>,
    >,
    visible: Query<(Entity, &GlobalTransform, &Collider, &Visible)>,
) {
    let mut to_update = HashSet::new();
    for (entity, transform, aabb, _) in &query {
        // println!(
        //     "Updating {entity} at {:?}",
        //     transform.translation().truncate().as_ivec2()
        // );
        let mut prev = HashSet::new();
        for (k, v) in opacity_map.iter() {
            if v.1.contains(&entity) {
                // println!("Also updating {k} because it contained entity");
                prev.insert(*k);
            }
        }
        let mut current = HashSet::new();
        current.insert(transform.translation().truncate().as_ivec2());
        for x in aabb.min.x as i32..aabb.max.x as i32 {
            for y in aabb.min.y as i32..aabb.max.y as i32 {
                // println!(
                //     "Also updating coordinates at {:?}",
                //     IVec2::new(x as i32, y as i32)
                // );
                current.insert(IVec2::new(x, y));
            }
        }
        if prev != current {
            to_update.extend(prev);
            to_update.extend(current);
        }
    }
    opacity_map.update(&spatial_query, to_update, &visible);
}

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

#[derive(Event)]
pub enum VisibilityChanged {
    Gained(Entity),
    Lost(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<MapData: 'static + Clone + Default + Send + Sync>(
    mut commands: Commands,
    map_config: Res<MapPlugin<MapData>>,
    query: Query<(Entity, &Map<MapData>), (Added<Map<MapData>>, Without<RevealedTiles>)>,
) {
    for (entity, map) in &query {
        let count = map.width * map.height;
        commands
            .entity(entity)
            .insert(RevealedTiles(vec![map_config.start_revealed; count]));
    }
}

fn viewshed_removed(
    mut commands: Commands,
    mut query: RemovedComponents<Viewshed>,
    visible_entities: Query<&VisibleEntities>,
) {
    for entity in query.read() {
        if let Ok(visible) = visible_entities.get(entity) {
            for e in visible.iter() {
                commands.trigger_targets(VisibilityChanged::Lost(*e), entity);
            }
        }
    }
}

struct VisibilityGrid(Entity, IVec2, u32, OpacityMap);

impl InputGrid for VisibilityGrid {
    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 {
        // println!("Checking {coord:?}");
        let coord = IVec2::new(coord.x, coord.y);
        if coord.distance(&self.1) > self.2 as f32 {
            return 0;
        }
        if let Some((opacity, entities)) = self.3.get(&coord) {
            if entities.contains(&self.0) {
                // println!("Hit viewer, 0");
                0
            } else {
                // println!("{opacity:?}");
                *opacity
            }
        } else {
            // println!("Miss, 0");
            0
        }
    }
}

fn update_viewshed(
    mut commands: Commands,
    physics_time: Res<Time<Physics>>,
    mut viewers: Query<(
        Entity,
        &mut Viewshed,
        &mut VisibleEntities,
        &GlobalTransform,
    )>,
    opacity_map: Res<OpacityMap>,
    sensors: Query<&Sensor>,
    spatial_query: SpatialQuery,
) {
    if physics_time.is_paused() {
        return;
    }
    for (viewer_entity, mut viewshed, mut visible_entities, viewer_transform) in &mut viewers {
        viewshed.update(
            &mut commands,
            &viewer_entity,
            &mut visible_entities,
            &viewer_transform.translation().truncate(),
            &opacity_map,
            &sensors,
            &spatial_query,
        );
    }
}

fn remove_visible(
    mut commands: Commands,
    mut removed: RemovedComponents<Visible>,
    mut viewers: Query<(
        Entity,
        &mut Viewshed,
        &mut VisibleEntities,
        &GlobalTransform,
    )>,
    sensors: Query<&Sensor>,
    spatial_query: SpatialQuery,
    mut opacity_map: ResMut<OpacityMap>,
    visible: Query<(Entity, &GlobalTransform, &Collider, &Visible)>,
) {
    if !removed.is_empty() {
        let mut to_update = HashSet::new();
        for e in removed.read() {
            for (k, v) in opacity_map.iter() {
                if v.1.contains(&e) {
                    to_update.insert(*k);
                }
            }
        }
        opacity_map.update(&spatial_query, to_update, &visible);
        for removed in removed.read() {
            for (viewer_entity, mut viewshed, mut visible_entities, start) in &mut viewers {
                if !visible_entities.contains(&removed) {
                    continue;
                }
                viewshed.update(
                    &mut commands,
                    &viewer_entity,
                    &mut visible_entities,
                    &start.translation().truncate(),
                    &opacity_map,
                    &sensors,
                    &spatial_query,
                );
            }
        }
    }
}

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 {
        for (map, mut revealed_tiles) in &mut map {
            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(
    trigger: Trigger<VisibilityChanged>,
    viewers: Query<(&GlobalTransform, Option<&Collider>), (With<Player>, With<Viewshed>)>,
    visible: Query<(&Name, &GlobalTransform, Option<&Collider>), Without<DontLogWhenVisible>>,
    mut log: Query<&mut Log>,
) {
    if let VisibilityChanged::Gained(entity) = trigger.event() {
        if *entity == trigger.entity() {
            return;
        }
        if let Ok((viewer_transform, viewer_collider)) = viewers.get(trigger.entity()) {
            if let Ok((name, viewed_transform, viewed_collider)) = visible.get(*entity) {
                let location = if let (Some(viewer_collider), Some(viewed_collider)) =
                    (viewer_collider, viewed_collider)
                {
                    viewer_transform.collider_direction_and_distance(
                        viewer_collider,
                        viewed_transform,
                        viewed_collider,
                    )
                } else {
                    let yaw = viewer_transform.yaw();
                    viewer_transform.direction_and_distance(viewed_transform, Some(yaw))
                };
                if let Ok(mut log) = log.get_single_mut() {
                    log.push(format!("{}: {location}", name));
                }
            }
        }
    }
}

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

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

impl<MapData: 'static + Clone + Default + Send + Sync> Plugin for VisibilityPlugin<MapData> {
    fn build(&self, app: &mut App) {
        app.init_resource::<OpacityMap>()
            .add_systems(FixedPreUpdate, update_opacity_map.before(update_viewshed))
            .add_systems(
                FixedPreUpdate,
                (
                    add_visibility_indices::<MapData>,
                    update_viewshed,
                    update_revealed_tiles::<MapData>,
                ),
            )
            .add_systems(FixedPostUpdate, (viewshed_removed, remove_visible))
            .add_observer(log_visible);
    }
}