traaains/editor/train/placing/

mod.rs

use bevy::{
    app::{App, Plugin, Update},
    camera::visibility::Visibility,
    ecs::{
        component::Component,
        entity::Entity,
        query::{With, Without},
        resource::Resource,
        schedule::IntoScheduleConfigs,
        system::{Commands, If, Query, Res, ResMut, Single},
    },
    input::{
        ButtonInput, common_conditions::input_just_pressed, keyboard::KeyCode, mouse::MouseButton,
    },
    log::tracing,
    math::{Quat, Vec2},
    state::{
        app::AppExtStates,
        condition::in_state,
        state::{NextState, OnEnter, OnExit, State},
    },
    transform::components::Transform,
    ui::widget::Text,
};
use shapes::{Circle, Intersect, lerp::Lerp};

use crate::{
    editor::{
        StateText,
        train::{
            TrainEditorState,
            placing::states::{
                color_virtual_carriages_green, color_virtual_carriages_red, on_enter_placing_train,
                on_exit_placing_train,
                placed_state::{PlacedTrainState, place_train},
            },
        },
    },
    sim::{
        track::{
            node::{ConnectionLocation, TrackNode},
            pos::TrackPos,
            segment::TrackSegment,
        },
        train::{COUPLER_LENGTH, CarriageBogies, DOUBLE_BOGIE_IDX, SINGLE_BOGIE_IDX},
    },
    utils::Dir,
};

use self::{
    states::hovering::{HoveringTrack, IsHoveringTrack, on_enter_hovering, on_exit_hovering},
    templates::CarriageTemplate,
};

mod observers;
mod states;
pub mod templates;

#[derive(Resource)]
struct CurrentPlacingDir(Dir);

/// Stores the template of the currently placing train. Used to create the virtual pieces and later the actually placed ones.
#[derive(Clone, Resource)]
pub struct CurrentPlacingTrain {
    pub carriages: Vec<CarriageTemplate>,
}

// Default train when activating this editor
impl Default for CurrentPlacingTrain {
    fn default() -> Self {
        // random default train
        Self {
            carriages: vec![CarriageTemplate::default(); 2],
        }
    }
}

/// Marker component grouping all of the virtual pieces of the trains together. The virtual pieces store only the information
/// needed to provide visual feedback to the user, all other information is stored in [CurrentPlacingTrain].
#[derive(Component, Default)]
struct Virtual;

#[derive(Component)]
#[require(Virtual)]
struct VirtualTrain {
    carriages: Vec<Entity>,
}

#[derive(Component)]
#[require(Virtual)]
struct VirtualCarriage {
    length: f32,
    bogies: CarriageBogies,
}

#[derive(Component)]
#[require(Virtual)]
struct VirtualBogie;

/// Handles switching the direction of the train.
fn input_handler(input: Res<ButtonInput<KeyCode>>, mut placing_dir: ResMut<CurrentPlacingDir>) {
    if input.just_pressed(KeyCode::KeyR) {
        match placing_dir.0 {
            Dir::Positive => placing_dir.0 = Dir::Negative,
            Dir::Negative => placing_dir.0 = Dir::Positive,
        }
    }
}

fn get_intersection(ints: impl Iterator<Item = (f32, Vec2)>, dir: Dir) -> Option<(f32, Vec2)> {
    match dir {
        Dir::Positive => ints.min_by(|a, b| a.0.partial_cmp(&b.0).unwrap()),
        Dir::Negative => ints.max_by(|a, b| a.0.partial_cmp(&b.0).unwrap()),
    }
}

fn get_filtered_intersection(
    ints: impl Iterator<Item = (f32, Vec2)>,
    dir: Dir,
    curr_idx: f32,
) -> Option<(f32, Vec2)> {
    get_intersection(
        ints.filter(|(int_idx, _)| match dir {
            Dir::Positive => *int_idx > curr_idx,
            Dir::Negative => *int_idx < curr_idx,
        }),
        dir,
    )
}

