annatomic 0.4.0

use std::collections::HashMap;

use egui::{Rect, Vec2};

use crate::app::data::Token;

#[derive(Clone, Debug, PartialEq)]
enum LayoutPhase {
    Initial,
    Layouting,
    Valid,
}

#[derive(Clone)]
pub(crate) struct LayoutCache {
    phase: LayoutPhase,
    /// For each token offset, the different minimal widths as derived from each segmentation layer.
    min_token_widths_by_segmentation: Vec<HashMap<String, f32>>,
    base_token_sizes: Vec<Vec2>,
    total_width: f32,
}

impl LayoutCache {
    pub fn new() -> Self {
        Self {
            phase: LayoutPhase::Initial,
            min_token_widths_by_segmentation: Vec::default(),
            base_token_sizes: Vec::default(),
            total_width: 0.0,
        }
    }

    pub fn is_valid(&self) -> bool {
        self.phase == LayoutPhase::Valid
    }

    pub fn add_segmentation_width(
        &mut self,
        seg_name: &str,
        seg_token: &Token,
        width: f32,
        base_tokens: &[Token],
    ) {
        if self.switch_to_layouting_phase(base_tokens) {
            let widget_width_per_token = width / ((seg_token.end - seg_token.start + 1) as f32);

            for offset in seg_token.start..=seg_token.end {
                if offset < self.min_token_widths_by_segmentation.len()
                    && !self.min_token_widths_by_segmentation[offset].contains_key(seg_name)
                {
                    self.min_token_widths_by_segmentation[offset]
                        .entry(seg_name.to_string())
                        .or_insert(widget_width_per_token);
                }
            }
        }
    }

    pub fn add_token_rect(
        &mut self,
        token_position: usize,
        widget_rect: Rect,
        base_tokens: &[Token],
    ) {
        self.switch_to_layouting_phase(base_tokens);
        // get a minimal token width based on the segmentation spans covering this token
        let min_token_width = self.min_base_token_width(token_position);
        // Only fill cache with node position if the resizing based on the segmentation spans is finished for this token
        if min_token_width.is_some()
            || self.min_token_widths_by_segmentation[token_position].is_empty()
        {
            self.base_token_sizes[token_position] = widget_rect.size();
        }
    }

    pub fn rendering_step_finished(&mut self, item_spacing: Vec2) {
        if self.phase == LayoutPhase::Layouting {
            let mut all_base_token_added = true;
            let mut new_total_width = 0.0;
            for size in &self.base_token_sizes {
                if size.any_nan() {
                    all_base_token_added = false;
                    break;
                }
                new_total_width += size.x + item_spacing.x;
            }
            if all_base_token_added {
                self.total_width = new_total_width;
                self.phase = LayoutPhase::Valid;
            }
        }
    }

    pub fn min_base_token_width(&self, token_position: usize) -> Option<f32> {
        if self.phase == LayoutPhase::Layouting || self.phase == LayoutPhase::Valid {
            let min_widths_for_token = &self.min_token_widths_by_segmentation[token_position];
            if min_widths_for_token.is_empty() {
                None
            } else {
                Some(
                    self.min_token_widths_by_segmentation[token_position]
                        .values()
                        .fold(0.0, |old_value: f32, new_value: &f32| {
                            if old_value.total_cmp(new_value).is_lt() {
                                *new_value
                            } else {
                                old_value
                            }
                        }),
                )
            }
        } else {
            None
        }
    }

    pub fn cached_base_token_size(&self, token_position: usize) -> Option<Vec2> {
        if self.phase == LayoutPhase::Valid {
            self.base_token_sizes.get(token_position).copied()
        } else {
            None
        }
    }

    /// Switches from `Initial` to `Layouting` and resets the cache if
    /// switching. If the state is already `Valid`, this returns `false`.
    fn switch_to_layouting_phase(&mut self, base_tokens: &[Token]) -> bool {
        if self.phase == LayoutPhase::Valid {
            return false;
        } else if self.phase != LayoutPhase::Layouting {
            // Reset the internal caches
            self.min_token_widths_by_segmentation = vec![HashMap::new(); base_tokens.len()];
            self.base_token_sizes = vec![Vec2::NAN; base_tokens.len()];
            self.total_width = 0.0;
            self.phase = LayoutPhase::Layouting
        }
        true
    }
}