vibe-graph-bevy 0.1.0

use bevy::prelude::*;
use petgraph::stable_graph::StableDiGraph;
use petgraph::visit::{EdgeRef, IntoEdgeReferences};

use crate::benchmark::{self, GraphScale};
use crate::layout::{ForceLayout3D, LayoutConfig};

#[derive(Resource)]
pub struct GraphLayout {
    pub layout: ForceLayout3D,
    pub node_count: usize,
    pub edge_count: usize,
    pub running: bool,
    pub iterations_per_frame: usize,
    #[allow(dead_code)]
    pub labels: Vec<String>,
    #[allow(dead_code)]
    pub source_graph: Option<vibe_graph_core::SourceCodeGraph>,
}

impl GraphLayout {
    pub fn iterations(&self) -> u64 {
        self.layout.iterations
    }

    pub fn positions(&self) -> &[Vec3] {
        &self.layout.positions
    }

    pub fn edges(&self) -> &[(usize, usize)] {
        &self.layout.edges
    }
}

#[derive(Resource)]
pub struct LayoutSettings {
    pub config: LayoutConfig,
    pub iterations_per_frame: usize,
    pub scale: GraphScale,
    pub custom_graph_path: Option<String>,
    pub node_size: f32,
}

impl Default for LayoutSettings {
    fn default() -> Self {
        Self {
            config: LayoutConfig::default(),
            iterations_per_frame: 10,
            scale: GraphScale::Medium,
            custom_graph_path: None,
            node_size: 1.0,
        }
    }
}

impl GraphLayout {
    pub fn from_petgraph(g: &StableDiGraph<String, String>, settings: &LayoutSettings) -> Self {
        let node_indices: Vec<_> = g.node_indices().collect();
        let node_count = node_indices.len();

        let labels: Vec<String> = node_indices.iter().map(|&idx| g[idx].clone()).collect();

        let mut idx_map = std::collections::HashMap::new();
        for (i, &ni) in node_indices.iter().enumerate() {
            idx_map.insert(ni, i);
        }

        let edges: Vec<(usize, usize)> = g
            .edge_references()
            .filter_map(|e| {
                let src = idx_map.get(&e.source())?;
                let tgt = idx_map.get(&e.target())?;
                Some((*src, *tgt))
            })
            .collect();

        let edge_count = edges.len();

        let layout = ForceLayout3D::new(node_count, edges, settings.config.clone());

        Self {
            layout,
            node_count,
            edge_count,
            running: true,
            iterations_per_frame: settings.iterations_per_frame,
            labels,
            source_graph: None,
        }
    }

    pub fn from_source_code_graph(
        g: &vibe_graph_core::SourceCodeGraph,
        settings: &LayoutSettings,
    ) -> Self {
        let node_count = g.nodes.len();

        let mut idx_map = std::collections::HashMap::new();
        let mut labels = Vec::with_capacity(node_count);

        for (i, node) in g.nodes.iter().enumerate() {
            idx_map.insert(node.id, i);
            labels.push(node.name.clone());
        }

        let edges: Vec<(usize, usize)> = g
            .edges
            .iter()
            .filter_map(|e| {
                let src = idx_map.get(&e.from)?;
                let tgt = idx_map.get(&e.to)?;
                Some((*src, *tgt))
            })
            .collect();

        let edge_count = edges.len();

        let layout = ForceLayout3D::new(node_count, edges, settings.config.clone());

        Self {
            layout,
            node_count,
            edge_count,
            running: true,
            iterations_per_frame: settings.iterations_per_frame,
            labels,
            source_graph: Some(g.clone()),
        }
    }
}

pub fn init_graph(
    mut commands: Commands,
    settings: Res<LayoutSettings>,
    initial_graph: Option<Res<crate::InitialGraph>>,
) {
    let layout = if let Some(init_graph) = initial_graph {
        GraphLayout::from_source_code_graph(&init_graph.0, &settings)
    } else if let Some(path) = &settings.custom_graph_path {
        tracing::info!("Loading custom graph from {}", path);
        if let Ok(file_content) = std::fs::read_to_string(path) {
            match serde_json::from_str::<vibe_graph_core::SourceCodeGraph>(&file_content) {
                Ok(sc_graph) => GraphLayout::from_source_code_graph(&sc_graph, &settings),
                Err(e) => {
                    tracing::error!("Failed to parse custom graph JSON: {}", e);
                    let g = benchmark::generate_random_graph(settings.scale.node_count());
                    GraphLayout::from_petgraph(&g, &settings)
                }
            }
        } else {
            tracing::error!("Failed to read custom graph file: {}", path);
            let g = benchmark::generate_random_graph(settings.scale.node_count());
            GraphLayout::from_petgraph(&g, &settings)
        }
    } else {
        let g = benchmark::generate_random_graph(settings.scale.node_count());
        GraphLayout::from_petgraph(&g, &settings)
    };

    tracing::info!(
        nodes = layout.node_count,
        edges = layout.edge_count,
        "Graph initialized"
    );
    commands.insert_resource(layout);
}

pub fn step_layout(mut layout: ResMut<GraphLayout>) {
    if !layout.running {
        return;
    }

    // Scale down iterations per frame for large graphs to maintain FPS
    let iters = if layout.node_count >= 5000 {
        layout.iterations_per_frame.min(2)
    } else if layout.node_count >= 1000 {
        layout.iterations_per_frame.min(5)
    } else {
        layout.iterations_per_frame
    };

    for _ in 0..iters {
        layout.layout.step();
    }
}