annatomic 0.4.0

use std::{
    cmp::Ordering,
    collections::{BTreeMap, BTreeSet, HashMap, HashSet},
    ops::RangeInclusive,
    sync::Arc,
};

use crate::app::{
    actions::{EditorAction, GraphAction},
    data::{Span, SpanRow, Token},
    editors::spans::{
        actions::{
            AddBaseTokenAfter, AddBaseTokenBefore, AddSpanAction, DeleteSelectedNode,
            EditWhitespaceAfter, EditWhitespaceBefore, ExpandNodeLeft, ExpandNodeRight,
            GoToFirstToken, GoToLastToken, Search, ShrinkNodeLeft, ShrinkNodeRight,
        },
        cache::LayoutCache,
    },
    messages::Notifier,
    project::EditorStateUpdates,
    util::token_helper::{TOKEN_KEY, TokenHelper},
    views::Editor,
    widgets::{
        filter::{FilterWidget, FilterWidgetOutput},
        search::{SearchWidget, SearchWidgetOutput},
        span::SpanWidget,
        token::{SelectionType, TokenWidget},
    },
};
use anyhow::{Context, Result};
use egui::{
    Button, Color32, Direction, Key, Layout, Modifiers, Pos2, Rangef, Rect, RichText, ScrollArea,
    Sense, Stroke, TextBuffer, TextEdit, Ui, UiBuilder, Widget, mutex::RwLock,
    text_edit::TextEditState,
};
use egui_phosphor::regular as icons;
use graphannis::{
    AnnotationGraph,
    graph::{AnnoKey, NodeID},
    model::AnnotationComponentType,
};
use graphannis_core::{
    annostorage::ValueSearch,
    graph::{ANNIS_NS, NODE_NAME_KEY},
};
use itertools::Itertools;

mod actions;
mod cache;
#[cfg(test)]
mod tests;

#[derive(Clone, Debug)]
enum NodeIndexType {
    Token {
        token_index: usize,
    },
    Segmentation {
        seg_name: String,
        token_index: usize,
    },
    Span {
        index_of_row: usize,
        position_in_row: usize,
    },
}

#[non_exhaustive]
#[derive(Clone, Debug, PartialEq)]
enum ActiveInteraction {
    EmptyNodeAdded {
        node_name: String,
    },
    AddingNewSpan {
        namespace: String,
        name: String,
        is_segmentation: bool,
    },
    EditWhitespaceBefore {
        node_name: String,
    },
    EditWhitespaceAfter {
        node_name: String,
    },
    Search,
}

#[derive(Clone)]
pub(crate) struct SpanEditor {
    parent_name: String,
    render_spans: bool,
    graph: Arc<RwLock<AnnotationGraph>>,
    token: Vec<Token>,
    segmentations: BTreeMap<String, Vec<Token>>,
    span_rows: Vec<SpanRow>,
    node_index_by_name: HashMap<String, NodeIndexType>,
    selected_nodes: HashSet<String>,
    pending_graph_actions: Vec<GraphAction>,
    notifier: Notifier,
    layout_cache: LayoutCache,
    selected_span_rows: Vec<usize>,
    scroll_to_token: Option<usize>,
    active_interaction: Option<ActiveInteraction>,
    keep_search_window_open: bool,
}

type TokenList = Vec<Token>;

fn get_fields_from_graph(
    parent_name: &str,
    graph: &AnnotationGraph,
) -> Result<(TokenList, BTreeMap<String, TokenList>, Vec<SpanRow>)> {
    let mut token = Vec::new();
    let mut token_to_index = HashMap::new();
    let mut segmentations = BTreeMap::new();
    let mut is_segmentation_node = HashSet::new();

    let tok_helper = TokenHelper::new(graph)?;
    let token_ids = tok_helper.get_ordered_token(parent_name, None)?;
    for (idx, node_id) in token_ids.iter().enumerate() {
        let t = Token::from_graph(*node_id, idx, idx, graph)?;
        token.push(t);
        token_to_index.insert(*node_id, idx);
    }

    // Find all ordering components other than the base layer
    for ordering_component in
        graph.get_all_components(Some(AnnotationComponentType::Ordering), None)
    {
        if ordering_component.layer != ANNIS_NS || !ordering_component.name.is_empty() {
            let token_ids =
                tok_helper.get_ordered_token(parent_name, Some(&ordering_component.name))?;
            for node_id in token_ids.iter() {
                let covered = tok_helper.covered_token(*node_id)?;
                let start = covered.first().and_then(|t| token_to_index.get(t));
                let end = covered.last().and_then(|t| token_to_index.get(t));
                if let (Some(start), Some(end)) = (start, end) {
                    let t = Token::from_graph(*node_id, *start, *end, graph)?;
                    is_segmentation_node.insert(*node_id);
                    segmentations
                        .entry(ordering_component.name.to_string())
                        .or_insert_with(Vec::default)
                        .push(t);
                }
            }
        }
    }

    // Get all segmentation token that don't belong to any of the existing ordering components
    // This is the case if there is only single segmentation token not connected to another one.
    for node_with_token_key in graph.get_node_annos().exact_anno_search(
        Some(&TOKEN_KEY.ns),
        &TOKEN_KEY.name,
        ValueSearch::Any,
    ) {
        let node_with_token_key = node_with_token_key?.node;
        // Not a base token and not already known from the ordering component
        if !token_to_index.contains_key(&node_with_token_key)
            && !is_segmentation_node.contains(&node_with_token_key)
            && let Some(seg) = tok_helper.get_segmentation_name(node_with_token_key)?
        {
            let covered = tok_helper.covered_token(node_with_token_key)?;
            let start = covered.first().and_then(|t| token_to_index.get(t));
            let end = covered.last().and_then(|t| token_to_index.get(t));
            if let (Some(start), Some(end)) = (start, end) {
                let t = Token::from_graph(node_with_token_key, *start, *end, graph)?;
                is_segmentation_node.insert(node_with_token_key);
                segmentations
                    .entry(seg)
                    .or_insert_with(Vec::default)
                    .push(t);
            }
        }
    }

    // Get all other nodes as spans
    let mut coverage_node_ids = HashSet::new();
    for t in token_ids {
        for gs in tok_helper.get_coverage_gs() {
            for n in gs.get_ingoing_edges(t) {
                let n = n?;
                coverage_node_ids.insert(n);
            }
        }
    }
    let mut dom_gs = Vec::new();
    for c in graph.get_all_components(Some(AnnotationComponentType::Dominance), None) {
        if let Some(gs) = graph.get_graphstorage(&c) {
            dom_gs.push(gs);
        }
    }

    // Create a row for each annotation key and assign the spans with the keys to this row
    let mut span_rows_by_key: BTreeMap<AnnoKey, Vec<SpanRow>> = BTreeMap::new();

    for n in coverage_node_ids {
        if !graph.get_node_annos().has_value_for_item(&n, &TOKEN_KEY)? {
            // Check if there is no outgoing dominance edge
            let mut has_outgoing_dom_rel = false;
            for gs in &dom_gs {
                if gs.has_outgoing_edges(n)? {
                    has_outgoing_dom_rel = true;
                    break;
                }
            }
            if !has_outgoing_dom_rel && !is_segmentation_node.contains(&n) {
                let span = Span::from_graph(n, &tok_helper, graph)?;
                let first_non_annis_key = span
                    .labels
                    .keys()
                    .find(|k| k.ns != ANNIS_NS)
                    .cloned()
                    .unwrap_or_else(|| AnnoKey {
                        name: "".into(),
                        ns: "".into(),
                    });
                let span_rows_for_key = span_rows_by_key.entry(first_non_annis_key).or_default();
                if span_rows_for_key.is_empty() {
                    // No rows yet, directly add the span as row
                    span_rows_for_key.push(span.into());
                } else {
                    // Check if we can add this span to any existing row
                    let mut merged = false;

                    for existing_row in span_rows_for_key.iter_mut() {
                        let was_merged = existing_row.merge_span(&span)?;
                        if was_merged {
                            // Span was merged into a row
                            merged = true;
                            break;
                        }
                    }

                    if !merged {
                        // Add another row and add this span to the new empty row
                        span_rows_for_key.push(span.into());
                    }
                }
            }
        }
    }
    let all_span_rows: Vec<_> = span_rows_by_key
        .into_values()
        .flatten()
        .filter(|row| !row.spans.is_empty())
        .collect();

    Ok((token, segmentations, all_span_rows))
}

fn calculate_node_index(
    token: &[Token],
    segmentations: &BTreeMap<String, Vec<Token>>,
    span_rows: &[SpanRow],
) -> HashMap<String, NodeIndexType> {
    let mut result: HashMap<String, NodeIndexType> = token
        .iter()
        .enumerate()
        .map(|(idx, t)| {
            (
                t.node_name.clone(),
                NodeIndexType::Token { token_index: idx },
            )
        })
        .collect();

    for (seg_name, seg_token) in segmentations.iter() {
        for (idx, t) in seg_token.iter().enumerate() {
            result.insert(
                t.node_name.clone(),
                NodeIndexType::Segmentation {
                    seg_name: seg_name.clone(),
                    token_index: idx,
                },
            );
        }
    }

    for (index_of_row, span_row) in span_rows.iter().enumerate() {
        for (position_in_row, s) in span_row.spans.iter().enumerate() {
            result.insert(
                s.node_name.clone(),
                NodeIndexType::Span {
                    index_of_row,
                    position_in_row,
                },
            );
        }
    }

    result
}

