rmf_site_editor 0.0.3

/*
 * Copyright (C) 2022 Open Source Robotics Foundation
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
*/

use crate::{recency::RecencyRanking, site::*, WorkspaceMarker};
use bevy::{
    ecs::{
        hierarchy::ChildOf,
        system::{BoxedSystem, SystemParam, SystemState},
    },
    prelude::*,
};
use rmf_site_format::legacy::{building_map::BuildingMap, PortingError};
use smallvec::SmallVec;
use std::{
    collections::{HashMap, HashSet},
    error::Error,
    path::PathBuf,
    sync::Arc,
};
use thiserror::Error as ThisError;

/// This component is given to the site to keep track of what file it should be
/// saved to by default.
#[derive(Component, Clone, Debug, Deref)]
pub struct DefaultFile(pub PathBuf);

#[derive(Event, Clone)]
pub struct LoadSite {
    /// The site data to load
    pub site: rmf_site_format::Site,
    /// Should the application switch focus to this new site
    pub focus: bool,
    /// The default file path that should be assigned to the site
    pub default_file: Option<PathBuf>,
}

impl LoadSite {
    #[allow(non_snake_case)]
    pub fn blank_L1(name: String, default_file: Option<PathBuf>) -> Self {
        Self {
            site: rmf_site_format::Site::blank_L1(name),
            default_file,
            focus: true,
        }
    }

    /// Create a `LoadSite` instance from raw data and optionally a file name.
    ///
    /// Note that this function can take some time to run if the file data is
    /// large, so it's best to use this in an async context.
    pub fn from_data(data: &Vec<u8>, default_file: Option<PathBuf>) -> Result<Self, LoadSiteError> {
        if let Some(path) = &default_file {
            let Some(filename) = path.file_name().and_then(|f| f.to_str()) else {
                return Err(LoadSiteError::IncompatibleFilename(path.clone()));
            };

            // If the default file is specified, we should only try to parse
            // based on formats that match the name, if it's possible to identify
            // one.
            let site = if filename.ends_with(".building.yaml") {
                match BuildingMap::from_bytes(data) {
                    Ok(building) => building
                        .to_site()
                        .map_err(LoadSiteError::LegacyConversion)?,
                    Err(err) => {
                        return Err(LoadSiteError::CorruptedBuildingFile {
                            path: path.clone(),
                            err,
                        });
                    }
                }
            } else if filename.ends_with(".json") {
                Site::from_bytes_json(data)?
            } else {
                return Err(LoadSiteError::UnrecognizedFileType(path.clone()));
            };

            return Ok(Self {
                site,
                focus: false,
                default_file,
            });
        }

        // No file type was indicated, so try parsing the data with each option
        // in order of how likely it will be used
        let site = Site::from_bytes_json(data)
            .map_err(|_| LoadSiteError::UnknownDataFormat)
            .or_else(|_| {
                BuildingMap::from_bytes(data)
                    .map_err(|_| LoadSiteError::UnknownDataFormat)
                    .and_then(|building| {
                        building
                            .to_site()
                            .map_err(|_| LoadSiteError::UnknownDataFormat)
                    })
            })?;

        Ok(Self {
            site,
            focus: false,
            default_file,
        })
    }
}

