disposition_input_ir_rt 0.3.0

use disposition_ir_model::{
    edge::EdgeId,
    entity::EntityType,
    node::{NodeId, NodeInbuilt},
    IrDiagram,
};
use disposition_model_common::Map;
use disposition_taffy_model::{
    taffy::{self, AvailableSpace, Size, TaffyTree},
    DimensionAndLod, EdgeDescriptionTaffyNodes, EdgeSpacerTaffyNodes, IrToTaffyError,
    MdNodeTaffyIds, ProcessesIncluded, TaffyNodeMappings, TEXT_FONT_SIZE,
};
use typed_builder::TypedBuilder;

use self::{
    edge_description_builder::{EdgeDescriptionBuildResult, EdgeDescriptionBuilder},
    edge_label_builder::EdgeLabelBuilder,
    edge_lca_sibling_distance::EdgeLcaSiblingDistance,
    edge_spacer_builder::EdgeSpacerBuilder,
    highlighted_spans_computer::HighlightedSpansComputer,
    md_spans_computer::MdSpansComputer,
    taffy_build_ctx::TaffyBuildCtx,
    taffy_build_state::TaffyBuildState,
    taffy_container_builder::TaffyContainerBuilder,
    taffy_diagram_node_builder::{FirstLevelNodesBuilt, TaffyDiagramNodeBuilder},
    taffy_node_build_context::NodeMeasureContext,
    text_measure::MONOSPACE_CHAR_WIDTH_RATIO,
};

mod edge_description_builder;
mod edge_label_builder;
mod edge_lca_sibling_distance;
mod edge_spacer_builder;
mod highlighted_spans_computer;
mod md_node_builder;
mod md_spans_computer;
mod taffy_build_ctx;
mod taffy_build_state;
mod taffy_container_builder;
mod taffy_diagram_node_builder;
mod taffy_envelope_builder;
mod taffy_node_build_context;
mod text_measure;

/// Maps an intermediate representation diagram to a `TaffyNodeMappings`.
///
/// # Examples
///
/// ```rust
/// # use disposition_input_ir_rt::IrToTaffyBuilder;
/// # use disposition_ir_model::IrDiagram;
/// # use disposition_taffy_model::DimensionAndLod;
/// #
/// let ir_diagram = IrDiagram::new();
/// let dimension_and_lods = vec![DimensionAndLod::default_lg()];
///
/// let mut taffy_trees = IrToTaffyBuilder::builder()
///     .with_ir_diagram(&ir_diagram)
///     .with_dimension_and_lods(dimension_and_lods)
///     .build();
/// ```
#[derive(Debug, TypedBuilder)]
pub struct IrToTaffyBuilder<'builder> {
    /// The intermediate representation of the diagram to render the taffy trees
    /// for.
    #[builder(setter(prefix = "with_"))]
    ir_diagram: &'builder IrDiagram<'static>,
    /// The dimensions at which elements should be repositioned.
    #[builder(setter(prefix = "with_"), default = vec![
        DimensionAndLod::default_sm(),
        DimensionAndLod::default_md(),
        DimensionAndLod::default_lg(),
    ])]
    dimension_and_lods: Vec<DimensionAndLod>,
    /// What processes to create diagrams for.
    #[builder(setter(prefix = "with_"), default = ProcessesIncluded::All)]
    processes_included: ProcessesIncluded,
}

