orrery 0.4.1

//! Basic component layout engine.
//!
//! This module provides a layout engine for component diagrams
//! using a simple, deterministic algorithm.

use std::{
    collections::{HashMap, HashSet, VecDeque},
    rc::Rc,
};

use orrery_core::{
    draw::{Drawable, PositionedArrowWithText, Shape, ShapeWithText, Text},
    geometry::{Insets, Point, Size},
    identifier::Id,
    semantic::{Block, Relation},
};

use crate::{
    error::RenderError,
    layout::{
        component::{self, ArrowPlacer, Component, Layout, SmartArrowPlacer},
        engines::{ComponentEngine, EmbeddedLayouts},
        layer::{ContentStack, PositionedContent},
    },
    structure::{ComponentGraph, ContainmentScope},
};

/// Basic layout engine for component diagrams.
///
/// Assigns components to layers using BFS traversal and positions them
/// in a left-to-right arrangement. Relation labels influence inter-layer
/// spacing, and containment scopes are processed from innermost to
/// outermost.
#[derive(Default)]
pub struct Engine {
    /// Padding inside component shapes.
    padding: Insets,
    /// Padding around text elements.
    text_padding: f32,
    /// Minimum spacing between adjacent components.
    min_spacing: f32,
    arrow_placer: SmartArrowPlacer,
}

impl Engine {
    /// Create a new basic component layout engine.
    pub fn new() -> Self {
        Self {
            text_padding: 20.0,
            ..Self::default()
        }
    }

    /// Set the padding around components.
    pub fn set_padding(&mut self, padding: Insets) -> &mut Self {
        self.padding = padding;
        self
    }

    /// Set the padding for text elements.
    #[allow(dead_code)]
    pub fn set_text_padding(&mut self, padding: f32) -> &mut Self {
        self.text_padding = padding;
        self
    }

    /// Set the minimum spacing between components.
    pub fn set_min_spacing(&mut self, spacing: f32) -> &mut Self {
        self.min_spacing = spacing;
        self
    }