#[derive(ThisError, Debug)]
pub enum LoadSiteError {
    #[error("Trying to load a site with an incompatible filename: {0}")]
    IncompatibleFilename(PathBuf),
    #[error("Failed to parse legacy building file named [{path}]: {err}")]
    CorruptedBuildingFile {
        path: PathBuf,
        err: serde_yaml::Error,
    },
    #[error("Failed to convert a legacy building into a site: {0}")]
    LegacyConversion(#[from] PortingError),
    #[error("Failed parsing ron site file: {0}")]
    RonParsingError(Box<dyn std::error::Error + Send + Sync + 'static>),
    #[error("Failed parsing json site file: {0}")]
    JsonParsingError(#[from] serde_json::Error),
    #[error("Unrecognized file type: {0}")]
    UnrecognizedFileType(PathBuf),
    #[error(
        "Cannot determine data format for raw data. It could not be parsed as .building.yaml or .site.json"
    )]
    UnknownDataFormat,
}

trait LoadResult<T> {
    fn as_broken_error(self, site: Entity, cause: &'static str) -> Result<T, SiteLoadingError>;
}

impl<T> LoadResult<T> for Result<T, u32> {
    fn as_broken_error(self, site: Entity, cause: &'static str) -> Result<T, SiteLoadingError> {
        self.map_err(|broken| SiteLoadingError {
            site,
            error: BrokenEntityError::new(broken, cause).into(),
        })
    }
}

pub type LoadSiteResult = Result<LoadSite, LoadSiteError>;

struct SiteLoadingError {
    site: Entity,
    error: LoadingError,
}

/// Errors that can occur while loading a site.
#[derive(ThisError, Debug)]
enum LoadingError {
    #[error(transparent)]
    BrokenEntity(#[from] BrokenEntityError),
    #[error("Your site editor application is missing a mandatory extension: {0}")]
    MissingMandatoryExtension(Arc<str>),
    #[error("Extension [{extension}] encountered an error: {error}")]
    ExtensionError {
        extension: Arc<str>,
        error: Arc<dyn Error>,
    },
}

#[derive(ThisError, Debug)]
#[error("The site has a broken internal reference: {broken} for {cause}")]
struct BrokenEntityError {
    broken: u32,
    cause: &'static str,
    // TODO(@mxgrey): reintroduce Backtrack when it's supported on stable
    // backtrace: Backtrace,
}

impl BrokenEntityError {
    fn new(broken: u32, cause: &'static str) -> Self {
        Self { broken, cause }
    }
}

pub struct LoadingArgs {
    /// The entity of the site that is being loaded
    pub site: Entity,
    /// The data for the extension that is being loaded
    pub data: serde_json::Value,
}

pub type LoadingResult<E> = Result<(), E>;

pub(crate) type LoadingSystem = BoxedSystem<In<LoadingArgs>, Result<(), Arc<dyn Error>>>;

fn load_single_site(
    world: &mut World,
    entity_generation_params_state: &mut SystemState<EntityGenerationParams>,
    site: &rmf_site_format::Site,
) -> Result<Entity, SiteLoadingError> {
    let mut entity_generation_params = entity_generation_params_state.get_mut(world);
    let r = generate_site_entities(
        &mut entity_generation_params.commands,
        &mut entity_generation_params.model_loader,
        site,
    );
    entity_generation_params_state.apply(world);
    let site_entity = r?;

    world
        .resource_scope::<ExtensionHooks, Result<(), LoadingError>>(|world, mut hooks| {
            for (extension, data) in &site.extensions.data {
                let prevent_loading_on_error = site
                    .properties
                    .extension_settings
                    .get(extension)
                    .map(|settings| settings.prevent_loading_on_error)
                    .unwrap_or(false);

                let hook = match hooks.hooks.get_mut(extension) {
                    Some(hook) => hook,
                    None => {
                        if prevent_loading_on_error {
                            return Err(LoadingError::MissingMandatoryExtension(Arc::clone(
                                extension,
                            )));
                        } else {
                            warn!(
                                "Missing extension {extension} which was expected by the site data"
                            );
                            continue;
                        }
                    }
                };

                if let Some(loading) = &mut hook.loading {
                    let args = LoadingArgs {
                        site: site_entity,
                        data: data.clone(),
                    };
                    let r = loading.run(args, world);
                    loading.apply_deferred(world);

                    if let Err(err) = r {
                        if prevent_loading_on_error {
                            return Err(LoadingError::ExtensionError {
                                extension: Arc::clone(extension),
                                error: err,
                            });
                        } else {
                            warn!(
                                "Error occurred while extension [{extension}] was loading: {err}"
                            );
                        }
                    }
                }
            }
            Ok(())
        })
        .map_err(|error| SiteLoadingError {
            site: site_entity,
            error,
        })?;

    Ok(site_entity)
}

fn generate_site_entities(
    commands: &mut Commands,
    model_loader: &mut ModelLoader,
    site_data: &rmf_site_format::Site,
) -> Result<Entity, SiteLoadingError> {
    let mut id_to_entity = HashMap::new();
    let mut highest_id = 0_u32;
    let mut consider_id = |consider| {
        if consider > highest_id {
            highest_id = consider;
        }
    };

    let site_id = commands
        .spawn((Transform::IDENTITY, Visibility::Hidden))
        .insert(Category::Site)
        .insert(WorkspaceMarker)
        .id();

    commands
        .spawn(InfiniteGridBundle {
            transform: Transform {
                translation: Vec3::new(0., 0., -0.01),
                rotation: Quat::from_rotation_x(90_f32.to_radians()),
                scale: Vec3::splat(1.0),
            },
            settings: InfiniteGridSettings {
                minor_line_color: Color::srgb(0.2, 0.2, 0.2),
                major_line_color: Color::srgb(0.4, 0.4, 0.4),
                ..Default::default()
            },
            ..Default::default()
        })
        .insert(ChildOf(site_id));

    for (anchor_id, anchor) in &site_data.anchors {
        let anchor_entity = commands
            .spawn(AnchorBundle::new(anchor.clone()))
            .insert(SiteID(*anchor_id))
            .insert(ChildOf(site_id))
            .id();
        id_to_entity.insert(*anchor_id, anchor_entity);
        consider_id(*anchor_id);
    }

    for (group_id, group) in &site_data.fiducial_groups {
        let group_entity = commands
            .spawn(group.clone())
            .insert(SiteID(*group_id))
            .insert(ChildOf(site_id))
            .id();
        id_to_entity.insert(*group_id, group_entity);
        consider_id(*group_id);
    }

    for (group_id, group) in &site_data.textures {
        let group_entity = commands
            .spawn(group.clone())
            .insert(SiteID(*group_id))
            .insert(ChildOf(site_id))
            .id();
        id_to_entity.insert(*group_id, group_entity);
        consider_id(*group_id);
    }

    for (level_id, level_data) in &site_data.levels {
        let level_entity = commands
            .spawn(SiteID(*level_id))
            .insert(ChildOf(site_id))
            .id();

        for (anchor_id, anchor) in &level_data.anchors {
            let anchor_entity = commands
                .spawn(AnchorBundle::new(anchor.clone()))
                .insert(SiteID(*anchor_id))
                .insert(ChildOf(level_entity))
                .id();
            id_to_entity.insert(*anchor_id, anchor_entity);
            consider_id(*anchor_id);
        }

        for (door_id, door) in &level_data.doors {
            let door_entity = commands
                .spawn(
                    door.convert(&id_to_entity)
                        .as_broken_error(site_id, "door")?,
                )
                .insert(SiteID(*door_id))
                .insert(ChildOf(level_entity))
                .id();
            id_to_entity.insert(*door_id, door_entity);
            consider_id(*door_id);
        }

        for (drawing_id, drawing) in &level_data.drawings {
            let drawing_entity = commands
                .spawn(DrawingBundle::new(drawing.properties.clone()))
                .insert(SiteID(*drawing_id))
                .insert(ChildOf(level_entity))
                .id();

            for (anchor_id, anchor) in &drawing.anchors {
                let anchor_entity = commands
                    .spawn(AnchorBundle::new(anchor.clone()))
                    .insert(SiteID(*anchor_id))
                    .insert(ChildOf(drawing_entity))
                    .id();
                id_to_entity.insert(*anchor_id, anchor_entity);
                consider_id(*anchor_id);
            }

            for (fiducial_id, fiducial) in &drawing.fiducials {
                let fiducial_entity = commands
                    .spawn(
                        fiducial
                            .convert(&id_to_entity)
                            .as_broken_error(site_id, "drawing fiducial")?,
                    )
                    .insert(SiteID(*fiducial_id))
                    .insert(ChildOf(drawing_entity))
                    .id();
                id_to_entity.insert(*fiducial_id, fiducial_entity);
                consider_id(*fiducial_id);
            }

            for (measurement_id, measurement) in &drawing.measurements {
                let measurement_entity = commands
                    .spawn(
                        measurement
                            .convert(&id_to_entity)
                            .as_broken_error(site_id, "measurement")?,
                    )
                    .insert(SiteID(*measurement_id))
                    .insert(ChildOf(drawing_entity))
                    .id();
                id_to_entity.insert(*measurement_id, measurement_entity);
                consider_id(*measurement_id);
            }

            consider_id(*drawing_id);
        }

        for (floor_id, floor) in &level_data.floors {
            commands
                .spawn(
                    floor
                        .convert(&id_to_entity)
                        .as_broken_error(site_id, "floor")?,
                )
                .insert(SiteID(*floor_id))
                .insert(ChildOf(level_entity));
            consider_id(*floor_id);
        }

        for (wall_id, wall) in &level_data.walls {
            commands
                .spawn(
                    wall.convert(&id_to_entity)
                        .as_broken_error(site_id, "wall")?,
                )
                .insert(SiteID(*wall_id))
                .insert(ChildOf(level_entity));
            consider_id(*wall_id);
        }

        commands
            .entity(level_entity)
            .insert((Transform::IDENTITY, Visibility::Hidden))
            .insert(level_data.properties.clone())
            .insert(Category::Level)
            .with_children(|level| {
                // These don't need a return value so can be wrapped in a with_children
                for (light_id, light) in &level_data.lights {
                    level.spawn(light.clone()).insert(SiteID(*light_id));
                    consider_id(*light_id);
                }

                for (physical_camera_id, physical_camera) in &level_data.physical_cameras {
                    level
                        .spawn(physical_camera.clone())
                        .insert(SiteID(*physical_camera_id));
                    consider_id(*physical_camera_id);
                }

                for (camera_pose_id, camera_pose) in &level_data.user_camera_poses {
                    level
                        .spawn(camera_pose.clone())
                        .insert(SiteID(*camera_pose_id));
                    consider_id(*camera_pose_id);
                }
            });

        // TODO(MXG): Log when a RecencyRanking fails to load correctly.
        commands
            .entity(level_entity)
            .insert(
                RecencyRanking::<FloorMarker>::from_u32(&level_data.rankings.floors, &id_to_entity)
                    .unwrap_or(RecencyRanking::new()),
            )
            .insert(
                RecencyRanking::<DrawingMarker>::from_u32(
                    &level_data.rankings.drawings,
                    &id_to_entity,
                )
                .unwrap_or(RecencyRanking::new()),
            );
        id_to_entity.insert(*level_id, level_entity);
        consider_id(*level_id);
    }

    for (lift_id, lift_data) in &site_data.lifts {
        let lift_entity = commands
            .spawn(SiteID(*lift_id))
            .insert(ChildOf(site_id))
            .id();

        commands.entity(lift_entity).with_children(|lift| {
            lift.spawn((Transform::default(), Visibility::default()))
                .insert(CabinAnchorGroupBundle::default())
                .with_children(|anchor_group| {
                    for (anchor_id, anchor) in &lift_data.cabin_anchors {
                        let anchor_entity = anchor_group
                            .spawn(AnchorBundle::new(anchor.clone()))
                            .insert(SiteID(*anchor_id))
                            .id();
                        id_to_entity.insert(*anchor_id, anchor_entity);
                        consider_id(*anchor_id);
                    }
                });
        });

        for (door_id, door) in &lift_data.cabin_doors {
            let door_entity = commands
                .spawn(
                    door.convert(&id_to_entity)
                        .as_broken_error(site_id, "cabin door")?,
                )
                .insert(Dependents::single(lift_entity))
                .insert(SiteID(*door_id))
                .insert(ChildOf(lift_entity))
                .id();
            id_to_entity.insert(*door_id, door_entity);
            consider_id(*door_id);
        }

        commands.entity(lift_entity).insert(Category::Lift).insert(
            lift_data
                .properties
                .convert(&id_to_entity)
                .as_broken_error(site_id, "lift")?,
        );

        id_to_entity.insert(*lift_id, lift_entity);
        consider_id(*lift_id);
    }

    for (fiducial_id, fiducial) in &site_data.fiducials {
        let fiducial_entity = commands
            .spawn(
                fiducial
                    .convert(&id_to_entity)
                    .as_broken_error(site_id, "site fiducial")?,
            )
            .insert(SiteID(*fiducial_id))
            .insert(ChildOf(site_id))
            .id();
        id_to_entity.insert(*fiducial_id, fiducial_entity);
        consider_id(*fiducial_id);
    }

    for (group_id, group) in &site_data.navigation.guided.mutex_groups {
        let group_entity = commands
            .spawn(group.clone())
            .insert(SiteID(*group_id))
            .insert(ChildOf(site_id))
            .id();
        id_to_entity.insert(*group_id, group_entity);
        consider_id(*group_id);
    }

    for (nav_graph_id, nav_graph_data) in &site_data.navigation.guided.graphs {
        let nav_graph = commands
            .spawn((Transform::default(), Visibility::default()))
            .insert(nav_graph_data.clone())
            .insert(SiteID(*nav_graph_id))
            .insert(ChildOf(site_id))
            .id();
        id_to_entity.insert(*nav_graph_id, nav_graph);
        consider_id(*nav_graph_id);
    }

    for (lane_id, lane_data) in &site_data.navigation.guided.lanes {
        let lane = commands
            .spawn(
                lane_data
                    .convert(&id_to_entity)
                    .as_broken_error(site_id, "lane")?,
            )
            .insert(SiteID(*lane_id))
            .insert(ChildOf(site_id))
            .id();
        id_to_entity.insert(*lane_id, lane);
        consider_id(*lane_id);
    }

    for (location_id, location_data) in &site_data.navigation.guided.locations {
        let location = commands
            .spawn(
                location_data
                    .convert(&id_to_entity)
                    .as_broken_error(site_id, "location")?,
            )
            .insert(SiteID(*location_id))
            .insert(ChildOf(site_id))
            .id();
        id_to_entity.insert(*location_id, location);
        consider_id(*location_id);
    }
    // Properties require the id_to_entity map to be fully populated to load suppressed issues
    commands.entity(site_id).insert(
        site_data
            .properties
            .convert(&id_to_entity)
            .as_broken_error(site_id, "site properties")?,
    );

    let mut model_description_dependents = HashMap::<Entity, HashSet<Entity>>::new();
    let mut model_description_to_source = HashMap::<Entity, AssetSource>::new();
    for (model_description_id, model_description) in &site_data.model_descriptions {
        let model_description_entity = commands
            .spawn(model_description.clone())
            .insert(SiteID(*model_description_id))
            .insert(Category::ModelDescription)
            .insert(ChildOf(site_id))
            .id();
        id_to_entity.insert(*model_description_id, model_description_entity);
        consider_id(*model_description_id);
        model_description_dependents.insert(model_description_entity, HashSet::new());
        model_description_to_source
            .insert(model_description_entity, model_description.source.0.clone());
    }

    for (robot_id, robot_data) in &site_data.robots {
        // Robot IDs are pointing to model description entities
        if let Some(model_description_entity) = id_to_entity
            .get(robot_id)
            .filter(|e| model_description_to_source.contains_key(*e))
        {
            commands
                .entity(*model_description_entity)
                .insert(ModelProperty(robot_data.clone()));
        } else {
            // Robot is affiliated to a non-existent model description,
            // create a description entity for users to modify after loading
            commands
                .spawn(ModelDescriptionBundle::default())
                .insert(Category::ModelDescription)
                .insert(ModelProperty(robot_data.clone()))
                .insert(ChildOf(site_id));
            error!(
                "Robot {} with properties {:?} is pointing to a non-existent \
                model description! Assigning robot to the default model description \
                with an empty asset source.",
                robot_id, robot_data
            );
        };
    }

    for (model_instance_id, parented_model_instance) in &site_data.model_instances {
        let model_instance = parented_model_instance
            .bundle
            .convert(&id_to_entity)
            .as_broken_error(site_id, "model instance")?;

        // The parent id is invalid, we do not spawn this model instance and generate
        // an error instead
        let parent = id_to_entity
            .get(&parented_model_instance.parent)
            .ok_or_else(|| SiteLoadingError {
                site: site_id,
                error: BrokenEntityError::new(
                    parented_model_instance.parent,
                    "model instance parent",
                )
                .into(),
            })?;

        let model_instance_entity = model_loader
            .spawn_model_instance(*parent, model_instance.clone())
            .insert((Category::Model, SiteID(*model_instance_id)))
            .id();
        id_to_entity.insert(*model_instance_id, model_instance_entity);
        consider_id(*model_instance_id);

        if let Some(instances) = model_instance
            .description
            .0
            .map(|e| model_description_dependents.get_mut(&e))
            .flatten()
        {
            instances.insert(model_instance_entity);
        } else {
            error!(
                "Model description missing for instance {}. This should \
                not happen, please report this bug to the maintainers of \
                rmf_site_editor.",
                model_instance.name.0,
            );
        }
    }

    for (model_description_entity, dependents) in model_description_dependents {
        commands
            .entity(model_description_entity)
            .insert(Dependents(dependents));
    }

    for (task_id, task_data) in &site_data.tasks {
        let task_entity = commands
            .spawn(task_data.clone())
            .insert(SiteID(*task_id))
            .insert(Category::Task)
            .insert(ChildOf(site_id))
            .id();
        id_to_entity.insert(*task_id, task_entity);
        consider_id(*task_id);
    }

    for (scenario_id, scenario_data) in &site_data.scenarios {
        let parent = match scenario_data.properties.parent_scenario.0 {
            Some(parent_id) => *id_to_entity.get(&parent_id).unwrap_or(&site_id),
            None => site_id,
        };

        let scenario = scenario_data
            .convert(&id_to_entity)
            .as_broken_error(site_id, "scenario data")?;
        let scenario_entity = commands
            .spawn(scenario.properties.clone())
            .insert(SiteID(*scenario_id))
            .insert(ChildOf(parent))
            .id();
        id_to_entity.insert(*scenario_id, scenario_entity);
        consider_id(*scenario_id);

        // Spawn instance modifier entities
        for (instance_id, instance_modifier) in scenario_data.instances.iter() {
            let instance_entity = id_to_entity
                .get(instance_id)
                .ok_or(*instance_id)
                .as_broken_error(site_id, "instance modifier")?;

            if let Some(pose) = instance_modifier.pose {
                commands.trigger(UpdateModifier::modify(
                    scenario_entity,
                    *instance_entity,
                    pose,
                ));
            }
            if let Some(inclusion) = instance_modifier.inclusion {
                commands.trigger(UpdateModifier::modify(
                    scenario_entity,
                    *instance_entity,
                    inclusion,
                ));
            }
            if let Some(level_entity) = instance_modifier
                .on_level
                .and_then(|level_id| id_to_entity.get(&level_id))
            {
                commands.trigger(UpdateModifier::modify(
                    scenario_entity,
                    *instance_entity,
                    OnLevel(Some(*level_entity)),
                ));
            }
        }
        // Spawn task modifier entities
        for (task_id, task_modifier) in scenario_data.tasks.iter() {
            let task_entity = id_to_entity
                .get(task_id)
                .ok_or(*task_id)
                .as_broken_error(site_id, "task modifier")?;
            if let Some(inclusion) = task_modifier.inclusion {
                commands.trigger(UpdateModifier::modify(
                    scenario_entity,
                    *task_entity,
                    inclusion,
                ));
            }
            if let Some(params) = &task_modifier.params {
                commands.trigger(UpdateModifier::modify(
                    scenario_entity,
                    *task_entity,
                    params.clone(),
                ));
            }
        }
    }

    let nav_graph_rankings = match RecencyRanking::<NavGraphMarker>::from_u32(
        &site_data.navigation.guided.ranking,
        &id_to_entity,
    ) {
        Ok(r) => r,
        Err(id) => {
            error!(
                "ERROR: Nav Graph ranking could not load because a graph with \
                id {id} does not exist."
            );
            RecencyRanking::new()
        }
    };

    commands
        .entity(site_id)
        .insert(nav_graph_rankings)
        .insert(NextSiteID(highest_id + 1));

    // Make the lift cabin anchors that are used by doors subordinate
    for (lift_id, lift_data) in &site_data.lifts {
        for (_, door) in &lift_data.cabin_doors {
            for anchor in door.reference_anchors.array() {
                commands
                    .entity(
                        *id_to_entity
                            .get(&anchor)
                            .ok_or(anchor)
                            .as_broken_error(site_id, "lift door")?,
                    )
                    .insert(Subordinate(Some(
                        *id_to_entity
                            .get(lift_id)
                            .ok_or(*lift_id)
                            .as_broken_error(site_id, "lift")?,
                    )));
            }
        }
    }

    return Ok(site_id);
}

#[derive(SystemParam)]
pub struct EntityGenerationParams<'w, 's> {
    commands: Commands<'w, 's>,
    model_loader: ModelLoader<'w, 's>,
}

#[derive(SystemParam)]
pub struct SiteLoadingParams<'w, 's> {
    load_sites: EventReader<'w, 's, LoadSite>,
    change_current_site: EventWriter<'w, ChangeCurrentSite>,
}

pub fn load_site(
    world: &mut World,
    entity_generation_params_state: &mut SystemState<EntityGenerationParams>,
    loading_params_state: &mut SystemState<SiteLoadingParams>,
) {
    let mut loading_params = loading_params_state.get_mut(world);
    let sites_to_load: SmallVec<[_; 8]> = loading_params.load_sites.read().cloned().collect();
    for cmd in sites_to_load {
        let site = match load_single_site(world, entity_generation_params_state, &cmd.site) {
            Ok(site) => site,
            Err(err) => {
                world.entity_mut(err.site).despawn();
                error!(
                    "Failed to load the site entities because the file had an \
                    internal inconsistency:\n{}\n---\nSite Data:\n{:#?}",
                    err.error, &cmd.site,
                );
                continue;
            }
        };
        if let Some(path) = &cmd.default_file {
            world.entity_mut(site).insert(DefaultFile(path.clone()));
        }

        if cmd.focus {
            let mut loading_params = loading_params_state.get_mut(world);
            loading_params.change_current_site.write(ChangeCurrentSite {
                site,
                level: None,
                scenario: None,
            });
        }
    }

    loading_params_state.apply(world);
}

#[derive(ThisError, Debug, Clone)]
pub enum ImportNavGraphError {
    #[error("The site we are importing into has a broken reference")]
    BrokenSiteReference,
    #[error("The nav graph that is being imported has a broken reference inside of it")]
    BrokenInternalReference(u32),
    #[error("The existing site is missing a level name required by the nav graphs: {0}")]
    MissingLevelName(String),
    #[error("The existing site is missing a lift name required by the nav graphs: {0}")]
    MissingLiftName(String),
    #[error("The existing site has a lift without a cabin anchor group: {0}")]
    MissingCabinAnchorGroup(String),
}

#[derive(Event)]
pub struct ImportNavGraphs {
    pub into_site: Entity,
    pub from_site: rmf_site_format::Site,
}

#[derive(SystemParam)]
pub struct ImportNavGraphParams<'w, 's> {
    commands: Commands<'w, 's>,
    sites: Query<'w, 's, &'static Children, With<NameOfSite>>,
    levels: Query<
        'w,
        's,
        (
            Entity,
            &'static NameInSite,
            &'static ChildOf,
            &'static Children,
        ),
        With<LevelElevation>,
    >,
    lifts: Query<
        'w,
        's,
        (
            Entity,
            &'static NameInSite,
            &'static ChildOf,
            &'static Children,
        ),
        With<LiftCabin<Entity>>,
    >,
    cabin_anchor_groups: Query<'w, 's, &'static Children, With<CabinAnchorGroup>>,
    anchors: Query<'w, 's, (Entity, &'static Anchor)>,
}

fn generate_imported_nav_graphs(
    params: &mut ImportNavGraphParams,
    into_site: Entity,
    from_site_data: &rmf_site_format::Site,
) -> Result<(), ImportNavGraphError> {
    let site_children = match params.sites.get(into_site) {
        Ok(c) => c,
        _ => return Err(ImportNavGraphError::BrokenSiteReference),
    };

    let mut level_name_to_entity = HashMap::new();
    for (e, name, child_of, _) in &params.levels {
        if child_of.parent() != into_site {
            continue;
        }

        level_name_to_entity.insert(name.clone().0, e);
    }

    let mut lift_name_to_entity = HashMap::new();
    for (e, name, child_of, _) in &params.lifts {
        if child_of.parent() != into_site {
            continue;
        }

        lift_name_to_entity.insert(name.clone().0, e);
    }

    let mut id_to_entity = HashMap::new();
    for (level_id, level_data) in &from_site_data.levels {
        if let Some(e) = level_name_to_entity.get(&level_data.properties.name.0) {
            id_to_entity.insert(*level_id, *e);
        } else {
            return Err(ImportNavGraphError::MissingLevelName(
                level_data.properties.name.0.clone(),
            ));
        }
    }

    let mut lift_to_anchor_group = HashMap::new();
    for (lift_id, lift_data) in &from_site_data.lifts {
        if let Some(e) = lift_name_to_entity.get(&lift_data.properties.name.0) {
            id_to_entity.insert(*lift_id, *e);
            if let Some(e_group) = params
                .lifts
                .get(*e)
                .unwrap()
                .3
                .iter()
                .find(|child| params.cabin_anchor_groups.contains(*child))
            {
                lift_to_anchor_group.insert(*e, e_group);
            } else {
                return Err(ImportNavGraphError::MissingCabinAnchorGroup(
                    lift_data.properties.name.0.clone(),
                ));
            }
        } else {
            return Err(ImportNavGraphError::MissingLiftName(
                lift_data.properties.name.0.clone(),
            ));
        }
    }

    let anchor_close_enough = 0.05;
    for (lift_id, lift_data) in &from_site_data.lifts {
        let lift_e = *id_to_entity.get(lift_id).unwrap();
        let anchor_group = *lift_to_anchor_group.get(&lift_e).unwrap();
        let existing_lift_anchors: Vec<(Entity, &Anchor)> = params
            .cabin_anchor_groups
            .get(anchor_group)
            .unwrap()
            .iter()
            .filter_map(|child| params.anchors.get(child).ok())
            .collect();

        for (anchor_id, anchor) in &lift_data.cabin_anchors {
            let mut already_existing = false;
            for (existing_id, existing_anchor) in &existing_lift_anchors {
                if anchor.is_close(*existing_anchor, anchor_close_enough) {
                    id_to_entity.insert(*anchor_id, *existing_id);
                    already_existing = true;
                    break;
                }
            }
            if !already_existing {
                params.commands.entity(anchor_group).with_children(|group| {
                    let e_anchor = group.spawn(AnchorBundle::new(anchor.clone())).id();
                    id_to_entity.insert(*anchor_id, e_anchor);
                });
            }
        }
    }

    for (level_id, level_data) in &from_site_data.levels {
        let level_e = *id_to_entity.get(level_id).unwrap();
        let existing_level_anchors: Vec<(Entity, &Anchor)> = params
            .levels
            .get(level_e)
            .unwrap()
            .3
            .iter()
            .filter_map(|child| params.anchors.get(child).ok())
            .collect();
        for (anchor_id, anchor) in &level_data.anchors {
            let mut already_existing = false;
            for (existing_id, existing_anchor) in &existing_level_anchors {
                if anchor.is_close(*existing_anchor, anchor_close_enough) {
                    id_to_entity.insert(*anchor_id, *existing_id);
                    already_existing = true;
                    break;
                }
            }
            if !already_existing {
                params.commands.entity(level_e).with_children(|level| {
                    let e_anchor = level.spawn(AnchorBundle::new(anchor.clone())).id();
                    id_to_entity.insert(*anchor_id, e_anchor);
                });
            }
        }
    }

    {
        let existing_site_anchors: Vec<(Entity, &Anchor)> = site_children
            .iter()
            .filter_map(|child| params.anchors.get(child).ok())
            .collect();
        for (anchor_id, anchor) in &from_site_data.anchors {
            let mut already_existing = false;
            for (existing_id, existing_anchor) in &existing_site_anchors {
                if anchor.is_close(*existing_anchor, anchor_close_enough) {
                    id_to_entity.insert(*anchor_id, *existing_id);
                    already_existing = true;
                    break;
                }
            }
            if !already_existing {
                params.commands.entity(into_site).with_children(|site| {
                    let e_anchor = site.spawn(AnchorBundle::new(anchor.clone())).id();
                    id_to_entity.insert(*anchor_id, e_anchor);
                });
            }
        }
    }

    for (nav_graph_id, nav_graph_data) in &from_site_data.navigation.guided.graphs {
        params.commands.entity(into_site).with_children(|site| {
            let e = site
                .spawn((Transform::default(), Visibility::default()))
                .insert(nav_graph_data.clone())
                .id();
            id_to_entity.insert(*nav_graph_id, e);
        });
    }

    for (lane_id, lane_data) in &from_site_data.navigation.guided.lanes {
        let lane_data = lane_data
            .convert(&id_to_entity)
            .map_err(ImportNavGraphError::BrokenInternalReference)?;
        params.commands.entity(into_site).with_children(|site| {
            let e = site.spawn(lane_data).id();
            id_to_entity.insert(*lane_id, e);
        });
    }

    for (location_id, location_data) in &from_site_data.navigation.guided.locations {
        let location_data = location_data
            .convert(&id_to_entity)
            .map_err(ImportNavGraphError::BrokenInternalReference)?;
        params.commands.entity(into_site).with_children(|site| {
            let e = site.spawn(location_data).id();
            id_to_entity.insert(*location_id, e);
        });
    }

    Ok(())
}

pub fn import_nav_graph(
    mut params: ImportNavGraphParams,
    mut import_requests: EventReader<ImportNavGraphs>,
) {
    for r in import_requests.read() {
        if let Err(err) = generate_imported_nav_graphs(&mut params, r.into_site, &r.from_site) {
            error!("Failed to import nav graph: {err}");
        }
    }
}