impl SpanEditor {
    pub fn create_from_graph(
        selected_corpus_node: String,
        render_spans: bool,
        graph: Arc<RwLock<AnnotationGraph>>,
        notifier: Notifier,
    ) -> Result<Self> {
        let graph_lock = graph.read();

        let (token, segmentations, span_rows) =
            get_fields_from_graph(&selected_corpus_node, &graph_lock)?;

        let node_index_by_name = calculate_node_index(&token, &segmentations, &span_rows);
        std::mem::drop(graph_lock);

        Ok(Self {
            parent_name: selected_corpus_node,
            render_spans,
            graph,
            layout_cache: LayoutCache::new(),
            token,
            node_index_by_name,
            segmentations,
            span_rows,
            selected_nodes: HashSet::new(),
            pending_graph_actions: Vec::new(),
            notifier,
            active_interaction: None,
            keep_search_window_open: false,
            scroll_to_token: None,
            selected_span_rows: Vec::default(),
        })
    }

    fn recreate_from_graph(&mut self) -> Result<()> {
        let graph_lock = self.graph.read();
        let (token, segmentations, span_rows) =
            get_fields_from_graph(&self.parent_name, &graph_lock)?;

        let node_index_by_name = calculate_node_index(&token, &segmentations, &span_rows);

        self.layout_cache = LayoutCache::new();
        self.token = token;
        self.node_index_by_name = node_index_by_name;
        self.segmentations = segmentations;
        self.span_rows = span_rows;
        Ok(())
    }

    fn get_labels_for_node_mut(
        &mut self,
        node_name: &str,
    ) -> Option<&mut BTreeMap<AnnoKey, String>> {
        if let Some(node_index) = self.node_index_by_name.get(node_name) {
            match node_index {
                NodeIndexType::Token { token_index } => {
                    self.token.get_mut(*token_index).map(|t| &mut t.labels)
                }
                NodeIndexType::Segmentation {
                    seg_name,
                    token_index,
                } => self
                    .segmentations
                    .get_mut(seg_name)
                    .and_then(|tokens| tokens.get_mut(*token_index))
                    .map(|t| &mut t.labels),
                NodeIndexType::Span {
                    index_of_row,
                    position_in_row,
                } => self
                    .span_rows
                    .get_mut(*index_of_row)
                    .and_then(|span_row| span_row.spans.get_mut(*position_in_row))
                    .map(|span| &mut span.labels),
            }
        } else {
            None
        }
    }

    /// Returns its position if the given node is a base token
    fn get_token_index_by_name(&self, node_name: &str) -> Option<usize> {
        match self.node_index_by_name.get(node_name) {
            Some(NodeIndexType::Token { token_index }) => Some(*token_index),
            _ => None,
        }
    }

    /// Returns its position if the given node is a base token
    fn get_token_index_by_id(&self, node_id: NodeID) -> Option<usize> {
        let graph_lock = self.graph.read();
        let node_name = graph_lock
            .get_node_annos()
            .get_value_for_item(&node_id, &NODE_NAME_KEY)
            .context("Missing node name");
        if let Some(Some(node_name)) = self.notifier.ok_or_report(node_name) {
            self.get_token_index_by_name(node_name.as_str())
        } else {
            None
        }
    }

    fn node_name_from_index_element(&self, element: &NodeIndexType) -> Option<String> {
        match element {
            NodeIndexType::Token { token_index } => {
                self.token.get(*token_index).map(|t| t.node_name.clone())
            }
            NodeIndexType::Segmentation {
                seg_name,
                token_index,
            } => self
                .segmentations
                .get(seg_name)
                .and_then(|row| row.get(*token_index))
                .map(|s| s.node_name.clone()),
            NodeIndexType::Span {
                index_of_row,
                position_in_row,
            } => self
                .span_rows
                .get(*index_of_row)
                .and_then(|row| row.spans.get(*position_in_row))
                .map(|span| span.node_name.clone()),
        }
    }

    fn selected_base_token_indices(&self) -> BTreeSet<usize> {
        self.selected_nodes
            .iter()
            .filter_map(|selected_node| self.get_token_index_by_name(selected_node))
            .collect()
    }

    fn show_single_token(&mut self, ui: &mut Ui, token_position: usize) -> Rect {
        let min_token_width = self.layout_cache.min_base_token_width(token_position);
        // Determine if we actually have to render the token, or if we can allocate empty invisible space
        let cached_size = self.layout_cache.cached_base_token_size(token_position);
        let estimated_widget_rect = cached_size.map(|cached| {
            let mut result = ui.cursor();
            result.set_width(cached.x);
            result.set_height(cached.y);
            result
        });

        let response = if let Some(cached_size) = cached_size
            && let Some(estimated_widget_rect) = estimated_widget_rect
            && !ui.is_rect_visible(estimated_widget_rect)
        {
            ui.allocate_exact_size(cached_size, Sense::empty()).1
        } else {
            let mut token_editor = TokenWidget::with_min_width(
                &self.token[token_position],
                self.token_selection_type(token_position),
                min_token_width,
                None,
            );
            if let Some(ActiveInteraction::EmptyNodeAdded { node_name }) = &self.active_interaction
                && node_name == &self.token[token_position].node_name
            {
                token_editor.request_edit_token_value(ui);
                self.active_interaction = None;
            } else if let Some(ActiveInteraction::EditWhitespaceBefore { node_name: n }) =
                &self.active_interaction
                && n == &self.token[token_position].node_name
            {
                token_editor.request_edit_whitespace_before_value(ui);
                self.active_interaction = None;
            } else if let Some(ActiveInteraction::EditWhitespaceAfter { node_name: n }) =
                &self.active_interaction
                && n == &self.token[token_position].node_name
            {
                token_editor.request_edit_whitespace_after_value(ui);
                self.active_interaction = None;
            }

            let response = token_editor.show(ui);
            let (widget_response, widget_output) = (response.response, response.inner);

            self.layout_cache
                .add_token_rect(token_position, widget_response.rect, &self.token);

            if widget_output.add_token_after_action {
                AddBaseTokenAfter::perform(self);
            }

            if widget_output.selected {
                let select_range = ui.ctx().input(|i| i.modifiers.shift_only());
                let keep_selection = ui.ctx().input(|i| i.modifiers.command_only());
                if select_range {
                    self.select_base_token_range(token_position);
                } else {
                    if !keep_selection {
                        self.selected_nodes.clear();
                    }
                    self.selected_nodes
                        .insert(self.token[token_position].node_name.clone());
                }
                self.scroll_to_token = Some(token_position);
            }

            for anno in widget_output.changed_annotations.into_iter() {
                // create editor actions
                if anno.val.is_empty() {
                    self.pending_graph_actions.push(GraphAction::AnnoRemoved {
                        node_name: self.token[token_position].node_name.clone(),
                        anno_key: anno.key,
                    });
                } else {
                    self.pending_graph_actions
                        .push(GraphAction::AnnoValueChanged {
                            node_name: self.token[token_position].node_name.clone(),
                            anno,
                        });
                }
            }

            widget_response
        };

        if self.scroll_to_token == Some(token_position) && self.layout_cache.is_valid() {
            if !ui.clip_rect().contains_rect(response.rect) {
                response.scroll_to_me(None);
            }
            self.scroll_to_token = None;
        }

        response.rect
    }