impl IrToTaffyBuilder<'_> {
    /// Returns an iterator over `TaffyNodeMappings` instances for each
    /// dimension.
    pub fn build(
        &self,
    ) -> Result<impl Iterator<Item = TaffyNodeMappings<'static>>, IrToTaffyError> {
        let IrToTaffyBuilder {
            ir_diagram,
            dimension_and_lods,
            processes_included,
        } = self;

        let taffy_node_mappings_iter =
            dimension_and_lods
                .iter()
                .flat_map(move |dimension_and_lod| {
                    Self::build_taffy_trees_for_dimension(
                        ir_diagram,
                        dimension_and_lod,
                        processes_included,
                    )
                });

        Ok(taffy_node_mappings_iter)
    }

    /// Returns a `TaffyNodeMappings` with all processes as part of the diagram.
    ///
    /// This includes the processes container. Clicking on each process node
    /// reveals the process steps.
    fn build_taffy_trees_for_dimension(
        ir_diagram: &IrDiagram<'static>,
        dimension_and_lod: &DimensionAndLod,
        processes_included: &ProcessesIncluded,
    ) -> impl Iterator<Item = TaffyNodeMappings<'static>> {
        let IrDiagram {
            nodes,
            node_copy_text: _,
            node_hierarchy,
            node_ordering: _,
            edge_groups,
            thing_descs,
            edge_descs,
            edge_labels,
            entity_tooltips: _,
            entity_types,
            tailwind_classes: _,
            node_layouts,
            node_ranks_nested,
            node_nesting_infos,
            edge_face_assignments,
            node_face_edges,
            node_shapes,
            process_step_entities: _,
            process_step_edges: _,
            process_step_ranks,
            process_step_graphs,
            render_options,
            css: _,
        } = ir_diagram;

        let rank_dir = render_options.rank_dir;

        let DimensionAndLod { dimension, lod } = dimension_and_lod;

        let mut taffy_tree = TaffyTree::new();
        let mut node_id_to_taffy = Map::new();
        let mut taffy_id_to_node = Map::new();
        let mut node_id_to_envelope_taffy_node: Map<NodeId<'static>, taffy::NodeId> = Map::new();
        let mut edge_label_leaves = Vec::new();

        // Precompute monospace character width
        let char_width = TEXT_FONT_SIZE * MONOSPACE_CHAR_WIDTH_RATIO;

        let mut md_node_taffy_ids: Map<NodeId<'static>, MdNodeTaffyIds> = Map::new();

        // Precompute markdown / text content per node once, so the
        // node-building, size-measuring, and highlighted-span passes all share
        // the same text.
        let node_md_texts = TaffyBuildCtx::node_md_texts_compute(nodes, thing_descs, *lod);

        let ctx = TaffyBuildCtx {
            node_layouts,
            node_hierarchy,
            entity_types,
            thing_descs,
            edge_descs,
            node_shapes,
            node_ranks_nested,
            process_step_ranks,
            process_step_graphs,
            node_nesting_infos,
            node_face_edges,
            edge_groups,
            rank_dir,
            lod: *lod,
            char_width,
            node_md_texts: &node_md_texts,
        };

        let FirstLevelNodesBuilt {
            entity_type_to_rank_nodes: node_rank_to_nodes_by_entity_type,
            nested_edge_taffy_nodes: first_level_nested_edge_taffy_nodes,
        } = {
            let mut state = TaffyBuildState {
                taffy_tree: &mut taffy_tree,
                node_id_to_taffy: &mut node_id_to_taffy,
                taffy_id_to_node: &mut taffy_id_to_node,
                node_id_to_envelope_taffy_node: &mut node_id_to_envelope_taffy_node,
                edge_label_leaves: &mut edge_label_leaves,
                md_node_taffy_ids: &mut md_node_taffy_ids,
            };
            TaffyDiagramNodeBuilder::build_first_level_nodes(ctx, &mut state, processes_included)
        };
        let mut thing_rank_to_taffy_ids = node_rank_to_nodes_by_entity_type
            .get(&EntityType::ThingDefault)
            .cloned()
            .unwrap_or_default();
        let mut tag_rank_to_taffy_ids = node_rank_to_nodes_by_entity_type
            .get(&EntityType::TagDefault)
            .cloned()
            .unwrap_or_default();
        let mut process_rank_to_taffy_ids = node_rank_to_nodes_by_entity_type
            .get(&EntityType::ProcessDefault)
            .cloned()
            .unwrap_or_default();

        // === Insert spacer taffy nodes for cross-rank edges === //
        //
        // For each edge that crosses multiple ranks, we insert small spacer
        // leaf nodes at every intermediate rank. The edge path will later
        // be routed through these spacer positions to avoid overlapping
        // other nodes.
        let mut edge_spacer_taffy_nodes: Map<EdgeId<'static>, EdgeSpacerTaffyNodes> = Map::new();
        edge_spacer_taffy_nodes.extend(first_level_nested_edge_taffy_nodes.edge_spacer_taffy_nodes);
        edge_spacer_taffy_nodes.extend(EdgeSpacerBuilder::build(
            ctx,
            &mut taffy_tree,
            &EntityType::ThingDefault,
            &mut thing_rank_to_taffy_ids,
            None,
        ));
        edge_spacer_taffy_nodes.extend(EdgeSpacerBuilder::build(
            ctx,
            &mut taffy_tree,
            &EntityType::TagDefault,
            &mut tag_rank_to_taffy_ids,
            None,
        ));
        edge_spacer_taffy_nodes.extend(EdgeSpacerBuilder::build(
            ctx,
            &mut taffy_tree,
            &EntityType::ProcessDefault,
            &mut process_rank_to_taffy_ids,
            None,
        ));

        // === Build edge_description_container nodes for top-level described edges ===
        // //
        //
        // For each described edge at the top level (LCA = root), we create a
        // container node interleaved between rank containers, plus a leaf node
        // for text measurement.
        let mut edge_description_taffy_nodes: Map<EdgeId<'static>, EdgeDescriptionTaffyNodes> =
            Map::new();
        edge_description_taffy_nodes
            .extend(first_level_nested_edge_taffy_nodes.edge_description_taffy_nodes);

        let thing_rank_container_style = TaffyContainerBuilder::taffy_container_style(
            node_layouts,
            &NodeInbuilt::ThingsContainer.id(),
            Size::auto(),
        );
        let tag_rank_container_style = TaffyContainerBuilder::taffy_container_style(
            node_layouts,
            &NodeInbuilt::TagsContainer.id(),
            Size::auto(),
        );
        let process_rank_container_style = TaffyContainerBuilder::taffy_container_style(
            node_layouts,
            &NodeInbuilt::ProcessesContainer.id(),
            Size::auto(),
        );

        let EdgeDescriptionBuildResult {
            edge_description_taffy_nodes: thing_edge_desc_taffy_nodes,
            position_to_container_ids: thing_position_to_container_ids,
        } = EdgeDescriptionBuilder::build(
            ctx,
            &mut taffy_tree,
            &EntityType::ThingDefault,
            None,
            &thing_rank_container_style,
        );
        let EdgeDescriptionBuildResult {
            edge_description_taffy_nodes: tag_edge_desc_taffy_nodes,
            position_to_container_ids: tag_position_to_container_ids,
        } = EdgeDescriptionBuilder::build(
            ctx,
            &mut taffy_tree,
            &EntityType::TagDefault,
            None,
            &tag_rank_container_style,
        );
        let EdgeDescriptionBuildResult {
            edge_description_taffy_nodes: process_edge_desc_taffy_nodes,
            position_to_container_ids: process_position_to_container_ids,
        } = EdgeDescriptionBuilder::build(
            ctx,
            &mut taffy_tree,
            &EntityType::ProcessDefault,
            None,
            &process_rank_container_style,
        );
        edge_description_taffy_nodes.extend(thing_edge_desc_taffy_nodes);
        edge_description_taffy_nodes.extend(tag_edge_desc_taffy_nodes);
        edge_description_taffy_nodes.extend(process_edge_desc_taffy_nodes);

        // === Build edge_description_container spacers === //
        //
        // For each edge that crosses through an edge_description_container at
        // the top level, insert a spacer inside that container so the edge
        // path can route around it. Must run before position_to_container_ids
        // is consumed by `rank_containers_for_first_level_nodes_build`.
        for (target_entity_type, position_to_container_ids) in [
            (&EntityType::ThingDefault, &thing_position_to_container_ids),
            (&EntityType::TagDefault, &tag_position_to_container_ids),
            (
                &EntityType::ProcessDefault,
                &process_position_to_container_ids,
            ),
        ] {
            for (edge_id, new_spacers) in EdgeSpacerBuilder::build_edge_desc_container_spacers(
                ctx,
                &mut taffy_tree,
                target_entity_type,
                None,
                position_to_container_ids,
                &edge_description_taffy_nodes,
            ) {
                edge_spacer_taffy_nodes
                    .entry(edge_id)
                    .or_default()
                    .edge_desc_container_spacer_taffy_node_ids
                    .extend(new_spacers.edge_desc_container_spacer_taffy_node_ids);
            }
        }

        // Create rank sub-containers for top-level nodes, mirroring the
        // rank-based child container logic used inside
        // `TaffyDiagramNodeBuilder::build_node_with_child_hierarchy`.
        //
        // Each entity type gets its own set of rank containers using the
        // style of its parent container, interleaved with any
        // edge_description_containers at the same level.
        let thing_rank_container_ids =
            TaffyContainerBuilder::rank_containers_for_first_level_nodes_build(
                &mut taffy_tree,
                thing_rank_to_taffy_ids,
                thing_rank_container_style,
                thing_position_to_container_ids,
            );
        let tag_rank_container_ids =
            TaffyContainerBuilder::rank_containers_for_first_level_nodes_build(
                &mut taffy_tree,
                tag_rank_to_taffy_ids,
                tag_rank_container_style,
                tag_position_to_container_ids,
            );
        let process_rank_container_ids =
            TaffyContainerBuilder::rank_containers_for_first_level_nodes_build(
                &mut taffy_tree,
                process_rank_to_taffy_ids,
                process_rank_container_style,
                process_position_to_container_ids,
            );

        let node_inbuilt_to_taffy = TaffyContainerBuilder::build(
            &mut taffy_tree,
            &mut taffy_id_to_node,
            node_layouts,
            dimension,
            &thing_rank_container_ids,
            &process_rank_container_ids,
            &tag_rank_container_ids,
        );

        let Some(root) = node_inbuilt_to_taffy.get(&NodeInbuilt::Root).copied() else {
            panic!("`root` node not present in `node_inbuilt_to_taffy`.");
        };

        // Pre-compute edge endpoint node IDs for edge label slot sizing.
        let edge_id_to_endpoint_node_ids = EdgeLabelBuilder::edge_id_to_node_ids_build(edge_groups);

        // Compute layout (size measurement only, no syntax highlighting)
        let mut node_measure_context = NodeMeasureContext {
            ctx,
            edge_labels,
            edge_id_to_endpoint_node_ids: &edge_id_to_endpoint_node_ids,
        };

        taffy_tree
            .compute_layout_with_measure(
                root,
                Size::<AvailableSpace> {
                    width: AvailableSpace::Definite(dimension.width()),
                    height: AvailableSpace::Definite(dimension.height()),
                },
                |known_dimensions, available_space, _taffy_node_id, taffy_node_ctx, style| {
                    node_measure_context.size_measure(
                        known_dimensions,
                        available_space,
                        taffy_node_ctx,
                        style,
                    )
                },
            )
            .expect("Expected layout computation to succeed.");

        // Merge collected edge label leaf nodes into the edge label taffy node
        // map now that all envelope nodes have been built.
        let edge_label_taffy_nodes =
            EdgeLabelBuilder::build(edge_label_leaves, edge_face_assignments, edge_groups);

        // Compute highlighted spans *after* layout is complete.
        //
        // This is done once per node instead of multiple times during layout
        // measurement
        let entity_highlighted_spans = HighlightedSpansComputer::compute(
            ctx,
            &taffy_tree,
            &node_id_to_taffy,
            &edge_label_taffy_nodes,
            edge_labels,
        );

        let edge_description_highlighted_spans =
            HighlightedSpansComputer::compute_edge_desc_containers(
                &taffy_tree,
                &edge_description_taffy_nodes,
                edge_descs,
                char_width,
                lod,
            );

        // Compute highlighted text spans and image spans for markdown nodes
        // after layout is complete.
        //
        // Markdown nodes are skipped by `HighlightedSpansComputer::compute`
        // (their `text_node_id` holds a `TaffyNodeCtx::MdToken` /
        // `TaffyNodeCtx::MdImage` rather than a `TaffyNodeCtx::DiagramNode`),
        // so their spans must be computed separately here and merged in.
        let (md_entity_spans, entity_image_spans) = MdSpansComputer::compute(
            &taffy_tree,
            &node_id_to_taffy,
            &md_node_taffy_ids,
            char_width,
        );

        let mut entity_highlighted_spans = entity_highlighted_spans;
        for (node_id, spans) in md_entity_spans.into_inner() {
            entity_highlighted_spans.insert(node_id, spans);
        }

        // Compute highlighted text spans and image spans for edge descriptions
        // that used the markdown path.
        //
        // Edge descriptions with markdown are skipped by
        // `HighlightedSpansComputer::compute_edge_desc_containers` (their
        // `md_node_taffy_ids` is `Some`), so their spans must be computed
        // separately here and merged in.
        let (md_edge_desc_spans, edge_description_image_spans) =
            MdSpansComputer::compute_edge_descs(
                &taffy_tree,
                &edge_description_taffy_nodes,
                char_width,
            );

        let mut edge_description_highlighted_spans = edge_description_highlighted_spans;
        for (edge_id, spans) in md_edge_desc_spans {
            edge_description_highlighted_spans.insert(edge_id, spans);
        }

        std::iter::once(TaffyNodeMappings {
            taffy_tree,
            node_inbuilt_to_taffy,
            node_id_to_taffy,
            taffy_id_to_node,
            edge_spacer_taffy_nodes,
            entity_highlighted_spans,
            edge_label_taffy_nodes,
            edge_description_taffy_nodes,
            edge_description_highlighted_spans,
            node_id_to_envelope_taffy_node,
            md_node_taffy_ids,
            entity_image_spans,
            edge_description_image_spans,
        })
    }
}