orrery 0.2.0

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

use std::{cmp::Ordering, collections::HashMap, rc::Rc};

use orrery_core::{
    draw::{self, Drawable as _},
    geometry::{Insets, Point, Size},
    identifier::Id,
    semantic,
};

use crate::{
    error::RenderError,
    layout::{
        component::Component,
        engines::{EmbeddedLayouts, SequenceEngine},
        layer::{ContentStack, PositionedContent},
        sequence::{ActivationBox, ActivationTiming, FragmentTiming, Layout, Message, Participant},
    },
    structure::{SequenceEvent, SequenceGraph},
};

/// Collected output from [`Engine::process_events`]: messages, activation boxes,
/// fragments, notes, and the final Y coordinate.
type ProcessEventsResult<'a> = (
    Vec<Message<'a>>,
    Vec<ActivationBox>,
    Vec<draw::PositionedDrawable<draw::Fragment>>,
    Vec<draw::PositionedDrawable<draw::Note>>,
    f32,
);

/// Basic deterministic layout engine for sequence diagrams.
///
/// Distributes participants horizontally, processes events sequentially
/// to place messages and activations, and builds lifelines from
/// participant boxes to the final event position.
pub struct Engine {
    /// Minimum horizontal space between participants.
    min_spacing: f32,
    /// Vertical padding between consecutive events.
    event_padding: f32,
    /// Vertical margin above participant boxes.
    top_margin: f32,
    /// Padding inside participant shapes.
    padding: Insets,
    /// Extra horizontal padding to accommodate message labels.
    label_padding: f32,
}

impl Engine {
    /// Create a new basic sequence layout engine
    pub fn new() -> Self {
        Self {
            min_spacing: 40.0, // Minimum spacing between participants
            event_padding: 15.0,
            top_margin: 60.0,
            padding: Insets::uniform(15.0),
            label_padding: 20.0, // Extra padding for labels
        }
    }

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

    /// Sets the vertical padding between sequence events.
    pub fn set_event_padding(&mut self, padding: f32) -> &mut Self {
        self.event_padding = padding;
        self
    }

    /// Set the top margin of the diagram
    #[allow(dead_code)]
    pub fn set_top_margin(&mut self, margin: f32) -> &mut Self {
        self.top_margin = margin;
        self
    }

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

    /// Set the padding for message labels
    #[allow(dead_code)]
    pub fn set_label_padding(&mut self, padding: f32) -> &mut Self {
        self.label_padding = padding;
        self
    }

    /// Calculate additional spacing needed between participants based on message label sizes
    fn calculate_message_label_spacing(
        &self,
        source_id: Id,
        target_id: Id,
        messages: &[&semantic::Relation],
    ) -> f32 {
        // Filter messages to only those between the two participants
        let relevant_messages = messages.iter().filter(|relation| {
            (relation.source() == source_id && relation.target() == target_id)
                || (relation.source() == target_id && relation.target() == source_id)
        });

        // Extract labels from relations and use shared function to calculate spacing
        let labels = relevant_messages.map(|relation| relation.text());
        crate::layout::positioning::calculate_label_spacing(labels, self.label_padding)
    }