    fn show_segmentation_layers(
        &mut self,
        ui: &mut Ui,
        token_offset_to_rect: &[Option<Rect>],
        highlighted_token_positions: &BTreeSet<usize>,
        mut current_span_offset: f32,
    ) -> f32 {
        let ui_style = ui.style().clone();
        let mut any_annotation_changed = false;

        let segmentation_keys: Vec<_> = self.segmentations.keys().cloned().collect();
        for (seg_index, seg_name) in segmentation_keys.iter().enumerate() {
            let mut max_node_height = 0.0;
            if let Some(seg_token) = self.segmentations.get_mut(seg_name) {
                for current_seg_token in seg_token.iter_mut() {
                    // Get the base token covered by this span and use them to create a rectangle
                    let mut covered_span = Rangef::NOTHING;

                    for token_rect in token_offset_to_rect
                        .iter()
                        .take(current_seg_token.end + 1)
                        .skip(current_seg_token.start)
                        .flatten()
                    {
                        covered_span.min = covered_span.min.min(token_rect.left());
                        covered_span.max = covered_span.max.max(token_rect.right());
                    }
                    if covered_span.span() > 0.0 {
                        let min_pos = Pos2::new(covered_span.min, current_span_offset);
                        let max_pos = Pos2::new(covered_span.max, current_span_offset);
                        let segmentation_rectangle = Rect::from_min_max(min_pos, max_pos);

                        if !self.layout_cache.is_valid()
                            || ui.is_rect_visible(segmentation_rectangle)
                        {
                            let selected =
                                if self.selected_nodes.contains(&current_seg_token.node_name) {
                                    SelectionType::Selected
                                } else {
                                    SelectionType::None
                                };
                            let mut segmentation_editor = TokenWidget::with_exact_width(
                                current_seg_token,
                                selected,
                                Some(segmentation_rectangle.width()),
                                Some(seg_name.clone()),
                            );

                            if let Some(ActiveInteraction::EmptyNodeAdded { node_name: n }) =
                                &self.active_interaction
                                && n == &current_seg_token.node_name
                            {
                                segmentation_editor.request_edit_token_value(ui);
                                self.active_interaction = None;
                            } else if let Some(ActiveInteraction::EditWhitespaceBefore {
                                node_name: n,
                            }) = &self.active_interaction
                                && n == &current_seg_token.node_name
                            {
                                segmentation_editor.request_edit_whitespace_before_value(ui);
                                self.active_interaction = None;
                            } else if let Some(ActiveInteraction::EditWhitespaceAfter {
                                node_name: n,
                            }) = &self.active_interaction
                                && n == &current_seg_token.node_name
                            {
                                segmentation_editor.request_edit_whitespace_after_value(ui);
                                self.active_interaction = None;
                            }

                            let ui_builder = UiBuilder::new()
                                .max_rect(segmentation_rectangle)
                                .layout(Layout::centered_and_justified(Direction::TopDown));
                            let ui_response =
                                ui.scope_builder(ui_builder, |ui| segmentation_editor.show(ui));

                            let (ui_response, widget_outout) =
                                (ui_response.response, ui_response.inner.inner);

                            max_node_height = ui_response.rect.height().max(max_node_height);

                            handle_widget_selection_event(
                                &mut self.selected_nodes,
                                selected,
                                widget_outout.selected,
                                &current_seg_token.node_name,
                                ui,
                            );

                            for anno in widget_outout.changed_annotations.into_iter() {
                                any_annotation_changed = true;
                                ui_response.scroll_to_me(Some(egui::Align::Center));

                                if anno.val.is_empty() {
                                    self.pending_graph_actions.push(GraphAction::AnnoRemoved {
                                        node_name: current_seg_token.node_name.clone(),
                                        anno_key: anno.key,
                                    });
                                } else {
                                    self.pending_graph_actions.push(
                                        GraphAction::AnnoValueChanged {
                                            node_name: current_seg_token.node_name.clone(),
                                            anno,
                                        },
                                    );
                                }
                            }

                            let widget_width = ui_response
                                .intrinsic_size
                                .map(|s| s.x)
                                .unwrap_or_else(|| ui_response.rect.width());

                            self.layout_cache.add_segmentation_width(
                                seg_name,
                                current_seg_token,
                                widget_width,
                                &self.token,
                            );
                        }
                    }
                }
            }
            // Interface to add a new segmentation, when a token is highlighed
            if let Some(min_position) = highlighted_token_positions.first()
                && let Some(max_position) = highlighted_token_positions.last()
                && let Some(Some(min_token_offset)) = token_offset_to_rect.get(*min_position)
                && let Some(Some(max_token_offset)) = token_offset_to_rect.get(*max_position)
            {
                let min_pos = Pos2::new(min_token_offset.left(), current_span_offset);
                let max_pos = Pos2::new(max_token_offset.right(), current_span_offset);
                let span_rectangle = Rect::from_min_max(min_pos, max_pos);

                let ui_builder = UiBuilder::new()
                    .max_rect(span_rectangle)
                    .layout(Layout::centered_and_justified(Direction::TopDown));
                ui.scope_builder(ui_builder, |ui| {
                    let button_text = if seg_index < 10 {
                        format!("{} - {}", seg_index + 1, &seg_name)
                    } else {
                        seg_name.clone()
                    };
                    let bt = Button::new(RichText::new(button_text).size(16.0))
                        .stroke(Stroke::new(3.0, Color32::BLACK))
                        .corner_radius(12.0)
                        .wrap_mode(egui::TextWrapMode::Truncate)
                        .ui(ui);
                    if bt.clicked() {
                        self.add_segmentation_for_selection(seg_index);
                    }
                });
            }
            current_span_offset += max_node_height + ui_style.spacing.item_spacing.y;
        }
        if any_annotation_changed {
            // Reset layout information
            self.layout_cache = LayoutCache::new();
        }
        current_span_offset
    }

    fn show_spans(
        &mut self,
        ui: &mut Ui,
        selected_node_annos: &BTreeSet<AnnoKey>,
        token_offset_to_rect: &[Option<Rect>],
        highlighted_token_positions: &BTreeSet<usize>,
        mut current_span_offset: f32,
    ) {
        let ui_style = ui.style().clone();

        self.selected_span_rows.clear();
        for row_idx in 0..self.span_rows.len() {
            // Assume a certail default height if the span was not rendered
            let mut max_node_height = 31.0;

            // Ignore span rows that don't contain the selected node annos
            if self.span_rows[row_idx]
                .anno_keys
                .intersection(selected_node_annos)
                .next()
                .is_none()
            {
                continue;
            }

            for current_span in self.span_rows[row_idx].spans.iter() {
                if current_span
                    .labels
                    .keys()
                    .any(|k| selected_node_annos.contains(k))
                {
                    // Get the base token covered by this span and use them to create a rectangle
                    let mut covered_token_span = Rangef::NOTHING;
                    if let (Some(start_id), Some(end_id)) = (
                        current_span.sorted_covered_token_ids.first(),
                        current_span.sorted_covered_token_ids.last(),
                    ) {
                        let current_span_start = self.get_token_index_by_id(*start_id).unwrap_or(0);
                        let current_span_end = self
                            .get_token_index_by_id(*end_id)
                            .unwrap_or(current_span_start);
                        for token_rect in token_offset_to_rect
                            .iter()
                            .take(current_span_end + 1)
                            .skip(current_span_start)
                            .flatten()
                        {
                            covered_token_span.min = covered_token_span.min.min(token_rect.left());
                            covered_token_span.max = covered_token_span.max.max(token_rect.right());
                        }
                    }

                    if covered_token_span.span() > 0.0 {
                        // The resulting rectangle will have a fixed left/top
                        // point and width, but a the height will be determined
                        // by the widget itself.
                        let span_min_rect = Rect::from_x_y_ranges(
                            covered_token_span,
                            Rangef::new(current_span_offset, current_span_offset),
                        );
                        let span_max_rect = Rect::from_x_y_ranges(
                            covered_token_span,
                            Rangef::new(current_span_offset, f32::INFINITY),
                        );

                        if ui.is_rect_visible(span_min_rect) {
                            let selected = if self.selected_nodes.contains(&current_span.node_name)
                            {
                                SelectionType::Selected
                            } else {
                                SelectionType::None
                            };
                            let mut span_editor = SpanWidget::new(current_span, selected);

                            if let Some(ActiveInteraction::EmptyNodeAdded { node_name }) =
                                &self.active_interaction
                                && node_name == &current_span.node_name
                            {
                                span_editor.request_edit_label_value(ui);
                                self.active_interaction = None;
                            }

                            let widget_response = ui.scope_builder(
                                UiBuilder::new().max_rect(span_max_rect),
                                move |ui| span_editor.show(ui).inner,
                            );

                            let (widget_response, widget_output) =
                                (widget_response.response, widget_response.inner);

                            for anno in widget_output.changed_annotations {
                                self.pending_graph_actions
                                    .push(GraphAction::AnnoValueChanged {
                                        node_name: current_span.node_name.clone(),
                                        anno,
                                    });
                            }

                            handle_widget_selection_event(
                                &mut self.selected_nodes,
                                selected,
                                widget_output.selected,
                                &current_span.node_name,
                                ui,
                            );

                            max_node_height = widget_response.rect.height().max(max_node_height);
                        }
                    }
                }
            }

            self.selected_span_rows.push(row_idx);

            // Add the button to add a new span to this row
            if let Some(min_position) = highlighted_token_positions.first()
                && let Some(max_position) = highlighted_token_positions.last()
                && let Some(Some(min_token_offset)) = token_offset_to_rect.get(*min_position)
                && let Some(Some(max_token_offset)) = token_offset_to_rect.get(*max_position)
            {
                let min_pos = Pos2::new(min_token_offset.left(), current_span_offset);
                let max_pos = Pos2::new(max_token_offset.right(), current_span_offset);
                let span_rectangle = Rect::from_min_max(min_pos, max_pos);

                let ui_builder = UiBuilder::new()
                    .max_rect(span_rectangle)
                    .layout(Layout::centered_and_justified(Direction::TopDown));
                ui.scope_builder(ui_builder, |ui| {
                    let anno_key_text = self.span_rows[row_idx]
                        .anno_keys
                        .iter()
                        .filter(|k| k.ns != ANNIS_NS)
                        .map(|k| &k.name)
                        .join(", ");
                    let button_text =
                        if (self.segmentations.len() + self.selected_span_rows.len()) < 10 {
                            format!(
                                "{} - {}",
                                self.segmentations.len() + self.selected_span_rows.len(),
                                &anno_key_text
                            )
                        } else {
                            anno_key_text.clone()
                        };
                    let bt = Button::new(RichText::new(button_text).size(16.0))
                        .stroke(Stroke::new(3.0, Color32::BLACK))
                        .corner_radius(12.0)
                        .wrap_mode(egui::TextWrapMode::Truncate)
                        .ui(ui)
                        .on_hover_text(anno_key_text);
                    if bt.clicked() {
                        self.add_span_for_row(self.selected_span_rows.len() - 1);
                    }
                });
            }

            // Always add an offset even if the row was not rendered.
            current_span_offset += max_node_height + ui_style.spacing.item_spacing.y;
        }
    }