/// Function which tries to move n meters in global space along active connections.
fn move_n_meters_along_active(
    n_meters: f32,
    curr_idx: f32,
    curr_segment_entity: Entity,
    curr_dir: Dir,
    q_segments: &Query<(Entity, &TrackSegment, &Transform)>,
    q_nodes: &Query<(Entity, &TrackNode)>,
) -> Option<(Entity, f32, Dir)> {
    let curr_segment = q_segments.get(curr_segment_entity).unwrap();
    let curr_curve_point = curr_segment.1.shape.lerp(curr_idx);
    let offset_circle = Circle::new(curr_curve_point, n_meters);

    tracing::debug!(
        "moving from segment {} and idx {}",
        curr_segment.0,
        curr_idx
    );
    // active segment refers to the segment we are currently trying to intersect to find a viable position
    let mut active_segment = Some(curr_segment);
    let mut active_dir = Some(curr_dir);

    while let Some(r#as) = active_segment
        && let Some(ad) = active_dir
    {
        // get vector by which we want to translate everything
        let curr_offset = r#as.2.translation.truncate() - curr_segment.2.translation.truncate();
        // curr segment is treated as world origin for the moment, which means that we need to translate the active shape
        let mut active_shape = r#as.1.shape.clone();
        active_shape.translate_by(curr_offset);

        // the segment isn't always a straight line and thus all of the possible points where the bogie mount can be is a circle
        let intersections = active_shape.intersections(&offset_circle);
        tracing::debug!(
            "intersections (with segment {} {:?}, circle {:?}): {:?}",
            r#as.0,
            active_shape,
            offset_circle,
            intersections
        );

        let intersection = intersections.and_then(|ints| {
            if r#as.0 == curr_segment.0 {
                return get_filtered_intersection(ints.into_iter(), curr_dir, curr_idx);
            }

            get_intersection(ints.into_iter(), curr_dir)
        });

        match intersection {
            // we have found an intersection and thus can happily place the bogie
            Some((int_idx, _)) => return Some((r#as.0, int_idx, ad)),
            // no intersection was found, try to get next active segment and try again
            None => {
                tracing::debug!("no intersection found, trying again");
                let ac = q_nodes
                    .get(match curr_dir {
                        Dir::Positive => r#as.1.to_node,
                        Dir::Negative => r#as.1.from_node,
                    })
                    .unwrap()
                    .1
                    .connections
                    .get_active_for(r#as.0);

                active_segment = ac.map(|ac| q_segments.get(ac.segment).unwrap());
                active_dir = ac.map(|ac| match ac.location {
                    ConnectionLocation::Start => Dir::Positive,
                    ConnectionLocation::End => Dir::Negative,
                });
            }
        }
    }

    tracing::debug!("no active connection found, cannot continue");
    None
}

/// Tries to find and calculate `n` (length of `spacings`) points along active track starting at segment and idx, moving along dir.
fn get_n_points_along_active_track(
    spacings: &[f32],
    segment: Entity,
    idx: f32,
    dir: Dir,
    q_segments: &Query<(Entity, &TrackSegment, &Transform)>,
    q_nodes: &Query<(Entity, &TrackNode)>,
) -> Option<Vec<(Entity, f32, Dir, Vec2)>> {
    let mut curr_segment = q_segments.get(segment).unwrap();
    let mut curr_dir = dir;
    let mut curr_idx = idx;

    let mut points = Vec::with_capacity(spacings.len());
    points.push((segment, idx, dir, curr_segment.1.shape.lerp(idx)));

    for spacing in spacings {
        match move_n_meters_along_active(
            *spacing,
            curr_idx,
            curr_segment.0,
            curr_dir,
            q_segments,
            q_nodes,
        ) {
            Some((final_segment_ent, final_idx, final_dir)) => {
                curr_segment = q_segments.get(final_segment_ent).unwrap();
                curr_idx = final_idx;
                curr_dir = final_dir;

                points.push((
                    final_segment_ent,
                    curr_idx,
                    curr_dir,
                    curr_segment.1.shape.lerp(curr_idx),
                ));
            }
            // this whole thing failed, need to find all points, exit
            None => return None,
        }
    }

    Some(points)
}

/// System which updates virtual train's bogies positions according to head train position
fn update_placing_train(
    mut cmd: Commands,
    q_train: Single<(&VirtualTrain, &TrackPos), Without<VirtualBogie>>,
    q_carriages: Query<&VirtualCarriage>,
    mut q_bogies: Query<(Entity, &mut TrackPos), With<VirtualBogie>>,
    q_segments: Query<(Entity, &TrackSegment, &Transform)>,
    q_nodes: Query<(Entity, &TrackNode)>,
    mut next_state: ResMut<NextState<PlacedTrainState>>,
) {
    // curr_segment and friends refer to the segment where the last bogie was placed and from which we are trying to move
    let curr_segment = q_segments.get(q_train.1.segment).unwrap();
    let curr_dir = match q_train.1.dir.as_ref().unwrap() {
        Dir::Positive => Dir::Negative,
        Dir::Negative => Dir::Positive,
    };
    let curr_idx = q_train.1.idx;

    let mut spacings = Vec::new();
    for (i, carriage) in q_train.0.carriages.iter().enumerate() {
        let carriage = q_carriages.get(*carriage).unwrap();

        match carriage.bogies {
            CarriageBogies::Single(_) => {
                spacings.push(carriage.length * SINGLE_BOGIE_IDX);
                spacings.push(carriage.length * SINGLE_BOGIE_IDX);
            }
            CarriageBogies::Double { front: _, back: _ } => {
                spacings.push(carriage.length * DOUBLE_BOGIE_IDX.0);
                spacings.push(carriage.length * (DOUBLE_BOGIE_IDX.1 - DOUBLE_BOGIE_IDX.0));
                spacings.push(carriage.length * (1.0 - DOUBLE_BOGIE_IDX.1));
            }
        }

        if i != q_train.0.carriages.len() - 1 {
            spacings.push(COUPLER_LENGTH);
            spacings.push(COUPLER_LENGTH);
        }
    }

    match get_n_points_along_active_track(
        &spacings,
        curr_segment.0,
        curr_idx,
        curr_dir,
        &q_segments,
        &q_nodes,
    ) {
        Some(points) => {
            let mut points = points.into_iter();
            // starting at front of train, so move one
            points.next();

            for carriage_ent in q_train.0.carriages.iter() {
                let carriage = q_carriages.get(*carriage_ent).unwrap();

                match carriage.bogies {
                    CarriageBogies::Single(entity) => {
                        let (_, mut bogie_pos) = q_bogies.get_mut(entity).unwrap();

                        let target = points.next().unwrap();
                        bogie_pos.segment = target.0;
                        bogie_pos.idx = target.1;
                        bogie_pos.dir = Some(target.2.reverse());

                        // move to back of train, then two couplers to prepare next carriage
                        points.next();
                        points.next();
                        points.next();
                    }
                    CarriageBogies::Double { front, back } => {
                        let bogies = q_bogies.get_many_mut([front, back]).unwrap();

                        for mut bogie in bogies {
                            let target = points.next().unwrap();
                            bogie.1.segment = target.0;
                            bogie.1.idx = target.1;
                            // when traversing the connections, we go towards the back of the train,
                            // but the bogies should be facing towards the front
                            bogie.1.dir = Some(target.2.reverse());
                        }

                        // move to back of train, then two couplers to prepare next carriage
                        points.next();
                        points.next();
                        points.next();
                    }
                }

                cmd.entity(*carriage_ent).insert(Visibility::Visible);
            }

            next_state.set(PlacedTrainState::Valid);
        }
        None => {
            for carriage in q_train.0.carriages.iter() {
                cmd.entity(*carriage).insert(Visibility::Hidden);
            }

            next_state.set(PlacedTrainState::Invalid);
        }
    }
}

fn move_virtual_carriages_to_bogies(
    q_carriages: Query<(&VirtualCarriage, &mut Transform), Without<VirtualBogie>>,
    q_bogies: Query<&Transform, With<VirtualBogie>>,
) {
    for (carriage, mut transform) in q_carriages {
        match &carriage.bogies {
            CarriageBogies::Single(bogie) => {
                *transform = *q_bogies.get(*bogie).unwrap();
            }
            CarriageBogies::Double { front, back } => {
                let front_trans = *q_bogies.get(*front).unwrap();
                let back_trans = *q_bogies.get(*back).unwrap();

                transform.translation = (front_trans.translation + back_trans.translation) / 2.0;
                let dir = front_trans.translation.truncate() - back_trans.translation.truncate();
                transform.rotation = Quat::from_rotation_z(f32::atan2(dir.y, dir.x));
            }
        }
    }
}

fn update_state_text(
    mut text: Single<&mut Text, With<StateText>>,
    current_dir: If<Res<CurrentPlacingDir>>,
    valid: Option<Res<State<PlacedTrainState>>>,
) {
    text.0 = format!(
        "Placing train\n[Esc] to exit\nCurrent direction: {}\nCurrent state: {}",
        match current_dir.0.0 {
            Dir::Positive => "Positive",
            Dir::Negative => "Negative",
        },
        valid
            .map(|v| v.get().to_string())
            .unwrap_or(String::from("-"))
    );
}

fn register_state_handlers(app: &mut App) {
    app.add_sub_state::<HoveringTrack>()
        .add_computed_state::<IsHoveringTrack>()
        .add_sub_state::<PlacedTrainState>()
        .add_systems(
            OnEnter(TrainEditorState::PlacingTrain),
            on_enter_placing_train,
        )
        .add_systems(
            OnExit(TrainEditorState::PlacingTrain),
            on_exit_placing_train,
        )
        .add_systems(OnEnter(IsHoveringTrack), on_enter_hovering)
        .add_systems(OnExit(IsHoveringTrack), on_exit_hovering)
        .add_systems(
            OnEnter(PlacedTrainState::Valid),
            color_virtual_carriages_green,
        )
        .add_systems(
            OnEnter(PlacedTrainState::Invalid),
            color_virtual_carriages_red,
        );
}

pub struct PlacingTrainPlugin;

impl Plugin for PlacingTrainPlugin {
    fn build(&self, app: &mut App) {
        register_state_handlers(app);

        app.add_systems(
            Update,
            (
                input_handler,
                update_state_text,
                update_placing_train,
                move_virtual_carriages_to_bogies,
            )
                .run_if(in_state(TrainEditorState::PlacingTrain)),
        )
        .add_systems(
            Update,
            place_train
                .after(update_placing_train)
                .run_if(in_state(PlacedTrainState::Valid))
                .run_if(input_just_pressed(MouseButton::Left)),
        );
    }
}