    /// Calculate layout for a sequence diagram.
    ///
    /// # Arguments
    ///
    /// * `graph` - The sequence 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 SequenceGraph<'a>,
        embedded_layouts: &EmbeddedLayouts<'a>,
    ) -> Result<ContentStack<Layout<'a>>, RenderError> {
        // Create shapes with text for participants
        let mut participant_shapes: HashMap<_, _> = HashMap::new();
        for node in graph.nodes() {
            let mut shape = draw::Shape::new(Rc::clone(node.shape_definition()));
            shape.set_padding(self.padding);
            let text = draw::Text::new(node.shape_definition().text(), node.display_text());
            let mut shape_with_text = draw::ShapeWithText::new(shape, Some(text));

            if let semantic::Block::Diagram(_) = node.block() {
                // If this participant has an embedded diagram, use its layout size
                let content_size = if let Some(layout) = embedded_layouts.get(&node.id()) {
                    layout.calculate_size()
                } else {
                    Size::zero()
                };

                shape_with_text
                    .set_inner_content_size(content_size)
                    .map_err(|err| {
                        RenderError::Layout(format!(
                            "Failed to set content size for participant '{}': {err}",
                            node.display_text()
                        ))
                    })?;
            }
            // For non-Diagram blocks, don't call set_inner_content_size
            participant_shapes.insert(node.id(), shape_with_text);
        }

        // Collect all messages to consider their labels for spacing
        let mut messages_vec = Vec::new();
        for relation in graph.relations() {
            messages_vec.push(relation);
        }

        // Calculate additional spacings based on message labels
        let mut spacings = Vec::<f32>::new();
        let mut nodes_iter = graph.nodes();
        if let Some(mut last_node) = nodes_iter.next() {
            for node in nodes_iter {
                let spacing =
                    self.calculate_message_label_spacing(last_node.id(), node.id(), &messages_vec);
                spacings.push(spacing);
                last_node = node;
            }
        }

        // Get list of node indices and their sizes
        let mut sizes: Vec<Size> = Vec::new();
        for id in graph.node_ids() {
            let shape_with_text = participant_shapes.get(id).ok_or_else(|| {
                RenderError::Layout("Participant shape not found for node".to_string())
            })?;
            sizes.push(shape_with_text.size());
        }

        // Calculate horizontal positions using positioning algorithms
        let x_positions = crate::layout::positioning::distribute_horizontally(
            &sizes,
            self.min_spacing,
            Some(&spacings),
        );

        // Create participants and store their indices
        let mut components: HashMap<Id, Component> = HashMap::new();
        for (i, node) in graph.nodes().enumerate() {
            let shape_with_text = participant_shapes.remove(&node.id()).ok_or_else(|| {
                RenderError::Layout(format!("Participant shape not found for node '{node}'"))
            })?;
            let mut position = Point::new(x_positions[i], self.top_margin);

            // If this node contains an embedded diagram, adjust position to normalize
            // the embedded layout's coordinate system to start at origin
            if let semantic::Block::Diagram(_) = node.block()
                && let Some(layout) = embedded_layouts.get(&node.id())
            {
                position = position.add_point(layout.normalize_offset());
            }

            let component = Component::new(node, shape_with_text, position);
            components.insert(node.id(), component);
        }

        // Calculate message positions and update lifeline ends
        let participants_height = components
            .values()
            .map(|component| component.drawable().size().height())
            .max_by(|a, b| a.partial_cmp(b).unwrap_or(Ordering::Equal))
            .unwrap_or_default();

        let (messages, activations, fragments, notes, lifeline_end) =
            self.process_events(graph, participants_height, &components)?;

        // Update lifeline ends to match diagram height and finalize lifelines
        let participants: HashMap<Id, Participant<'a>> = components
            .into_iter()
            .map(|(id, component)| {
                // Rebuild the positioned lifeline with the final height
                let position = component.position();
                let lifeline_start_y = component.bounds().max_y();
                let height = (lifeline_end - lifeline_start_y).max(0.0);
                let lifeline = draw::PositionedDrawable::new(draw::Lifeline::new(
                    Rc::new(draw::LifelineDefinition::default()),
                    height,
                ))
                .with_position(Point::new(position.x(), lifeline_start_y));

                (id, Participant::new(component, lifeline))
            })
            .collect();

        let layout = Layout::new(
            participants,
            messages,
            activations,
            fragments,
            notes,
            lifeline_end,
        );

        let mut content_stack = ContentStack::new();
        content_stack.push(PositionedContent::new(layout));

        Ok(content_stack)
    }

    /// Process all sequence diagram events to create layout components.
    ///
    /// This method processes ordered events sequentially to create messages, activation boxes,
    /// and fragments with precise timing and positioning. It uses a HashMap-based stack approach
    /// (Id -> [`Vec<ActivationTiming>`]) to track activation periods per participant and converts
    /// them to ActivationBox objects when deactivation occurs.
    ///
    /// # Algorithm
    /// 1. Iterate through ordered events sequentially
    /// 2. For `SequenceEvent::Relation`: Create [`Message`] centered at `current_y + height/2`, update fragment bounds if inside a fragment, record last relation Y, then advance Y by message size + `event_padding`
    /// 3. For `SequenceEvent::Activate`: Create ActivationTiming at last relation Y position (aligning with the triggering message), push to participant's stack
    /// 4. For `SequenceEvent::Deactivate`: Pop activation, convert to ActivationBox ending at last relation Y position
    /// 5. For `SequenceEvent::FragmentStart`: Create FragmentTiming at `current_y` and push to fragment stack
    /// 6. For `SequenceEvent::FragmentSectionStart`: Start new section in current fragment, then advance Y by header height + `event_padding`
    /// 7. For `SequenceEvent::FragmentSectionEnd`: End current section
    /// 8. For `SequenceEvent::FragmentEnd`: Pop fragment, convert to Fragment with final bounds, then advance Y by bottom padding + `event_padding`
    /// 9. For `SequenceEvent::Note`: Create positioned [`Note`](draw::Note) at `current_y`, then advance Y by note size + `event_padding`
    ///
    /// # Parameters
    /// * `graph` - The sequence diagram graph containing ordered events
    /// * `participants_height` - Height of the participant boxes for calculating initial Y position
    /// * `components` - Map of participant IDs to their positioned components, used for fragment bounds tracking
    ///
    /// # Returns
    /// A tuple containing:
    /// * `Vec<Message<'a>>` - All messages with their positions and arrow information
    /// * `Vec<ActivationBox>` - All activation boxes with precise positioning and nesting levels
    /// * `Vec<draw::PositionedDrawable<draw::Fragment>>` - All fragments with their sections and bounds
    /// * `Vec<draw::PositionedDrawable<draw::Note>>` - All notes with their positions and content
    /// * `f32` - The final Y coordinate (lifeline end position)
    fn process_events<'a>(
        &self,
        graph: &SequenceGraph<'a>,
        participants_height: f32,
        components: &HashMap<Id, Component<'a>>,
    ) -> Result<ProcessEventsResult<'a>, RenderError> {
        let mut messages: Vec<Message<'a>> = Vec::new();
        let mut activation_boxes: Vec<ActivationBox> = Vec::new();
        let mut fragments: Vec<draw::PositionedDrawable<draw::Fragment>> = Vec::new();
        let mut notes: Vec<draw::PositionedDrawable<draw::Note>> = Vec::new();

        let mut activation_stack: HashMap<Id, Vec<ActivationTiming>> = HashMap::new();
        let mut fragment_stack: Vec<FragmentTiming> = Vec::new();

        // Initial Y is the top edge of the first event area.
        let mut current_y = self.top_margin + participants_height + self.event_padding;
        // Track the Y position of the last placed relation (before spacing advance).
        let mut last_relation_y = current_y;

        for event in graph.events() {
            match event {
                SequenceEvent::Relation(relation) => {
                    let mut message =
                        Message::from_ast(relation, relation.source(), relation.target());
                    let message_height = message.min_size().height();

                    // Center the arrow line within the message's vertical extent.
                    let center_y = current_y + message_height / 2.0;
                    message.set_y_position(center_y);
                    messages.push(message);

                    // Update fragment bounds if we're inside a fragment
                    // NOTE: For perfectly accurate bounds, this should use calculate_message_endpoint_x()
                    // to account for activation box offsets. Currently using participant center positions
                    // as a simpler approximation that is adequate for most cases.
                    if let Some(fragment_timing) = fragment_stack.last_mut() {
                        let source_x = components[&relation.source()].position().x();
                        let target_x = components[&relation.target()].position().x();
                        fragment_timing.update_x(source_x, target_x);
                    }

                    last_relation_y = center_y;
                    current_y += message_height + self.event_padding;
                }
                SequenceEvent::Activate(activate) => {
                    let node_id = activate.component();
                    // Calculate nesting level for this node
                    let nesting_level = activation_stack
                        .get(&node_id)
                        .map(|stack| stack.len() as u32)
                        .unwrap_or(0);

                    let new_timing = ActivationTiming::new(
                        node_id,
                        last_relation_y,
                        nesting_level,
                        Rc::clone(activate.definition()),
                    );

                    // Add to the stack for this node
                    activation_stack
                        .entry(node_id)
                        .or_default()
                        .push(new_timing);
                }
                SequenceEvent::Deactivate(node_id) => {
                    // Pop the most recent activation for this node
                    if let Some(node_stack) = activation_stack.get_mut(node_id) {
                        if let Some(completed_timing) = node_stack.pop() {
                            let activation_box =
                                completed_timing.to_activation_box(last_relation_y);
                            activation_boxes.push(activation_box);
                        }

                        // Clean up empty stacks to avoid memory bloat
                        if node_stack.is_empty() {
                            activation_stack.remove(node_id);
                        }
                    }
                }
                SequenceEvent::FragmentStart(fragment) => {
                    let fragment_timing = FragmentTiming::new(fragment, current_y);
                    fragment_stack.push(fragment_timing);
                    // No current_y advance here; handled in FragmentSectionStart
                }
                SequenceEvent::FragmentSectionStart(fragment_section) => {
                    let fragment_timing = fragment_stack
                        .last_mut()
                        .expect("fragment_timing stack is empty");
                    fragment_timing.start_section(fragment_section, current_y);
                    current_y += fragment_timing.section_header_height(fragment_section)
                        + self.event_padding;
                }
                SequenceEvent::FragmentSectionEnd => {
                    let fragment_timing = fragment_stack
                        .last_mut()
                        .expect("fragment_timing stack is empty");
                    fragment_timing.end_section(current_y).unwrap();
                }
                SequenceEvent::FragmentEnd => {
                    let fragment_timing = fragment_stack
                        .pop()
                        .expect("fragment_timing stack is empty");
                    let fragment_bottom_padding = fragment_timing.bottom_padding();
                    let fragment = fragment_timing.into_fragment(current_y);
                    fragments.push(fragment);
                    current_y += fragment_bottom_padding + self.event_padding;
                }
                SequenceEvent::Note(note) => {
                    let positioned_note =
                        self.create_positioned_note(note, components, current_y)?;
                    let note_height = positioned_note.size().height();

                    notes.push(positioned_note);
                    current_y += note_height + self.event_padding;
                }
            }
        }

        Ok((messages, activation_boxes, fragments, notes, current_y))
    }

    /// Create a positioned note drawable for a sequence diagram.
    ///
    /// Calculates the appropriate position and width for a note based on the participants
    /// it spans. If `note.on()` is empty, the note spans all participants in the diagram.
    ///
    /// # Arguments
    ///
    /// * `note` - The note element from the AST
    /// * `components` - Map of all participant components in the diagram
    /// * `current_y` - Current Y position in the sequence diagram
    ///
    /// # Returns
    ///
    /// A `PositionedDrawable<Note>` ready to be added to the layout
    fn create_positioned_note<'a>(
        &self,
        note: &semantic::Note,
        components: &HashMap<Id, Component<'a>>,
        current_y: f32,
    ) -> Result<draw::PositionedDrawable<draw::Note>, RenderError> {
        const NOTE_SPACING: f32 = 20.0; // Spacing between note and participant lifeline

        // Select appropriate components: all if on=[], otherwise specified ones
        let filtered_components: Vec<&Component> = if note.on().is_empty() {
            components.values().collect()
        } else {
            note.on()
                .iter()
                .map(|id| {
                    components.get(id).ok_or_else(|| {
                        RenderError::Layout("Component not found for note".to_string())
                    })
                })
                .collect::<Result<Vec<_>, _>>()?
        };

        let edge_map: fn(&Component) -> (f32, f32) = match note.align() {
            semantic::NoteAlign::Over => |component| {
                let center_x = component.position().x();
                let width = component.drawable().size().width();
                let left_edge = center_x - width / 2.0;
                let right_edge = center_x + width / 2.0;
                (left_edge, right_edge)
            },
            semantic::NoteAlign::Left | semantic::NoteAlign::Right => {
                |component| (component.position().x(), component.position().x())
            }
            semantic::NoteAlign::Top | semantic::NoteAlign::Bottom => {
                unreachable!("Alignments is not supported for sequence diagrams")
            }
        };

        let (min_x, max_x) = filtered_components
            .into_iter()
            .map(edge_map)
            .reduce(|(min_x, max_x), (left_x, right_x)| (min_x.min(left_x), max_x.max(right_x)))
            .ok_or_else(|| {
                RenderError::Layout("Note should have at least one participant".to_string())
            })?;

        let mut new_note_def = Rc::clone(note.definition());
        if note.align() == semantic::NoteAlign::Over {
            let note_def_mut = Rc::make_mut(&mut new_note_def);
            note_def_mut.set_min_width(Some(max_x - min_x));
        }

        let note_drawable = draw::Note::new(new_note_def, note.content().to_string());

        let note_size = note_drawable.size();
        let center_x = match note.align() {
            semantic::NoteAlign::Over => (min_x + max_x) / 2.0,
            semantic::NoteAlign::Left => min_x - (note_size.width() / 2.0) - NOTE_SPACING,
            semantic::NoteAlign::Right => max_x + (note_size.width() / 2.0) + NOTE_SPACING,
            semantic::NoteAlign::Top | semantic::NoteAlign::Bottom => {
                unreachable!("Alignments is not supported for sequence diagrams")
            }
        };
        let position = Point::new(center_x, current_y + note_size.height() / 2.0);

        Ok(draw::PositionedDrawable::new(note_drawable).with_position(position))
    }
}

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