    fn show_new_span_dialog(&mut self, ui: &egui::Ui) {
        let mut accepted_anno = None;
        if let Some(ActiveInteraction::AddingNewSpan {
            namespace,
            name,
            is_segmentation,
        }) = self.active_interaction.as_mut()
        {
            let response = egui::Modal::new("adding_new_span".into()).show(ui.ctx(), |ui| {
                let text_field_response = ui
                    .horizontal(|ui| {
                        let text_ns = TextEdit::singleline(namespace)
                            .hint_text("Namespace")
                            .horizontal_align(egui::Align::Max)
                            .desired_width(80.0);
                        ui.add_enabled(!(*is_segmentation), text_ns);
                        ui.label("::");

                        let text_name = TextEdit::singleline(name)
                            .hint_text("Name")
                            .desired_width(80.0)
                            .ui(ui);

                        if text_name.lost_focus() && ui.input(|i| i.key_pressed(egui::Key::Enter)) {
                            Some(true)
                        } else {
                            None
                        }
                    })
                    .inner;
                ui.checkbox(is_segmentation, "Is segmentation");
                ui.separator();
                let button_response = ui
                    .horizontal(|ui| {
                        let add_button = ui.button("Add");
                        let cancel_button = ui.button("Cancel");
                        if add_button.clicked() {
                            Some(true)
                        } else if cancel_button.clicked() {
                            Some(false)
                        } else {
                            None
                        }
                    })
                    .inner;
                text_field_response.or(button_response)
            });
            if let Some(accepted) = response.inner {
                if accepted {
                    accepted_anno = Some((
                        AnnoKey {
                            name: name.as_str().into(),
                            ns: namespace.as_str().into(),
                        },
                        *is_segmentation,
                    ));
                } else {
                    self.active_interaction = None;
                }
            }
        }

        if let Some((anno_key, is_segmentation)) = accepted_anno.take() {
            let selected_token: HashSet<_> = self
                .selected_nodes
                .iter()
                .filter(|selected_node| self.get_token_index_by_name(selected_node).is_some())
                .cloned()
                .collect();

            if is_segmentation {
                self.pending_graph_actions
                    .push(GraphAction::AddSegmentationSpan {
                        segmentation: anno_key.name.to_string(),
                        selected_token,
                        parent_name: self.parent_name.clone(),
                    });
            } else {
                self.pending_graph_actions.push(GraphAction::AddSpan {
                    selected_token,
                    parent_name: self.parent_name.clone(),
                    anno_key,
                });
            }

            self.active_interaction = None;
        }
    }

    fn show_search_dialog(&mut self, ui: &egui::Ui) {
        if let Some(ActiveInteraction::Search) = &self.active_interaction {
            let search_widget = SearchWidget::new("document-search-widget");
            let window_response = egui::Window::new("Search")
                .open(&mut self.keep_search_window_open)
                .collapsible(false)
                .show(ui.ctx(), |ui| search_widget.show(ui))
                .and_then(|response| response.inner);
            if let Some(widget_response) = window_response {
                let widget_output = widget_response.inner;

                // Try to find matches for the search input
                match widget_output {
                    SearchWidgetOutput::TokenIndex(position) => {
                        if let Some(t) = self.token.get(position) {
                            self.selected_nodes.clear();
                            self.selected_nodes.insert(t.node_name.clone());
                            self.scroll_to_token = Some(position);
                        }
                    }
                    SearchWidgetOutput::NodeName(search_text) => {
                        // Find the first node that matches the suffix
                        if !search_text.trim().is_empty()
                            && let Some((node_name, node_index)) = self
                                .node_index_by_name
                                .iter()
                                .find(|(key, _value)| key.ends_with(&search_text))
                        {
                            self.selected_nodes.clear();
                            self.selected_nodes.insert(node_name.clone());

                            let covered_token = self.get_token_range(node_index);
                            let range_length = covered_token.end() - covered_token.start();
                            let offset = ((range_length as f64) / 2.0).round() as usize;
                            let center_token = covered_token.start() + offset;
                            self.scroll_to_token = Some(center_token);
                        }
                    }
                    SearchWidgetOutput::CloseRequested => {
                        self.active_interaction = None;
                        self.keep_search_window_open = false;
                    }
                    SearchWidgetOutput::None => {}
                }
            }
            if !self.keep_search_window_open {
                self.active_interaction = None;
            }
        }
    }

    /// Get the token range a given node covers.
    fn get_token_range(&self, node: &NodeIndexType) -> RangeInclusive<usize> {
        match node {
            NodeIndexType::Token { token_index } => *token_index..=*token_index,
            NodeIndexType::Segmentation {
                seg_name,
                token_index,
            } => {
                if let Some(seg) = self.segmentations.get(seg_name) {
                    seg[*token_index].start..=seg[*token_index].end
                } else {
                    0..=0
                }
            }
            NodeIndexType::Span {
                index_of_row,
                position_in_row,
            } => {
                let span = &self.span_rows[*index_of_row].spans[*position_in_row];
                let start_idx = span
                    .sorted_covered_token_ids
                    .first()
                    .and_then(|id| self.get_token_index_by_id(*id))
                    .unwrap_or_default();
                let end_idx = span
                    .sorted_covered_token_ids
                    .last()
                    .and_then(|id| self.get_token_index_by_id(*id))
                    .unwrap_or_default();
                start_idx..=end_idx
            }
        }
    }

    fn handle_selection_key_press_event(
        &mut self,
        direction: Direction,
        modifiers: Modifiers,
    ) -> anyhow::Result<()> {
        let keep_selection = modifiers.shift_only();

        // Map the selected nodes to their node index types
        let mut selected: Vec<_> = self
            .selected_nodes
            .iter()
            .filter_map(|n| self.node_index_by_name.get(n))
            .cloned()
            .collect();

        // Sort them according to their covered token range
        selected.sort_by(|a, b| {
            let a_range = self.get_token_range(a);
            let b_range = self.get_token_range(b);
            a_range.cmp(b_range)
        });

        if let Some(first_selected) = selected.first()
            && let Some(last_selected) = selected.last()
        {
            // Get the node index types that should be selected next based on the direction and the current selection.
            let mut next_selected_indices = match direction {
                Direction::LeftToRight => selected
                    .last()
                    .and_then(|right_most_selected| self.next_node(right_most_selected))
                    .map(|n| vec![n])
                    .unwrap_or_default(),
                Direction::RightToLeft => selected
                    .first()
                    .and_then(|left_most_selected| self.previous_node(left_most_selected))
                    .map(|n| vec![n])
                    .unwrap_or_default(),
                Direction::TopDown => self.nodes_below(&selected),
                Direction::BottomUp => {
                    if modifiers.command_only() {
                        // Directly select the base tokens
                        let full_selected_token_range =
                            *self.get_token_range(first_selected).start()
                                ..=*self.get_token_range(last_selected).end();
                        full_selected_token_range
                            .filter_map(|token_position: usize| self.token.get(token_position))
                            .filter_map(|t| self.node_index_by_name.get(&t.node_name))
                            .cloned()
                            .collect()
                    } else {
                        self.nodes_above(&selected)
                    }
                }
            };

            // Sort the selection, so we can get the token range more easily
            next_selected_indices.sort_by(|a, b| {
                let a_range = self.get_token_range(a);
                let b_range = self.get_token_range(b);
                a_range.cmp(b_range)
            });

            // Only change the selection if we actually found new selected nodes
            if let Some(first_selected) = next_selected_indices.first()
                && let Some(last_selected) = next_selected_indices.last()
            {
                // Find the token index to scroll to by calculating the full
                // token range and retrieving the first/last token
                let full_selected_token_range = *self.get_token_range(first_selected).start()
                    ..=*self.get_token_range(last_selected).end();

                // Depending on the direction, select the first or last covered token
                if direction == Direction::RightToLeft {
                    self.scroll_to_token = Some(*full_selected_token_range.start());
                } else {
                    self.scroll_to_token = Some(*full_selected_token_range.end());
                }

                if !keep_selection {
                    self.selected_nodes.clear();
                }
                let next_selected_node_names: Vec<_> = next_selected_indices
                    .into_iter()
                    .filter_map(|element| self.node_name_from_index_element(&element))
                    .collect();
                self.selected_nodes.extend(next_selected_node_names);
            }
        }

        Ok(())
    }

    /// Find all segmentation nodes from a layer that overlap with the token range.
    fn segmentation_nodes_by_token_overlap(
        &self,
        covered_token_range: &RangeInclusive<usize>,
        seg_name: &str,
    ) -> HashSet<String> {
        let mut result = HashSet::new();

        let given_start = *covered_token_range.start();
        let given_end = *covered_token_range.end();

        if let Some(row) = self.segmentations.get(seg_name) {
            result.extend(
                row.iter()
                    .filter(|candidate| {
                        let candidate_start = candidate.start;
                        let candidate_end = candidate.end;
                        // Include spans where the start or end nodes are in the
                        // range (overlap) or where the range is fully included
                        // in a candidate
                        (given_start <= candidate_start && candidate_start <= given_end)
                            || (given_start <= candidate_end && candidate_end <= given_end)
                            || (given_start >= candidate_start && given_end <= candidate_end)
                    })
                    .map(|span| span.node_name.clone()),
            );
        }

        result
    }

