traaains/sim/track/segment/commands/

remove.rs

use bevy::{
    ecs::{entity::Entity, system::Command, world::World},
    log::tracing,
    state::state::{NextState, State},
};

use crate::{
    sim::{
        logical_graph::LogicalBlock,
        track::{
            node::{RemoveConnectionFor, TrackNode},
            segment::{TrackSegment, signals_manager::SignalsManager},
        },
    },
    ui::UiState,
};

/// Command which removes a track segment along with taking care of removing node and logical connections.
pub struct RemoveSegment {
    /// Target segment to remove
    target: Entity,
    /// Whether to despawn nodes. This should almost always be `true` unless one is splitting segments.
    despawn_nodes: bool,
    /// Whether to despawn block. This should almost always be `true` unless one is splitting segments.
    despawn_block: bool,
    /// Whether to despawn associated signals. This should almost always be `true` unless one is splitting segments.
    despawn_signals: bool,
}

impl RemoveSegment {
    pub fn new(target: Entity) -> Self {
        Self {
            target,
            despawn_nodes: true,
            despawn_block: true,
            despawn_signals: true,
        }
    }

    pub fn with_despawn_nodes(mut self, despawn_nodes: bool) -> Self {
        self.despawn_nodes = despawn_nodes;
        self
    }

    pub fn with_despawn_block(mut self, despawn_block: bool) -> Self {
        self.despawn_block = despawn_block;
        self
    }

    pub fn with_despawn_signals(mut self, despawn_signals: bool) -> Self {
        self.despawn_signals = despawn_signals;
        self
    }

    /// Decreases number of segments for the given block and despawns it if the number is zero.
    /// Always returns false if RemoveSegment's despawn_block is set to false.
    fn decrease_block_segments(world: &mut World, block: Entity, despawn: bool) -> bool {
        let mut lb = world
            .query::<&mut LogicalBlock>()
            .get_mut(world, block)
            .unwrap();

        lb.number_of_segments -= 1;

        if despawn && lb.number_of_segments == 0 {
            tracing::debug!(
                "Despawning block {} due to having no attached segments",
                block
            );
            world.despawn(block);
            true
        } else {
            false
        }
    }
}

impl Command for RemoveSegment {
    type Out = Result<TrackSegment, &'static str>;

    fn apply(self, world: &mut World) -> Self::Out {
        // despawn UI if it belongs to our segment
        let ui_state = world.resource::<State<UiState>>();
        if let UiState::SegmentWindow { target } = ui_state.get()
            && *target == self.target
        {
            world
                .resource_mut::<NextState<UiState>>()
                .set(UiState::Default);
        }

        let Some(ts) = world.entity_mut(self.target).take::<TrackSegment>() else {
            panic!("RemoveSegment applied on a non-track segment entity!");
        };

        let block_despawned =
            RemoveSegment::decrease_block_segments(world, ts.block, self.despawn_block);

        for node in [ts.from_node, ts.to_node] {
            if RemoveConnectionFor::new(node, self.target)
                .with_despawn(self.despawn_nodes)
                .apply(world)
                == Ok(true)
            {
                continue;
            }

            let Ok(tn) = world.query::<&mut TrackNode>().get(world, node) else {
                panic!("Applied RemoveSegment on a non-track node entity!")
            };

            let local_segments: Vec<_> = tn.connections.iter().map(|c| c.segment).collect();
            // true if this node contains more segments from the same block
            let same_local_segments = local_segments
                .into_iter()
                .filter(|s| *s != self.target)
                .map(|s| world.query::<&TrackSegment>().get(world, s).unwrap().block)
                .any(|b| b == ts.block);

            if !block_despawned && !same_local_segments {
                let mut lb = world
                    .query::<&mut LogicalBlock>()
                    .get_mut(world, ts.block)
                    .unwrap();

                // connections are ordered by node first and entities are just u64
                // get all connections from our block through `from_node` and delete them
                let node_logical_conns: Vec<_> = lb
                    .connections
                    .range_mut(
                        &(node, Entity::from_bits(1)).into()
                            ..&(node, Entity::from_bits(u64::MAX)).into(),
                    )
                    .map(|(k, _)| k.clone())
                    .collect();

                for key in node_logical_conns {
                    lb.connections.remove(&key);
                }
            }

            // get logical blocks on either nodes and try to remove their connections keyed by (node, self.target)
            let node_segments: Vec<_> = world
                .query::<&TrackNode>()
                .get(world, node)
                .unwrap()
                .connections
                .iter()
                .map(|c| c.segment)
                .collect();

            let node_blocks: Vec<_> = node_segments
                .into_iter()
                .map(|segment_id| {
                    world
                        .query::<&TrackSegment>()
                        .get(world, segment_id)
                        .unwrap()
                        .block
                })
                .collect();

            for block in node_blocks {
                let mut block = world
                    .query::<&mut LogicalBlock>()
                    .get_mut(world, block)
                    .unwrap();

                block.connections.remove(&(node, self.target).into());
            }
        }

        if self.despawn_signals {
            let signals_manager = world
                .entity_mut(self.target)
                .take::<SignalsManager>()
                .unwrap();

            for (_, signal) in signals_manager.inner {
                world.despawn(signal);
            }
        }

        // finally, despawn the target segment
        world.despawn(self.target);

        Ok(ts)
    }
}