    /// Calculate the layout for a component diagram.
    ///
    /// # Arguments
    ///
    /// * `graph` - The component diagram graph to lay out.
    /// * `embedded_layouts` - Pre-calculated layouts for embedded diagrams.
    ///
    /// # Errors
    ///
    /// Returns [`RenderError::Layout`] if position or shape calculation fails.
    pub fn calculate_layout<'a>(
        &self,
        graph: &'a ComponentGraph<'a, '_>,
        embedded_layouts: &EmbeddedLayouts<'a>,
    ) -> Result<ContentStack<Layout<'a>>, RenderError> {
        let mut content_stack = ContentStack::<Layout<'a>>::new();
        let mut positioned_content_sizes = HashMap::<Id, Size>::new();

        for containment_scope in graph.containment_scopes() {
            // Calculate component shapes - they contain all sizing information
            let mut component_shapes = self.calculate_component_shapes(
                graph,
                containment_scope,
                &positioned_content_sizes,
                embedded_layouts,
            )?;

            // Calculate positions for components
            let positions = self.positions(graph, containment_scope, &component_shapes)?;

            // Build the final component list using the pre-configured shapes
            let components: Vec<Component> = graph
                .scope_nodes(containment_scope)
                .map(|node| {
                    let position = *positions.get(&node.id()).ok_or_else(|| {
                        RenderError::Layout(format!(
                            "Position not found for node '{node}' during component layout"
                        ))
                    })?;
                    let shape_with_text = component_shapes.remove(&node.id()).ok_or_else(|| {
                        RenderError::Layout(format!(
                            "Shape not found for node '{node}' during component layout"
                        ))
                    })?;

                    Ok(Component::new(node, shape_with_text, position))
                })
                .collect::<Result<Vec<_>, RenderError>>()?;

            // Map node IDs to their component indices
            let component_indices: HashMap<_, _> = components
                .iter()
                .enumerate()
                .map(|(idx, component)| (component.node_id(), idx))
                .collect();

            let relations = self.place_scope_relations(
                graph,
                containment_scope,
                &components,
                &component_indices,
            );

            let positioned_content = PositionedContent::new(Layout::new(components, relations));

            if let Some(container) = containment_scope.container() {
                // If this layer is a container, we need to adjust its size based on its contents
                let size = positioned_content.layout_size();
                positioned_content_sizes.insert(container, size);
            }
            content_stack.push(positioned_content);
        }

        component::adjust_positioned_contents_offset(&mut content_stack, graph)?;

        Ok(content_stack)
    }

    /// Builds positioned arrows for all visible relations in a containment scope.
    ///
    /// Parallel/reverse relations between the same component pair are placed
    /// together so the [`ArrowPlacer`] can offset them. Relations whose
    /// endpoints are not visible at this scope are skipped.
    fn place_scope_relations<'a>(
        &self,
        graph: &'a ComponentGraph<'a, '_>,
        containment_scope: &ContainmentScope,
        components: &[Component<'a>],
        component_indices: &HashMap<Id, usize>,
    ) -> Vec<PositionedArrowWithText<'a>> {
        let mut buckets: HashMap<(Id, Id), Vec<&'a Relation>> = HashMap::new();

        for relation in graph.scope_relations(containment_scope) {
            let source_id = relation.source();
            let target_id = relation.target();

            if !component_indices.contains_key(&source_id)
                || !component_indices.contains_key(&target_id)
            {
                continue;
            }

            // Canonicalize the pair key so (a,b) and (b,a) land in the same bucket.
            let key = if buckets.contains_key(&(target_id, source_id)) {
                (target_id, source_id)
            } else {
                (source_id, target_id)
            };
            buckets.entry(key).or_default().push(relation);
        }

        buckets
            .into_iter()
            .flat_map(|((src_id, tgt_id), bucket)| {
                let src = &components[component_indices[&src_id]];
                let tgt = &components[component_indices[&tgt_id]];
                self.arrow_placer.place(&bucket, src, tgt)
            })
            .collect()
    }

    /// Calculate component shapes with proper content size and padding.
    fn calculate_component_shapes<'a>(
        &self,
        graph: &'a ComponentGraph<'a, '_>,
        containment_scope: &ContainmentScope,
        positioned_content_sizes: &HashMap<Id, Size>,
        embedded_layouts: &EmbeddedLayouts<'a>,
    ) -> Result<HashMap<Id, ShapeWithText<'a>>, RenderError> {
        let mut component_shapes: HashMap<Id, ShapeWithText<'a>> = HashMap::new();

        // TODO: move it to the best place.
        for node in graph.scope_nodes(containment_scope) {
            let mut shape = Shape::new(Rc::clone(node.shape_definition()));
            shape.set_padding(self.padding);
            let text = Text::new(node.shape_definition().text(), node.display_text());
            let mut shape_with_text = ShapeWithText::new(shape, Some(text));

            match node.block() {
                Block::Diagram(_) => {
                    // Since we process in post-order (innermost to outermost),
                    // embedded diagram layouts should already be calculated and available
                    let layout = embedded_layouts.get(&node.id()).ok_or_else(|| {
                        RenderError::Layout(format!(
                            "Embedded layout not found for diagram block '{node}'"
                        ))
                    })?;

                    let content_size = layout.calculate_size();
                    shape_with_text
                        .set_inner_content_size(content_size)
                        .map_err(|err| {
                            RenderError::Layout(format!(
                                "Failed to set content size for diagram block '{node}': {err}"
                            ))
                        })?;
                }
                Block::Scope(_) => {
                    let content_size =
                        *positioned_content_sizes.get(&node.id()).ok_or_else(|| {
                            RenderError::Layout(format!("Scope size not found for node '{node}'"))
                        })?;
                    shape_with_text
                        .set_inner_content_size(content_size)
                        .map_err(|err| {
                            RenderError::Layout(format!(
                                "Failed to set content size for scope block '{node}': {err}"
                            ))
                        })?;
                }
                Block::None => {
                    // No content to size, so don't call set_inner_content_size
                }
            };
            component_shapes.insert(node.id(), shape_with_text);
        }

        Ok(component_shapes)
    }

    /// Calculate positions for components in a containment scope.
    fn positions<'a>(
        &self,
        graph: &'a ComponentGraph<'a, '_>,
        containment_scope: &ContainmentScope,
        component_shapes: &HashMap<Id, ShapeWithText<'a>>,
    ) -> Result<HashMap<Id, Point>, RenderError> {
        // Step 1: Assign layers for the top-level nodes
        let layers = Self::assign_layers_for_containment_scope_graph(graph, containment_scope)?;
        // Step 2: Calculate layer metrics (widths and spacings)
        let (layer_widths, layer_spacings) =
            self.calculate_layer_metrics(graph, containment_scope, &layers, component_shapes)?;
        // Step 3: Calculate X positions for each layer
        let layer_x_positions = self.calculate_layer_x_positions(&layer_widths, &layer_spacings);
        // Step 4: Position nodes within their layers
        self.position_nodes_in_layers(&layers, &layer_x_positions, component_shapes)
    }

    /// Calculate metrics for each layer: widths and spacings between layers.
    fn calculate_layer_metrics<'a>(
        &self,
        graph: &'a ComponentGraph<'a, '_>,
        containment_scope: &ContainmentScope,
        layers: &[Vec<Id>],
        component_shapes: &HashMap<Id, ShapeWithText<'a>>,
    ) -> Result<(Vec<f32>, Vec<f32>), RenderError> {
        // Calculate max width for each layer
        let layer_widths: Vec<f32> = layers
            .iter()
            .map(|layer| {
                layer.iter().try_fold(0.0_f32, |max_width, &node_idx| {
                    let width = component_shapes
                        .get(&node_idx)
                        .map(|shape| shape.size().width())
                        .ok_or_else(|| {
                            RenderError::Layout(format!(
                                "Component shape not found for node '{}' during layer metrics calculation",
                                node_idx
                            ))
                        })?;
                    Ok(max_width.max(width))
                })
            })
            .collect::<Result<Vec<f32>, RenderError>>()?;

        // Initialize spacings with default padding
        let mut layer_spacings =
            vec![self.padding.horizontal_sum() / 2.0; layers.len().saturating_sub(1)];

        // Adjust spacings based on relation labels
        for relation in graph.scope_relations(containment_scope) {
            if let Some(text) = relation.text() {
                let label_width = text.calculate_size().width();

                // Find layers for source and target nodes
                let (source_layer, target_layer) = self.find_node_layers(relation, layers);

                if let (Some(src), Some(tgt)) = (source_layer, target_layer)
                    && src != tgt
                {
                    // Only adjust spacing for relations between different layers
                    let min_layer = src.min(tgt);
                    let needed_spacing = label_width + 30.0; // Add some padding

                    // Update spacing if label requires more space
                    if min_layer < layer_spacings.len() {
                        layer_spacings[min_layer] = layer_spacings[min_layer].max(needed_spacing);
                    }
                }
            }
        }

        Ok((layer_widths, layer_spacings))
    }

    /// Find which layer contains nodes for a given relation.
    // PERF: Deprecate this method in favor of a more efficient approach.
    fn find_node_layers(
        &self,
        relation: &Relation,
        layers: &[Vec<Id>],
    ) -> (Option<usize>, Option<usize>) {
        let mut source_layer = None;
        let mut target_layer = None;

        for (layer_idx, layer_nodes) in layers.iter().enumerate() {
            for node_id in layer_nodes {
                if *node_id == relation.source() {
                    source_layer = Some(layer_idx);
                }
                if *node_id == relation.target() {
                    target_layer = Some(layer_idx);
                }
            }
        }

        (source_layer, target_layer)
    }

    /// Calculate X positions for each layer based on widths and spacings.
    fn calculate_layer_x_positions(
        &self,
        layer_widths: &[f32],
        layer_spacings: &[f32],
    ) -> Vec<f32> {
        let mut layer_x_positions = Vec::with_capacity(layer_widths.len());
        let mut x_pos = 0.0;

        for (i, width) in layer_widths.iter().enumerate() {
            layer_x_positions.push(x_pos + width / 2.0);
            let spacing = if i < layer_spacings.len() {
                layer_spacings[i]
            } else {
                self.padding.horizontal_sum() / 2.0
            };
            x_pos += width + spacing;
        }

        layer_x_positions
    }

    /// Position nodes within their layers.
    fn position_nodes_in_layers<'a>(
        &self,
        layers: &[Vec<Id>],
        layer_x_positions: &[f32],
        component_shapes: &HashMap<Id, ShapeWithText<'a>>,
    ) -> Result<HashMap<Id, Point>, RenderError> {
        let mut positions = HashMap::new();

        for (layer_idx, layer_nodes) in layers.iter().enumerate() {
            let x = layer_x_positions[layer_idx];

            // Calculate heights for vertical positioning
            let mut y_pos = 0.0;
            for (j, &node_idx) in layer_nodes.iter().enumerate() {
                let node_height = component_shapes
                    .get(&node_idx)
                    .ok_or_else(|| {
                        RenderError::Layout(format!(
                            "Component shape not found for node '{}' during position calculation",
                            node_idx
                        ))
                    })?
                    .size()
                    .height();

                if j > 0 {
                    y_pos += self.padding.vertical_sum() / 2.0; // Space between components
                }

                let y = y_pos + node_height / 2.0;
                positions.insert(node_idx, Point::new(x, y));

                y_pos += node_height;
            }
        }

        Ok(positions)
    }

    /// Helper method to assign layers for a specific graph.
    fn assign_layers_for_containment_scope_graph(
        graph: &ComponentGraph,
        containment_scope: &ContainmentScope,
    ) -> Result<Vec<Vec<Id>>, RenderError> {
        let mut layers = Vec::new();
        let mut visited = HashSet::new();

        // Find root nodes
        let root_nodes: Vec<_> = graph.scope_roots(containment_scope).collect();

        let start_nodes = if root_nodes.is_empty() {
            graph.scope_nodes(containment_scope).take(1).collect()
        } else {
            root_nodes
        };

        // Perform BFS to assign layers
        let mut queue = VecDeque::new();
        for node in start_nodes {
            queue.push_back((node, 0));
        }

        while let Some((node, layer)) = queue.pop_front() {
            if !visited.insert(node.id()) {
                continue;
            }
            while layers.len() <= layer {
                layers.push(Vec::new());
            }

            layers[layer].push(node.id());

            for neighbor in graph
                .scope_outgoing_neighbors(containment_scope, node.id())
                .filter(|node| !visited.contains(&node.id()))
            {
                queue.push_back((neighbor, layer + 1));
            }
        }

        // Handle disconnected components by processing any remaining unvisited nodes
        while visited.len() < containment_scope.nodes_count() {
            // Find an unvisited node to start a new component
            // PERF: This can be an outer loop for a single iteration.
            let unvisited_node_id = containment_scope
                .node_ids()
                .find(|id| !visited.contains(id))
                .ok_or_else(|| {
                    RenderError::Layout("Expected unvisited nodes but none found".to_string())
                })?;

            let unvisited_node = graph.node_by_id(unvisited_node_id).ok_or_else(|| {
                RenderError::Layout(format!(
                    "Node '{}' from containment scope not found in graph",
                    unvisited_node_id
                ))
            })?;

            queue.push_back((unvisited_node, 0));

            // Process this disconnected component using the same BFS logic
            // TODO: this is a duplicated code.
            while let Some((node, layer)) = queue.pop_front() {
                if !visited.insert(node.id()) {
                    continue;
                }
                while layers.len() <= layer {
                    layers.push(Vec::new());
                }

                layers[layer].push(node.id());

                for neighbor in graph
                    .scope_outgoing_neighbors(containment_scope, node.id())
                    .filter(|node| !visited.contains(&node.id()))
                {
                    queue.push_back((neighbor, layer + 1));
                }
            }
        }

        Ok(layers)
    }
}

impl ComponentEngine for Engine {
    fn calculate<'a>(
        &self,
        graph: &'a ComponentGraph<'a, '_>,
        embedded_layouts: &EmbeddedLayouts<'a>,
    ) -> Result<ContentStack<Layout<'a>>, RenderError> {
        self.calculate_layout(graph, embedded_layouts)
    }
}