    /// Find all span nodes on the span row that overlap with the token range.
    fn span_nodes_by_token_overlap(
        &self,
        covered_token_range: &RangeInclusive<usize>,
        row_idx: usize,
    ) -> HashSet<String> {
        let mut result = HashSet::new();

        let given_start = *covered_token_range.start();
        let given_end = *covered_token_range.end();

        if let Some(row) = self.span_rows.get(row_idx) {
            result.extend(
                row.spans
                    .iter()
                    .filter(|candidate| {
                        let candidate_start = candidate
                            .sorted_covered_token_ids
                            .first()
                            .and_then(|id| self.get_token_index_by_id(*id))
                            .unwrap_or_default();
                        let candidate_end = candidate
                            .sorted_covered_token_ids
                            .last()
                            .and_then(|id| self.get_token_index_by_id(*id))
                            .unwrap_or_default();
                        // Include spans where the start or end nodes are in the
                        // range (overlap) or where the range is fully included
                        // in a candidate
                        (given_start <= candidate_start && candidate_start <= given_end)
                            || (given_start <= candidate_end && candidate_end <= given_end)
                            || (given_start >= candidate_start && given_end <= candidate_end)
                    })
                    .map(|span| span.node_name.clone()),
            );
        }

        result
    }

    fn select_base_token_range(&mut self, new_selected_position: usize) {
        // Mark a range of token, find a suitable token as start for the range first
        let already_selected_indices = self.selected_base_token_indices();

        if let (Some(first), Some(last)) = (
            already_selected_indices.first(),
            already_selected_indices.last(),
        ) {
            if new_selected_position > *last {
                for i in *last..=new_selected_position {
                    self.selected_nodes.insert(self.token[i].node_name.clone());
                }
            } else if new_selected_position < *first {
                for i in new_selected_position..*first {
                    self.selected_nodes.insert(self.token[i].node_name.clone());
                }
            }
        } else {
            // If there is no current selection, just select the single token
            self.selected_nodes
                .insert(self.token[new_selected_position].node_name.clone());
        }
    }

    /// Adds an empty segmentation node that spans the currently selected token.
    ///
    /// - `layer_idx` The segmentation layer to add the new node to. **Starts with 0.**
    fn add_segmentation_for_selection(&mut self, layer_idx: usize) {
        if let Some((seg_name, _token)) = self.segmentations.iter().nth(layer_idx)
            && !self.selected_nodes.is_empty()
        {
            // Apply changes to internal data model
            {
                let graph = self.graph.read();
                if let Ok(tok_helper) = TokenHelper::new(&graph) {
                    // Schedule an update of the underlaying graph
                    let selected_token: HashSet<_> = self
                        .selected_nodes
                        .iter()
                        .filter_map(|node_name| {
                            let node_id = self.get_node_id_from_name(node_name, &graph)?;
                            if tok_helper.is_token(node_id).unwrap_or_default() {
                                Some(node_name.to_string())
                            } else {
                                None
                            }
                        })
                        .collect();

                    self.pending_graph_actions
                        .push(GraphAction::AddSegmentationSpan {
                            segmentation: seg_name.clone(),
                            selected_token,
                            parent_name: self.parent_name.clone(),
                        });
                }
            }
        }
    }

    /// Add an empty span that covers the currently selected tokens. The span
    /// will have an empty annotation value matching the annotations from the
    /// given the row index (starting from 0).
    fn add_span_for_row(&mut self, visible_row_idx: usize) {
        if let Some(actual_row_idx) = self.selected_span_rows.get(visible_row_idx)
            && let Some(selected_row) = self.span_rows.get(*actual_row_idx)
        {
            for anno_key in &selected_row.anno_keys {
                self.pending_graph_actions.push(GraphAction::AddSpan {
                    selected_token: self.selected_nodes.clone(),
                    parent_name: self.parent_name.clone(),
                    anno_key: anno_key.clone(),
                });
            }
        }
    }

    fn token_selection_type(&self, token_position: usize) -> SelectionType {
        let node_name = self
            .token
            .get(token_position)
            .map(|t| t.node_name.as_str())
            .unwrap_or_default();
        if self.selected_nodes.contains(node_name) {
            let mut first = true;
            let mut last = true;

            if let Some(token_index) = self.get_token_index_by_name(node_name) {
                if let Some(t) = token_index
                    .checked_sub(1)
                    .and_then(|token_index| self.token.get(token_index))
                    && self.selected_nodes.contains(&t.node_name)
                {
                    first = false;
                }

                if let Some(t) = self.token.get(token_index + 1)
                    && self.selected_nodes.contains(&t.node_name)
                {
                    last = false;
                }
            }
            if first && last {
                SelectionType::Selected
            } else if first {
                SelectionType::SelectionStart
            } else if last {
                SelectionType::SelectionEnd
            } else {
                SelectionType::Selected
            }
        } else {
            SelectionType::None
        }
    }

    fn get_node_id_from_name(&self, node_name: &str, graph: &AnnotationGraph) -> Option<NodeID> {
        let result = graph.get_node_annos().get_node_id_from_name(node_name);
        let result = self
            .notifier
            .ok_or_report(result.context("Missing node name annotation."));
        if let Some(Some(node_id)) = result {
            Some(node_id)
        } else {
            None
        }
    }

    fn next_node(&self, node: &NodeIndexType) -> Option<NodeIndexType> {
        match node {
            NodeIndexType::Token { token_index } => token_index
                .checked_add(1)
                .filter(|idx| *idx < self.token.len())
                .map(|token_index| NodeIndexType::Token { token_index }),
            NodeIndexType::Segmentation {
                seg_name,
                token_index,
            } => token_index
                .checked_add(1)
                .filter(|idx| {
                    self.segmentations
                        .get(seg_name)
                        .map(|seg| *idx < seg.len())
                        .unwrap_or_default()
                })
                .map(|token_index| NodeIndexType::Segmentation {
                    seg_name: seg_name.clone(),
                    token_index,
                }),
            NodeIndexType::Span {
                index_of_row,
                position_in_row,
            } => {
                // Find the token right to our span
                let graph = self.graph.read();
                if let Some(last_token_id) = self.span_rows[*index_of_row].spans[*position_in_row]
                    .sorted_covered_token_ids
                    .last()
                    && let Some(last_token_index) = self.get_token_index_by_id(*last_token_id)
                {
                    for next_token_index in (last_token_index + 1)..self.token.len() {
                        if let Some(next_token) = self.token.get(next_token_index)
                            && let Some(next_token_id) =
                                self.get_node_id_from_name(&next_token.node_name, &graph)
                        {
                            // Get the span that covers this next token index
                            if let Some((next_span_position, _)) =
                                self.span_rows[*index_of_row].spans.iter().enumerate().find(
                                    |(_, span)| {
                                        span.sorted_covered_token_ids.contains(&next_token_id)
                                    },
                                )
                            {
                                return Some(NodeIndexType::Span {
                                    index_of_row: *index_of_row,
                                    position_in_row: next_span_position,
                                });
                            }
                        };
                    }
                }
                None
            }
        }
    }

    fn previous_node(&self, node: &NodeIndexType) -> Option<NodeIndexType> {
        match node {
            NodeIndexType::Token { token_index } => token_index
                .checked_sub(1)
                .map(|token_index| NodeIndexType::Token { token_index }),
            NodeIndexType::Segmentation {
                seg_name,
                token_index,
            } => token_index
                .checked_sub(1)
                .map(|token_index| NodeIndexType::Segmentation {
                    seg_name: seg_name.clone(),
                    token_index,
                }),
            NodeIndexType::Span {
                index_of_row,
                position_in_row,
            } => {
                // Find the token left to our span
                let graph = self.graph.read();
                if let Some(first_token_id) = self.span_rows[*index_of_row].spans[*position_in_row]
                    .sorted_covered_token_ids
                    .first()
                    && let Some(first_token_index) = self.get_token_index_by_id(*first_token_id)
                {
                    for previous_token_index in (0..first_token_index).rev() {
                        if let Some(previous_token) = self.token.get(previous_token_index)
                            && let Some(previous_token_id) =
                                self.get_node_id_from_name(&previous_token.node_name, &graph)
                        {
                            // Get the span that covers this next token index
                            if let Some((previous_position, _)) =
                                self.span_rows[*index_of_row].spans.iter().enumerate().find(
                                    |(_, span)| {
                                        span.sorted_covered_token_ids.contains(&previous_token_id)
                                    },
                                )
                            {
                                return Some(NodeIndexType::Span {
                                    index_of_row: *index_of_row,
                                    position_in_row: previous_position,
                                });
                            }
                        }
                    }
                }
                None
            }
        }
    }

    /// For each layer that comes next, find a node where the token coverage
    /// overlaps with the one of the selected nodes. The nodes must be sorted
    /// according to their covered token.
    fn nodes_below(&self, sorted_nodes: &[NodeIndexType]) -> Vec<NodeIndexType> {
        let mut selected_node_names = HashSet::new();
        let selected_span_rows: HashSet<usize> = self.selected_span_rows.iter().copied().collect();

        if let Some(first_selected) = sorted_nodes.first()
            && let Some(last_selected) = sorted_nodes.last()
        {
            let token_range = *self.get_token_range(first_selected).start()
                ..=*self.get_token_range(last_selected).end();
            let segmentation_names: Vec<_> = self.segmentations.keys().cloned().collect();
            let start_seg_idx = sorted_nodes
                .iter()
                .filter_map(|node| match node {
                    NodeIndexType::Token { .. } => Some(0),
                    NodeIndexType::Segmentation { seg_name, .. } => segmentation_names
                        .iter()
                        .position(|probe| probe == seg_name)
                        .map(|idx| idx + 1),
                    _ => None,
                })
                .max();

            if let Some(start_seg_idx) = start_seg_idx {
                for seg_name in segmentation_names.iter().skip(start_seg_idx) {
                    let found_nodes =
                        self.segmentation_nodes_by_token_overlap(&token_range, seg_name);
                    if !found_nodes.is_empty() {
                        selected_node_names = found_nodes;
                        break;
                    }
                }
            }
            if selected_node_names.is_empty() {
                let start_span_idx = sorted_nodes
                    .iter()
                    .map(|node| match node {
                        NodeIndexType::Token { .. } | NodeIndexType::Segmentation { .. } => 0,
                        NodeIndexType::Span { index_of_row, .. } => index_of_row + 1,
                    })
                    .max()
                    .unwrap_or_default();
                for row_idx in start_span_idx..self.span_rows.len() {
                    if selected_span_rows.contains(&row_idx) {
                        let found_nodes = self.span_nodes_by_token_overlap(&token_range, row_idx);
                        if !found_nodes.is_empty() {
                            selected_node_names = found_nodes;
                            break;
                        }
                    }
                }
            }
        }

        selected_node_names
            .into_iter()
            .filter_map(|n| self.node_index_by_name.get(&n))
            .cloned()
            .collect()
    }

