use crate::types::{GraphData, GraphEdge, GraphNode};
use bevy::prelude::*;
use petgraph::graph::NodeIndex;
use std::collections::HashMap;
#[must_use]
pub fn get_node_appearance(node_type: Option<&str>) -> (Color, f32) {
match node_type {
Some("organization") => (Color::srgb(0.8, 0.2, 0.2), 1.5), Some("line_of_business") => (Color::srgb(0.8, 0.5, 0.2), 1.2), Some("site") => (Color::srgb(0.2, 0.6, 0.8), 1.0), Some("team") => (Color::srgb(0.2, 0.8, 0.5), 0.8), Some("user") => (Color::srgb(0.6, 0.4, 0.8), 0.6),
Some("database") => (Color::srgb(0.2, 0.4, 0.7), 1.0), Some("actor:participant") => (Color::srgb(0.3, 0.7, 0.3), 0.8), Some(t) if t.starts_with("actor:") => (Color::srgb(0.9, 0.6, 0.2), 0.9), Some("process") => (Color::srgb(0.7, 0.7, 0.2), 0.8), Some("external") => (Color::srgb(0.5, 0.2, 0.7), 0.9),
_ => (Color::srgb(0.5, 0.5, 0.5), 0.7), }
}
pub fn create_graph_visualization(
commands: &mut Commands,
meshes: &mut ResMut<Assets<Mesh>>,
materials: &mut ResMut<Assets<StandardMaterial>>,
graph_data: &GraphData,
) -> HashMap<NodeIndex, Entity> {
let mut node_entities = HashMap::new();
let mut level_counts = HashMap::new();
let mut level_indices = HashMap::new();
for node_idx in graph_data.graph.node_indices() {
let node_info = &graph_data.graph[node_idx];
*level_counts.entry(node_info.level).or_insert(0) += 1;
}
for node_idx in graph_data.graph.node_indices() {
let node_info = &graph_data.graph[node_idx];
let (color, size_mult) = get_node_appearance(node_info.node_type.as_deref());
let level_idx = level_indices.entry(node_info.level).or_insert(0);
let count_at_level = level_counts[&node_info.level];
let level_radius = (node_info.level as f32).mul_add(2.0, 5.0);
let angle = 2.0 * std::f32::consts::PI * (*level_idx as f32) / count_at_level as f32;
let x = level_radius * angle.cos();
let z = level_radius * angle.sin();
let y = node_info.level as f32 * 2.0;
*level_idx += 1;
let node_material = materials.add(StandardMaterial {
base_color: color,
emissive: LinearRgba::BLACK,
..default()
});
let mesh = match node_info.node_type.as_deref() {
Some("organization") => meshes.add(Cuboid::new(1.0, 1.0, 1.0)), Some("line_of_business") => meshes.add(Cylinder::new(0.5, 1.0)), Some("site") => meshes.add(Torus::new(0.3, 0.5)), Some("team") => meshes.add(Sphere::new(0.6)), Some("user") => meshes.add(Capsule3d::new(0.3, 0.4)),
Some("database") => meshes.add(Cylinder::new(0.6, 0.8)), Some("actor:participant") => meshes.add(Cuboid::new(0.8, 0.8, 0.8)), Some(t) if t.starts_with("actor:") => {
meshes.add(Capsule3d::new(0.4, 0.6))
}
Some("process") => meshes.add(Sphere::new(0.5)), Some("external") => meshes.add(Torus::new(0.25, 0.5)),
_ => meshes.add(Sphere::new(0.4)), };
let node_entity = commands
.spawn((
Mesh3d(mesh),
MeshMaterial3d(node_material),
Transform::from_xyz(x, y, z).with_scale(Vec3::splat(size_mult)),
GraphNode {
name: node_info.name.clone(),
index: node_idx,
},
Name::new(node_info.name.clone()),
))
.id();
node_entities.insert(node_idx, node_entity);
}
let edge_material = materials.add(StandardMaterial {
base_color: Color::srgb(0.4, 0.4, 0.4),
..default()
});
for edge in graph_data.graph.edge_indices() {
if let Some((from_idx, to_idx)) = graph_data.graph.edge_endpoints(edge) {
if let (Some(&from_entity), Some(&to_entity)) =
(node_entities.get(&from_idx), node_entities.get(&to_idx))
{
let edge_info = graph_data.graph.edge_weight(edge);
spawn_edge(
commands,
meshes,
materials,
edge_material.clone(),
from_entity,
to_entity,
from_idx,
to_idx,
edge_info,
);
}
}
}
node_entities
}
#[allow(clippy::too_many_arguments)]
fn spawn_edge(
commands: &mut Commands,
meshes: &mut ResMut<Assets<Mesh>>,
materials: &mut ResMut<Assets<StandardMaterial>>,
default_material: Handle<StandardMaterial>,
_from_entity: Entity,
_to_entity: Entity,
from_idx: NodeIndex,
to_idx: NodeIndex,
edge_info: Option<&crate::graph_state::EdgeInfo>,
) {
let (color, thickness) = edge_info.map_or_else(
|| (Color::srgb(0.4, 0.4, 0.4), 0.02), |info| {
match info.edge_type.as_deref() {
Some("sync") => (Color::srgb(0.2, 0.4, 0.8), 0.03), Some("async") => (Color::srgb(0.8, 0.4, 0.2), 0.02), Some("return") => (Color::srgb(0.4, 0.8, 0.4), 0.015), _ => (Color::srgb(0.4, 0.4, 0.4), 0.02), }
},
);
let edge_material = if edge_info.is_some() {
materials.add(StandardMaterial {
base_color: color,
emissive: LinearRgba::from(color) * 0.2, ..default()
})
} else {
default_material
};
let edge_component = GraphEdge {
from: from_idx,
to: to_idx,
label: edge_info.and_then(|info| info.label.clone()),
edge_type: edge_info.and_then(|info| info.edge_type.clone()),
sequence: edge_info.and_then(|info| info.sequence),
};
let edge_entity = commands
.spawn((
Mesh3d(meshes.add(Cylinder::new(thickness, 1.0))),
MeshMaterial3d(edge_material.clone()),
Transform::default(),
edge_component,
))
.id();
if edge_info.is_some() {
commands.spawn((
Mesh3d(meshes.add(Cone {
radius: thickness * 3.0,
height: thickness * 8.0,
})),
MeshMaterial3d(edge_material),
Transform::default(),
crate::types::EdgeArrowHead { edge: edge_entity },
));
}
}
#[allow(clippy::type_complexity)]
pub fn update_edge_positions(
node_query: Query<(&Transform, &GraphNode)>,
mut edge_query: Query<(Entity, &mut Transform, &GraphEdge), Without<GraphNode>>,
mut arrow_query: Query<
(&mut Transform, &crate::types::EdgeArrowHead),
(Without<GraphEdge>, Without<GraphNode>),
>,
_graph_data: Res<GraphData>,
) {
let mut node_positions = HashMap::new();
for (transform, graph_node) in &node_query {
node_positions.insert(graph_node.index, transform.translation);
}
for (edge_entity, mut edge_transform, graph_edge) in &mut edge_query {
if let (Some(&from_pos), Some(&to_pos)) = (
node_positions.get(&graph_edge.from),
node_positions.get(&graph_edge.to),
) {
let direction = to_pos - from_pos;
let distance = direction.length();
let midpoint = from_pos + direction * 0.5;
let up = Vec3::Y;
let rotation = if direction.normalize().dot(up).abs() > 0.999 {
Quat::from_rotation_x(std::f32::consts::FRAC_PI_2)
} else {
Quat::from_rotation_arc(up, direction.normalize())
};
edge_transform.translation = midpoint;
edge_transform.rotation = rotation;
edge_transform.scale = Vec3::new(1.0, distance, 1.0);
for (mut arrow_transform, arrow_head) in &mut arrow_query {
if arrow_head.edge == edge_entity {
let arrow_offset = direction.normalize() * 0.5; arrow_transform.translation = to_pos - arrow_offset;
arrow_transform.rotation = rotation;
}
}
}
}
}