    /// For each layer that comes before, find a node where the token coverage
    /// overlaps with the one of the selected nodes. The nodes must be sorted
    /// according to their covered token.
    fn nodes_above(&self, sorted_nodes: &[NodeIndexType]) -> Vec<NodeIndexType> {
        let mut selected_node_names = HashSet::new();
        let selected_span_rows: HashSet<usize> = self.selected_span_rows.iter().copied().collect();

        if let Some(first_selected) = sorted_nodes.first()
            && let Some(last_selected) = sorted_nodes.last()
        {
            let token_range = *self.get_token_range(first_selected).start()
                ..=*self.get_token_range(last_selected).end();

            let start_span_idx = sorted_nodes
                .iter()
                .filter_map(|node| match node {
                    NodeIndexType::Token { .. } | NodeIndexType::Segmentation { .. } => None,
                    NodeIndexType::Span { index_of_row, .. } => index_of_row.checked_sub(1),
                })
                .min();
            if let Some(start_span_idx) = start_span_idx {
                for row_idx in (0..=start_span_idx).rev() {
                    if selected_span_rows.contains(&row_idx) {
                        let found_nodes = self.span_nodes_by_token_overlap(&token_range, row_idx);
                        if !found_nodes.is_empty() {
                            selected_node_names = found_nodes;
                            break;
                        }
                    }
                }
            }
            if selected_node_names.is_empty() {
                let segmentation_names: Vec<_> = self.segmentations.keys().cloned().collect();
                let start_seg_idx = sorted_nodes
                    .iter()
                    .filter_map(|node| match node {
                        NodeIndexType::Token { .. } => None,
                        NodeIndexType::Segmentation { seg_name, .. } => segmentation_names
                            .iter()
                            .position(|probe| probe == seg_name)
                            .and_then(|idx| idx.checked_sub(1)),
                        NodeIndexType::Span { .. } => segmentation_names.len().checked_sub(1),
                    })
                    .min();

                if let Some(start_seg_idx) = start_seg_idx {
                    for seg_idx in (0..=start_seg_idx).rev() {
                        let found_nodes = self.segmentation_nodes_by_token_overlap(
                            &token_range,
                            &segmentation_names[seg_idx],
                        );
                        if !found_nodes.is_empty() {
                            selected_node_names = found_nodes;
                            break;
                        }
                    }
                }
            }

            if selected_node_names.is_empty() {
                // Select the covered token
                for token_position in token_range {
                    selected_node_names.insert(self.token[token_position].node_name.clone());
                }
            }
        }

        selected_node_names
            .into_iter()
            .filter_map(|n| self.node_index_by_name.get(&n))
            .cloned()
            .collect()
    }

    fn create_add_base_token_state_updates(
        &self,
        parent_name: &str,
        reference_node: &str,
        before: bool,
        state_updates: &mut EditorStateUpdates,
    ) -> anyhow::Result<()> {
        let graph = self.graph.read();

        let name_prefix = if let Some((prefix, _suffix)) = parent_name.split_once("#") {
            prefix
        } else {
            parent_name
        };
        let name_prefix = name_prefix.to_string();
        let number_of_token = graph
            .get_node_annos()
            .exact_anno_search(Some(ANNIS_NS), "tok", ValueSearch::Any)
            .count();
        let new_node_name = format!("{name_prefix}#tok{}", number_of_token + 1);
        let reference_node = reference_node.to_string();
        state_updates.set_before(move |editor: &mut SpanEditor| {
            if let Some(ref_tok_position) = editor.get_token_index_by_name(&reference_node) {
                let new_tok_pos = if before {
                    ref_tok_position
                } else {
                    ref_tok_position + 1
                };
                let new_token = Token {
                    node_name: new_node_name.clone(),
                    start: new_tok_pos,
                    end: new_tok_pos,
                    labels: [(TOKEN_KEY.as_ref().clone(), "".to_string())]
                        .into_iter()
                        .collect(),
                };
                editor.token.insert(new_tok_pos, new_token.clone());

                editor.node_index_by_name =
                    calculate_node_index(&editor.token, &editor.segmentations, &editor.span_rows);
                editor.layout_cache = LayoutCache::new();

                editor.selected_nodes.clear();
                editor.scroll_to_token = Some(new_tok_pos);
                editor.selected_nodes.insert(new_node_name.clone());
                editor.active_interaction = Some(ActiveInteraction::EmptyNodeAdded {
                    node_name: new_node_name,
                });
            }
        });
        Ok(())
    }

    fn create_add_segmentation_state_update(
        &self,
        parent_name: &str,
        segmentation: &str,
        selected_token: HashSet<String>,
        state_updates: &mut EditorStateUpdates,
    ) -> anyhow::Result<()> {
        let graph = self.graph.read();
        let new_node_name = format!(
            "{}#{}",
            &parent_name,
            graph
                .get_node_annos()
                .get_largest_item()?
                .map(|id| id + 1)
                .unwrap_or_default()
        );
        let tok_helper = TokenHelper::new(&graph)?;
        let mut sorted_covered_token = Vec::new();
        for node_name in selected_token {
            if let Some(n) = self.get_node_id_from_name(&node_name, &graph) {
                sorted_covered_token.push((n, node_name.to_string()));
            }
        }
        if let Some(gs) = tok_helper.get_ordering_gs(None) {
            sorted_covered_token.sort_by(|a, b| {
                if a == b {
                    Ordering::Equal
                } else if let Ok(connected) =
                    gs.is_connected(a.0, b.0, 1, std::ops::Bound::Unbounded)
                {
                    if connected {
                        Ordering::Less
                    } else {
                        Ordering::Greater
                    }
                } else {
                    Ordering::Less
                }
            });
        }

        let first_covered = sorted_covered_token.first().cloned();
        let last_covered = sorted_covered_token.last().cloned();

        let segmentation = segmentation.to_string();

        state_updates.set_before(move |editor: &mut SpanEditor| {
            editor
                .segmentations
                .entry(segmentation.clone())
                .or_default();
            let base_token_length = editor
                .segmentations
                .get("")
                .map(|token| token.len())
                .unwrap_or(0);
            if let (Some(seg_token), Some(first_covered), Some(last_covered)) = (
                editor.segmentations.get_mut(&segmentation),
                first_covered,
                last_covered,
            ) {
                // Insert the newly generated segmentation token at the approbiate position
                let first_covered_idx = match editor.node_index_by_name.get(&first_covered.1) {
                    Some(NodeIndexType::Token { token_index }) => Some(*token_index),
                    _ => None,
                }
                .unwrap_or(0);

                let last_covered_idx = match editor.node_index_by_name.get(&last_covered.1) {
                    Some(NodeIndexType::Token { token_index }) => Some(*token_index),
                    _ => None,
                }
                .unwrap_or(base_token_length);
                let mut new_token_labels = BTreeMap::new();
                new_token_labels.insert(TOKEN_KEY.as_ref().clone(), String::default());
                new_token_labels.insert(
                    AnnoKey {
                        name: segmentation.as_str().into(),
                        ns: ANNIS_NS.into(),
                    },
                    String::default(),
                );
                let new_token = Token {
                    node_name: new_node_name.clone(),
                    start: first_covered_idx,
                    end: last_covered_idx,
                    labels: new_token_labels,
                };
                let token_index =
                    match seg_token.binary_search_by(|probe| probe.end.cmp(&first_covered_idx)) {
                        Ok(idx) => idx + 1,
                        Err(idx) => idx,
                    };
                seg_token.insert(token_index, new_token);
                editor.node_index_by_name.insert(
                    new_node_name.clone(),
                    NodeIndexType::Segmentation {
                        seg_name: segmentation,
                        token_index,
                    },
                );

                editor.selected_nodes.clear();
                editor.scroll_to_token = Some(first_covered_idx);
                editor.selected_nodes.insert(new_node_name.clone());
                editor.active_interaction = Some(ActiveInteraction::EmptyNodeAdded {
                    node_name: new_node_name,
                });
            }
        });

        Ok(())
    }
}

fn handle_widget_selection_event(
    selected_nodes: &mut HashSet<String>,
    current_selection: SelectionType,
    widget_selected: bool,
    node_name: &str,
    ui: &Ui,
) {
    if widget_selected {
        let was_already_selected = current_selection != SelectionType::None;

        if ui.ctx().input(|i| i.modifiers.command_only()) {
            if was_already_selected {
                // Unselect
                selected_nodes.remove(node_name);
            } else {
                // Allow selection of multiple items
                selected_nodes.insert(node_name.to_string());
            }
        } else if !was_already_selected {
            // Select only one node
            selected_nodes.clear();
            selected_nodes.insert(node_name.to_string());
        }
    }
}

impl Editor for SpanEditor {
    fn show(&mut self, ui: &mut Ui) {
        let ui_style = ui.style().clone();

        if self.token.is_empty() {
            // Add a button to create the first token
            if ui
                .button(RichText::new(format!("{} Add first token", icons::EGG_CRACK)).heading())
                .clicked()
            {
                self.pending_graph_actions
                    .push(GraphAction::AddInitialToken {
                        parent_name: self.parent_name.clone(),
                    });
            }
            return;
        }

        self.show_new_span_dialog(ui);
        self.show_search_dialog(ui);

        let mut current_span_offset: f32 = 0.0;

        // Remember the location of each token, so we can paint the spans with
        // the same range later
        let mut token_offset_to_rect = vec![None; self.token.len()];

        let available_node_annos: BTreeSet<AnnoKey> = self
            .span_rows
            .iter()
            .flat_map(|row| row.anno_keys.iter())
            .cloned()
            .collect();

        ui.with_layout(
            egui::Layout::left_to_right(egui::Align::Min).with_cross_justify(true),
            |ui| {
                let mut filter_widget =
                    FilterWidget::new(available_node_annos, "document-filter-widget");
                if let Some(ActiveInteraction::AddingNewSpan {
                    namespace,
                    name,
                    is_segmentation,
                }) = &self.active_interaction
                    && !is_segmentation
                {
                    let anno_key = AnnoKey {
                        name: name.into(),
                        ns: namespace.into(),
                    };
                    filter_widget.select_anno_key(anno_key, ui);
                }
                let FilterWidgetOutput {
                    selected_node_annos,
                } = filter_widget.show(ui).inner;

                ScrollArea::horizontal()
                    .animated(false)
                    .show_viewport(ui, |ui, _viewport| {
                        // If we already calculated the token positions once, only render
                        // the token and their covering spans that are currently displayed
                        let last_token_index = self.token.len() - 1;

                        ui.horizontal_top(|ui| {
                            for token_position in 0..=last_token_index {
                                let token_start = self.token[token_position].start;

                                let token_rect = self.show_single_token(ui, token_position);
                                current_span_offset = current_span_offset.max(token_rect.bottom());
                                if token_start < token_offset_to_rect.len() {
                                    token_offset_to_rect[token_start] = Some(token_rect);
                                }
                            }
                        });
                        current_span_offset += ui_style.spacing.item_spacing.y;

                        ui.vertical(|ui| {
                            // If the command key is down, find the selected token and give
                            // them as argument, so we can add segmentations and spans for
                            // them.
                            let highlighted_token_positions: BTreeSet<usize> = if ui
                                .ctx()
                                .input(|i| i.modifiers.command_only())
                                && self.selected_nodes.iter().all(|selected_node| {
                                    self.get_token_index_by_name(selected_node).is_some()
                                }) {
                                self.selected_base_token_indices()
                            } else {
                                BTreeSet::new()
                            };
                            current_span_offset = self.show_segmentation_layers(
                                ui,
                                &token_offset_to_rect,
                                &highlighted_token_positions,
                                current_span_offset,
                            );
                            if self.render_spans {
                                self.show_spans(
                                    ui,
                                    &selected_node_annos,
                                    &token_offset_to_rect,
                                    &highlighted_token_positions,
                                    current_span_offset,
                                );
                            }
                            // Add some space to the horizontal scroll bar is not hiding the spans
                            ui.add_space(20.0);
                        });
                    });
            },
        );

        self.layout_cache
            .rendering_step_finished(ui_style.spacing.item_spacing);
    }

    fn has_pending_updates(&self) -> bool {
        !self.pending_graph_actions.is_empty()
    }

    fn take_pending_updates(&mut self) -> Vec<GraphAction> {
        std::mem::take(&mut self.pending_graph_actions)
    }

    fn editor_state_updates(&self, action: &GraphAction) -> anyhow::Result<EditorStateUpdates> {
        let mut result = EditorStateUpdates::default();
        match action {
            GraphAction::AddInitialToken { parent_name } => {
                let name_prefix = if let Some((prefix, _suffix)) = parent_name.split_once("#") {
                    prefix
                } else {
                    parent_name
                };
                let new_node_name = format!("{name_prefix}#tok{}", 1);

                result.set_after(move |editor: &mut SpanEditor| {
                    if let Err(err) = editor.recreate_from_graph() {
                        editor.notifier.report_error(err);
                    }
                    editor.active_interaction = Some(ActiveInteraction::EmptyNodeAdded {
                        node_name: new_node_name.clone(),
                    });
                    editor.selected_nodes.insert(new_node_name);
                });
            }
            GraphAction::AddBaseTokenBefore {
                reference_node,
                parent_name,
            } => {
                self.create_add_base_token_state_updates(
                    parent_name,
                    reference_node,
                    true,
                    &mut result,
                )?;
            }
            GraphAction::AddBaseTokenAfter {
                reference_node,
                parent_name,
            } => {
                self.create_add_base_token_state_updates(
                    parent_name,
                    reference_node,
                    false,
                    &mut result,
                )?;
            }
            GraphAction::ExpandNodeLeft { reference_node } => {
                if let Some(index_entry) = self.node_index_by_name.get(reference_node) {
                    match index_entry {
                        NodeIndexType::Segmentation {
                            seg_name,
                            token_index,
                        } => {
                            let seg_name = seg_name.clone();
                            let token_index = *token_index;
                            result.set_after(move |editor: &mut SpanEditor| {
                                if let Some(seg) = editor.segmentations.get_mut(&seg_name) {
                                    seg[token_index].start -= 1;
                                }
                            });
                        }
                        NodeIndexType::Span {
                            index_of_row,
                            position_in_row,
                        } => {
                            let graph = self.graph.read();
                            let tok_helper = TokenHelper::new(&graph)?;
                            if let Some(reference_node_id) = graph
                                .get_node_annos()
                                .get_node_id_from_name(reference_node)?
                                && !tok_helper.is_token(reference_node_id)?
                                && let Some(new_covered_token) =
                                    tok_helper.get_token_before(reference_node_id, None)?
                            {
                                let index_of_row = *index_of_row;
                                let position_in_row = *position_in_row;
                                result.set_before(move |editor: &mut SpanEditor| {
                                    editor.span_rows[index_of_row].spans[position_in_row]
                                        .sorted_covered_token_ids
                                        .insert_before(0, new_covered_token);

                                    editor.span_rows[index_of_row]
                                        .occupied_token
                                        .insert(new_covered_token);
                                });
                            }
                        }
                        _ => {}
                    }
                }
            }
            GraphAction::ExpandNodeRight { reference_node } => {
                if let Some(index_entry) = self.node_index_by_name.get(reference_node) {
                    match index_entry {
                        NodeIndexType::Segmentation {
                            seg_name,
                            token_index,
                        } => {
                            let seg_name = seg_name.clone();
                            let token_index = *token_index;
                            result.set_after(move |editor: &mut SpanEditor| {
                                if let Some(seg) = editor.segmentations.get_mut(&seg_name) {
                                    seg[token_index].end += 1;
                                }
                            });
                        }
                        NodeIndexType::Span {
                            index_of_row,
                            position_in_row,
                        } => {
                            let graph = self.graph.read();
                            let tok_helper = TokenHelper::new(&graph)?;
                            if let Some(reference_node_id) = graph
                                .get_node_annos()
                                .get_node_id_from_name(reference_node)?
                                && !tok_helper.is_token(reference_node_id)?
                                && let Some(new_covered_token) =
                                    tok_helper.get_token_after(reference_node_id, None)?
                            {
                                let index_of_row = *index_of_row;
                                let position_in_row = *position_in_row;
                                result.set_before(move |editor: &mut SpanEditor| {
                                    editor.span_rows[index_of_row].spans[position_in_row]
                                        .sorted_covered_token_ids
                                        .insert(new_covered_token);
                                    editor.span_rows[index_of_row]
                                        .occupied_token
                                        .insert(new_covered_token);
                                });
                            }
                        }
                        _ => {}
                    }
                }
            }
            GraphAction::ShrinkNodeLeft { reference_node } => {
                if let Some(index_entry) = self.node_index_by_name.get(reference_node) {
                    match index_entry {
                        NodeIndexType::Segmentation {
                            seg_name,
                            token_index,
                        } => {
                            let seg_name = seg_name.clone();
                            let token_index = *token_index;
                            result.set_after(move |editor: &mut SpanEditor| {
                                if let Some(seg) = editor.segmentations.get_mut(&seg_name) {
                                    seg[token_index].start += 1;
                                }
                            });
                        }
                        NodeIndexType::Span {
                            index_of_row,
                            position_in_row,
                        } => {
                            let graph = self.graph.read();
                            let tok_helper = TokenHelper::new(&graph)?;
                            if let Some(reference_node_id) = graph
                                .get_node_annos()
                                .get_node_id_from_name(reference_node)?
                                && !tok_helper.is_token(reference_node_id)?
                            {
                                let covered_token = tok_helper.covered_token(reference_node_id)?;
                                if let Some(first_covered_token) = covered_token.first().cloned() {
                                    let index_of_row = *index_of_row;
                                    let position_in_row = *position_in_row;
                                    result.set_before(move |editor: &mut SpanEditor| {
                                        editor.span_rows[index_of_row].spans[position_in_row]
                                            .sorted_covered_token_ids
                                            .remove(&first_covered_token);
                                        editor.span_rows[index_of_row]
                                            .occupied_token
                                            .remove(&first_covered_token);
                                    });
                                }
                            }
                        }
                        _ => {}
                    }
                }
            }
            GraphAction::ShrinkNodeRight { reference_node } => {
                if let Some(index_entry) = self.node_index_by_name.get(reference_node) {
                    match index_entry {
                        NodeIndexType::Segmentation {
                            seg_name,
                            token_index,
                        } => {
                            let seg_name = seg_name.clone();
                            let token_index = *token_index;
                            result.set_after(move |editor: &mut SpanEditor| {
                                if let Some(seg) = editor.segmentations.get_mut(&seg_name) {
                                    seg[token_index].end -= 1;
                                }
                            });
                        }
                        NodeIndexType::Span {
                            index_of_row,
                            position_in_row,
                        } => {
                            let graph = self.graph.read();
                            let tok_helper = TokenHelper::new(&graph)?;
                            if let Some(reference_node_id) = graph
                                .get_node_annos()
                                .get_node_id_from_name(reference_node)?
                                && !tok_helper.is_token(reference_node_id)?
                            {
                                let covered_token = tok_helper.covered_token(reference_node_id)?;
                                if let Some(last_covered_token) = covered_token.last().cloned() {
                                    let index_of_row = *index_of_row;
                                    let position_in_row = *position_in_row;
                                    result.set_before(move |editor: &mut SpanEditor| {
                                        editor.span_rows[index_of_row].spans[position_in_row]
                                            .sorted_covered_token_ids
                                            .remove(&last_covered_token);
                                        editor.span_rows[index_of_row]
                                            .occupied_token
                                            .remove(&last_covered_token);
                                    });
                                }
                            }
                        }
                        _ => {}
                    }
                }
            }
            GraphAction::AnnoValueChanged { node_name, anno } => {
                let node_name = node_name.clone();
                let anno = anno.clone();
                result.set_before(move |editor: &mut SpanEditor| {
                    if let Some(labels) = editor.get_labels_for_node_mut(&node_name) {
                        labels.insert(anno.key, anno.val.to_string());
                    }
                });
            }
            GraphAction::AnnoRemoved {
                node_name,
                anno_key,
            } => {
                let node_name = node_name.clone();
                let anno_key = anno_key.clone();
                result.set_before(move |editor: &mut SpanEditor| {
                    if let Some(labels) = editor.get_labels_for_node_mut(&node_name) {
                        labels.remove(&anno_key);
                    }
                });
            }
            GraphAction::AddSpan { parent_name, .. } => {
                let parent_name = parent_name.clone();
                result.set_after(move |editor: &mut SpanEditor| {
                    let result = editor.recreate_from_graph();
                    editor.notifier.ok_or_report(result);

                    let graph = editor.graph.read();

                    if let Ok(Some(new_node_id)) = graph.get_node_annos().get_largest_item() {
                        let new_node_name = format!("{parent_name}#n{new_node_id}");
                        editor.selected_nodes.clear();
                        editor.selected_nodes.insert(new_node_name.clone());
                        editor.active_interaction = Some(ActiveInteraction::EmptyNodeAdded {
                            node_name: new_node_name,
                        });
                    }
                });
            }
            GraphAction::AddSegmentationSpan {
                segmentation,
                selected_token,
                parent_name,
            } => {
                self.create_add_segmentation_state_update(
                    parent_name,
                    segmentation,
                    selected_token.clone(),
                    &mut result,
                )?;
            }
            GraphAction::DeleteNodes { node_names } => {
                let mut node_names_set = HashSet::new();
                let mut node_ids_set = HashSet::new();
                for node_name in node_names {
                    let graph = self.graph.read();
                    node_names_set.insert(node_name.clone());
                    if let Some(node_id) = self.get_node_id_from_name(node_name, &graph) {
                        node_ids_set.insert(node_id);
                    }
                }
                result.set_before(move |editor: &mut SpanEditor| {
                    // Remove node names an IDs from all internal representations
                    editor
                        .token
                        .retain(|t| !node_names_set.contains(t.node_name.as_str()));

                    for (_seg, token) in editor.segmentations.iter_mut() {
                        token.retain(|t| !node_names_set.contains(t.node_name.as_str()));
                    }

                    for row in editor.span_rows.iter_mut() {
                        for node_id in &node_ids_set {
                            row.occupied_token.remove(node_id);
                        }
                        row.spans
                            .retain(|span| !node_names_set.contains(span.node_name.as_str()));
                        for span in row.spans.iter_mut() {
                            for node_id in &node_ids_set {
                                span.sorted_covered_token_ids.remove(node_id);
                            }
                        }
                    }
                    editor
                        .selected_nodes
                        .retain(|n| !node_names_set.contains(n));
                    // TODO: try to merge rows

                    // Re-calculate the index from node name to index, because also the non-deleted nodes will have a different index
                    editor.node_index_by_name = calculate_node_index(
                        &editor.token,
                        &editor.segmentations,
                        &editor.span_rows,
                    );

                    // Clear size related caches
                    editor.layout_cache = LayoutCache::new();
                });
            }
            _ => {}
        }
        Ok(result)
    }

    fn get_edited_node(&self) -> &str {
        &self.parent_name
    }

    fn get_selected_nodes(&self) -> BTreeSet<String> {
        self.selected_nodes.iter().cloned().collect()
    }

    fn consume_shortcuts(&mut self, ctx: &egui::Context) {
        Search::consume_and_perform_shortcut(ctx, self);
        GoToFirstToken::consume_and_perform_shortcut(ctx, self);
        GoToLastToken::consume_and_perform_shortcut(ctx, self);
        AddSpanAction::consume_and_perform_shortcut(ctx, self);
        AddBaseTokenBefore::consume_and_perform_shortcut(ctx, self);
        AddBaseTokenAfter::consume_and_perform_shortcut(ctx, self);
        EditWhitespaceBefore::consume_and_perform_shortcut(ctx, self);
        EditWhitespaceAfter::consume_and_perform_shortcut(ctx, self);
        ShrinkNodeLeft::consume_and_perform_shortcut(ctx, self);
        ShrinkNodeRight::consume_and_perform_shortcut(ctx, self);
        ExpandNodeLeft::consume_and_perform_shortcut(ctx, self);
        ExpandNodeRight::consume_and_perform_shortcut(ctx, self);
        DeleteSelectedNode::consume_and_perform_shortcut(ctx, self);

        // Handle some navigation with the arrows
        if !self.selected_nodes.is_empty() {
            let any_text_editor_has_focus = ctx
                .memory(|m| m.focused())
                .map(|id| TextEditState::load(ctx, id).is_some())
                .unwrap_or_default();
            if !any_text_editor_has_focus {
                let direction = if ctx.input(|i| i.key_pressed(Key::ArrowRight)) {
                    Some(Direction::LeftToRight)
                } else if ctx.input(|i| i.key_pressed(Key::ArrowLeft)) {
                    Some(Direction::RightToLeft)
                } else if ctx.input(|i| i.key_pressed(Key::ArrowDown)) {
                    Some(Direction::TopDown)
                } else if ctx.input(|i| i.key_pressed(Key::ArrowUp)) {
                    Some(Direction::BottomUp)
                } else {
                    None
                };
                if let Some(direction) = direction {
                    let modifiers = ctx.input(|i| i.modifiers);
                    let result = self.handle_selection_key_press_event(direction, modifiers);
                    self.notifier.ok_or_report(result);
                }
            }

            // Add shortcuts for adding new spans.
            // The first layers are for segmentation spans, then add shortcuts for the remaining span rows.
            // At most 9 shortcuts can be created.
            let number_of_layers = (self.segmentations.len() + self.span_rows.len()).min(9);

            for layer_idx in 0..number_of_layers {
                if let Some(key) = Key::from_name(&(layer_idx + 1).to_string())
                    && ctx.input_mut(|i| i.consume_key(Modifiers::COMMAND, key))
                {
                    if layer_idx < self.segmentations.len() {
                        self.add_segmentation_for_selection(layer_idx);
                    } else {
                        self.add_span_for_row(layer_idx.saturating_sub(self.segmentations.len()));
                    }
                }
            }
        } // end at least one node selected
    }

    fn add_edit_menu_entries(&mut self, ui: &mut egui::Ui) {
        AddBaseTokenAfter::create_menu_button(ui, self);
        AddBaseTokenBefore::create_menu_button(ui, self);
        EditWhitespaceBefore::create_menu_button(ui, self);
        EditWhitespaceAfter::create_menu_button(ui, self);

        ui.separator();

        ExpandNodeLeft::create_menu_button(ui, self);
        ExpandNodeRight::create_menu_button(ui, self);
        ShrinkNodeLeft::create_menu_button(ui, self);
        ShrinkNodeRight::create_menu_button(ui, self);

        ui.separator();

        AddSpanAction::create_menu_button(ui, self);
        DeleteSelectedNode::create_menu_button(ui, self);

        ui.separator();

        Search::create_menu_button(ui, self);
    }

    fn any_mut(&mut self) -> &mut dyn std::any::Any {
        self
    }
}