annatomic 0.2.0

use std::f32;
use std::sync::Arc;

use egui::{Button, Key, ModalResponse, Modifiers, TextEdit};
use egui::{
    Context, Frame, Id, InnerResponse, Label, RichText, Sense, Widget, WidgetInfo, WidgetWithState,
};
use graphannis::graph::{AnnoKey, Annotation};
use graphannis_core::graph::ANNIS_NS;
use lazy_static::lazy_static;
use serde::{Deserialize, Serialize};

use crate::app::data::Token;
use crate::app::util::{make_whitespace_visible, token_helper::TOKEN_KEY};
use crate::app::widgets::create_single_anno_key_labels;

lazy_static! {
    static ref WHITESPACE_BEFORE: Arc<AnnoKey> = Arc::from(AnnoKey {
        ns: ANNIS_NS.into(),
        name: "tok-whitespace-before".into(),
    });
    static ref WHITESPACE_AFTER: Arc<AnnoKey> = Arc::from(AnnoKey {
        ns: ANNIS_NS.into(),
        name: "tok-whitespace-after".into(),
    });
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SelectionType {
    None,
    Selected,
    SelectionStart,
    SelectionEnd,
}

#[derive(Clone, Debug, Serialize, Deserialize, Default, PartialEq)]
pub enum EditedValue {
    #[default]
    None,
    Token(String),
    Label {
        key: AnnoKey,
        value: String,
    },
    WhitespaceBefore(String),
    WhitespaceAfter(String),
}

#[derive(Debug)]
pub struct TokenWidget<'t> {
    token: &'t Token,
    selected: SelectionType,
    min_width: Option<f32>,
    width: Option<f32>,
    value: String,
    whitespace_before: String,
    whitespace_after: String,
    seg_name: Option<String>,
}

#[derive(Debug, Default)]
pub struct TokenWidgetOutput {
    pub selected: bool,
    pub add_token_after_action: bool,
    pub changed_annotations: Vec<Annotation>,
}

#[derive(Clone, Default, Debug, Serialize, Deserialize)]
pub struct TokenWidgetState {
    pub edited_value: EditedValue,
    pub token_value_width: Option<f32>,
    pub adding_new_anno: Option<(String, String)>,
}

impl Widget for TokenWidget<'_> {
    fn ui(self, ui: &mut egui::Ui) -> egui::Response {
        self.show(ui).response
    }
}

impl WidgetWithState for TokenWidget<'_> {
    type State = TokenWidgetState;
}

impl<'t> TokenWidget<'t> {
    pub fn with_exact_width(
        token: &'t Token,
        selected: SelectionType,
        width: Option<f32>,
        seg_name: Option<String>,
    ) -> Self {
        TokenWidget {
            token,
            selected,
            min_width: None,
            width,
            value: token.labels.get(&TOKEN_KEY).cloned().unwrap_or_default(),
            whitespace_before: token
                .labels
                .get(&WHITESPACE_BEFORE)
                .cloned()
                .unwrap_or_default(),
            whitespace_after: token
                .labels
                .get(&WHITESPACE_AFTER)
                .cloned()
                .unwrap_or_default(),
            seg_name,
        }
    }

    pub fn with_min_width(
        token: &'t Token,
        selected: SelectionType,
        min_width: Option<f32>,
        seg_name: Option<String>,
    ) -> Self {
        TokenWidget {
            token,
            selected,
            min_width,
            width: None,
            value: token.labels.get(&TOKEN_KEY).cloned().unwrap_or_default(),
            whitespace_before: token
                .labels
                .get(&WHITESPACE_BEFORE)
                .cloned()
                .unwrap_or_default(),
            whitespace_after: token
                .labels
                .get(&WHITESPACE_AFTER)
                .cloned()
                .unwrap_or_default(),
            seg_name,
        }
    }

    pub fn load_state(&self, ctx: &Context) -> TokenWidgetState {
        let widget_id = Id::from(self.token.node_name.clone());
        ctx.data_mut(|d| d.get_persisted(widget_id))
            .unwrap_or_default()
    }

    pub fn store_state(&self, ctx: &Context, state: TokenWidgetState) {
        let widget_id = Id::from(self.token.node_name.clone());
        ctx.data_mut(|d| d.insert_persisted(widget_id, state));
    }

    /// Sets the state so that the token value is edited.
    pub fn request_edit_token_value(&mut self, ui: &egui::Ui) {
        let mut state = self.load_state(ui.ctx());
        if !matches!(state.edited_value, EditedValue::Token(_)) {
            state.edited_value = EditedValue::Token(self.value.clone());
        };
        self.store_state(ui.ctx(), state);
    }

    /// Sets the state so that the whitespace before value is edited.
    pub fn request_edit_whitespace_before_value(&mut self, ui: &egui::Ui) {
        let mut state = self.load_state(ui.ctx());
        if !matches!(state.edited_value, EditedValue::WhitespaceBefore(_)) {
            state.edited_value = EditedValue::WhitespaceBefore(self.whitespace_before.clone());
        };
        self.store_state(ui.ctx(), state);
    }

    /// Sets the state so that the whitespace after value is edited.
    pub fn request_edit_whitespace_after_value(&mut self, ui: &egui::Ui) {
        let mut state = self.load_state(ui.ctx());
        if !matches!(state.edited_value, EditedValue::WhitespaceAfter(_)) {
            state.edited_value = EditedValue::WhitespaceAfter(self.whitespace_after.clone());
        };
        self.store_state(ui.ctx(), state);
    }

    fn show_token_line(
        &self,
        state: &mut TokenWidgetState,
        output: &mut TokenWidgetOutput,
        ui: &mut egui::Ui,
    ) -> (Vec<Annotation>, bool) {
        ui.horizontal(|ui| {
            let mut result = Vec::new();
            let mut editing_finished_with_space = false;
            if let EditedValue::WhitespaceBefore(edited_whitespace_label) = &mut state.edited_value
            {
                let edit = TextEdit::singleline(edited_whitespace_label)
                    .clip_text(true)
                    .desired_width(50.0)
                    .cursor_at_end(true);
                let output = edit.show(ui);
                if output.response.lost_focus()
                    && ui.input_mut(|i| i.consume_key(Modifiers::NONE, Key::Enter))
                {
                    result.push(Annotation {
                        key: WHITESPACE_BEFORE.as_ref().clone(),
                        val: edited_whitespace_label.as_str().into(),
                    });
                    state.edited_value = EditedValue::None;
                } else if !output.response.has_focus() {
                    output.response.request_focus();
                }
            } else if !self.whitespace_before.is_empty() {
                // Put the whitespace and the actual value in one line
                ui.label(RichText::new(make_whitespace_visible(&self.whitespace_before)).weak());
            }

            if let EditedValue::Token(edited_token_value) = &mut state.edited_value {
                let mut edit = TextEdit::singleline(edited_token_value).cursor_at_end(true);
                if self.width.is_some() {
                    // Fixed width requested, do not grow text editor
                    edit = edit.clip_text(true);
                } else {
                    // Flexible width, grow text editor line as needed
                    edit = edit.clip_text(false).desired_width(0.0);
                }
                let output = edit.show(ui);

                // If this cursor is at the end, interpret a simple space press (without modifiers) as "finish editing and go to next token"
                if let Some(cursor) = output.cursor_range.and_then(|c| c.single())
                    && cursor.index == edited_token_value.len()
                    && ui.input_mut(|i| {
                        i.modifiers.is_none() && i.consume_key(Modifiers::NONE, Key::Space)
                    })
                {
                    editing_finished_with_space = true;
                    edited_token_value.pop();
                }

                if (output.response.lost_focus()
                    && ui.input_mut(|i| i.consume_key(Modifiers::NONE, Key::Enter)))
                    || editing_finished_with_space
                {
                    result.push(Annotation {
                        key: TOKEN_KEY.as_ref().clone(),
                        val: edited_token_value.as_str().into(),
                    });
                    if editing_finished_with_space {
                        result.push(Annotation {
                            key: WHITESPACE_AFTER.as_ref().clone(),
                            val: " ".into(),
                        });
                    }
                    state.edited_value = EditedValue::None;
                } else if !output.response.has_focus() {
                    output.response.request_focus();
                }
            } else {
                let label = ui.label(RichText::new(make_whitespace_visible(&self.value)).strong());
                if label.hovered() {
                    ui.ctx().set_cursor_icon(egui::CursorIcon::PointingHand);
                }
                if label.clicked() {
                    if self.selected == SelectionType::None {
                        output.selected = true;
                    } else {
                        state.edited_value = EditedValue::Token(self.value.clone());
                    }
                }
                state.token_value_width = Some(label.rect.width());
            }

            if let EditedValue::WhitespaceAfter(edited_whitespace_label) = &mut state.edited_value {
                let edit = TextEdit::singleline(edited_whitespace_label)
                    .clip_text(true)
                    .desired_width(50.0)
                    .cursor_at_end(true);
                let output = edit.show(ui);
                if output.response.lost_focus()
                    && ui.input_mut(|i| i.consume_key(Modifiers::NONE, Key::Enter))
                {
                    result.push(Annotation {
                        key: WHITESPACE_AFTER.as_ref().clone(),
                        val: edited_whitespace_label.as_str().into(),
                    });
                    state.edited_value = EditedValue::None;
                } else if !output.response.has_focus() {
                    output.response.request_focus();
                }
            } else if !self.whitespace_after.is_empty() {
                ui.label(RichText::new(make_whitespace_visible(&self.whitespace_after)).weak());
            }
            (result, editing_finished_with_space)
        })
        .inner
    }

    fn show_labels(
        &self,
        state: &mut TokenWidgetState,
        output: &mut TokenWidgetOutput,
        ui: &mut egui::Ui,
    ) -> Option<Annotation> {
        let mut result = None;
        for (current_key, current_value) in self.token.labels.iter() {
            // Do not render internal graphANNIS labels and also not the
            // segmentation value if it is the same as the annis:tok value
            let is_segmentation_value = if let Some(seg_name) = &self.seg_name
                && current_key.name.as_str() == seg_name
                && current_value == &self.value
            {
                true
            } else {
                false
            };
            if current_key.ns != ANNIS_NS && !is_segmentation_value {
                let edited_value =
                    if let EditedValue::Label { key, value } = &mut state.edited_value {
                        if key == current_key {
                            Some((key, value))
                        } else {
                            None
                        }
                    } else {
                        None
                    };
                if let Some((key, value)) = edited_value {
                    let value_edit = TextEdit::singleline(value)
                        .clip_text(true)
                        .cursor_at_end(true)
                        .ui(ui)
                        .on_hover_ui(|ui| {
                            create_single_anno_key_labels(current_key, ui);
                        });
                    if value_edit.lost_focus()
                        && ui.input_mut(|i| i.consume_key(Modifiers::NONE, Key::Enter))
                    {
                        result = Some(Annotation {
                            key: key.clone(),
                            val: value.as_str().into(),
                        });
                        state.edited_value = EditedValue::None;
                    }
                } else {
                    let value_label = Label::new(current_value)
                        .wrap_mode(egui::TextWrapMode::Extend)
                        .ui(ui)
                        .on_hover_ui(|ui| {
                            create_single_anno_key_labels(current_key, ui);
                        });
                    if value_label.hovered() {
                        ui.ctx().set_cursor_icon(egui::CursorIcon::PointingHand);
                    }
                    if value_label.clicked() {
                        if self.selected == SelectionType::None {
                            output.selected = true;
                        } else {
                            state.edited_value = EditedValue::Label {
                                key: current_key.clone(),
                                value: current_value.clone(),
                            };
                        }
                    }
                }
            }
        }
        // Add a row for adding new annotations
        ui.horizontal(|ui| {
            let add_button = ui
                .add_visible(
                    self.selected != SelectionType::None {},
                    Button::new(egui_phosphor::regular::PLUS_CIRCLE),
                )
                .on_hover_text("Add new label");
            if add_button.clicked() {
                state.adding_new_anno = Some((String::new(), String::new()));
            }
        });
        result
    }

    fn show_new_anno_dialog(
        &self,
        ns: &mut String,
        name: &mut String,
        ui: &egui::Ui,
    ) -> ModalResponse<Option<bool>> {
        egui::Modal::new("adding_new_anno".into()).show(ui.ctx(), move |ui| {
            let text_field_response = ui
                .horizontal(|ui| {
                    TextEdit::singleline(ns)
                        .hint_text("Namespace")
                        .horizontal_align(egui::Align::Max)
                        .desired_width(80.0)
                        .ui(ui);
                    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.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)
        })
    }

    pub fn show(self, ui: &mut egui::Ui) -> InnerResponse<TokenWidgetOutput> {
        let mut state = self.load_state(ui.ctx());

        if self.selected == SelectionType::None {
            state.edited_value = EditedValue::None;
            state.adding_new_anno = None;
        }
        if let Some((ns, name)) = state.adding_new_anno.as_mut() {
            let modal_response = self.show_new_anno_dialog(ns, name, ui);

            let mut output = TokenWidgetOutput::default();
            if let Some(accept) = modal_response.inner
                && let Some((namespace, name)) = state.adding_new_anno.take()
                && accept
            {
                output.changed_annotations.push(Annotation {
                    key: AnnoKey {
                        ns: namespace,
                        name,
                    },
                    val: "<new>".into(),
                });
            }
            self.store_state(ui.ctx(), state);
            InnerResponse::new(output, modal_response.response)
        } else {
            let group_response = ui.horizontal_top(|ui| {
                let mut output = TokenWidgetOutput::default();

                let inner_margin = 3.0;
                let mut group_widget = Frame::new()
                    .stroke(ui.style().visuals.widgets.noninteractive.bg_stroke)
                    .inner_margin(inner_margin);
                let width_offset = (inner_margin
                    + ui.style().visuals.widgets.noninteractive.bg_stroke.width)
                    * 2.0;

                if self.selected != SelectionType::None {
                    group_widget.fill = ui.style().visuals.selection.bg_fill;
                }
                group_widget.show(ui, |ui| {
                    if let Some(width) = self.width {
                        ui.set_width(width - width_offset);
                    } else if let Some(min_width) = self.min_width {
                        ui.set_min_width(min_width - width_offset);
                    } else if self.seg_name.is_none() {
                        // For the "add span" shortcut buttons to be displayed
                        // with at least their number, a minimal width is
                        // needed.
                        ui.set_min_width(50.0);
                    }

                    ui.vertical(|ui| {
                        // Show header as first line, then the token information itself and all other labels
                        if self.seg_name.is_none() && self.token.start == self.token.end {
                            ui.horizontal(|ui| {
                                ui.label(
                                    RichText::new(self.token.start.to_string()).weak().small(),
                                );
                            });
                        }
                        let (anno, finished_with_space) =
                            self.show_token_line(&mut state, &mut output, ui);
                        if let Some(seg_name) = &self.seg_name
                            && let Some(seg_key) = self
                                .token
                                .labels
                                .keys()
                                .find(|k| k.name.as_str() == seg_name)
                            && let Some(tok_anno) =
                                anno.iter().find(|a| &a.key == TOKEN_KEY.as_ref())
                        {
                            // Add the token value as segmentation label
                            let seg_anno = Annotation {
                                key: seg_key.clone(),
                                val: tok_anno.val.clone(),
                            };
                            output.changed_annotations.push(seg_anno);
                        }
                        output.changed_annotations.extend(anno);

                        if let Some(anno) = self.show_labels(&mut state, &mut output, ui) {
                            output.changed_annotations.push(anno);
                        }

                        if finished_with_space {
                            output.add_token_after_action = true;
                        }
                    });
                });
                output
            });

            let response = group_response.response.interact(Sense::click());
            let widget_label = if self.selected != SelectionType::None {
                format!(
                    "Selected token ranging from {} to {} ({})",
                    self.token.start, self.token.end, self.token.node_name
                )
            } else {
                format!(
                    "Token ranging from {} to {} ({})",
                    self.token.start, self.token.end, self.token.node_name
                )
            };
            response.widget_info(move || {
                WidgetInfo::labeled(egui::WidgetType::Other, true, widget_label.clone())
            });

            let mut output = group_response.inner;

            if response.clicked() {
                if self.selected != SelectionType::None {
                    state.edited_value = EditedValue::Token(self.value.clone());
                }
                output.selected = true;
            } else if self.selected == SelectionType::Selected
                && state.edited_value == EditedValue::None
                && ui.input_mut(|i| i.consume_key(Modifiers::NONE, Key::Enter))
            {
                state.edited_value = EditedValue::Token(self.value.clone());
            }

            self.store_state(ui.ctx(), state);
            InnerResponse {
                response,
                inner: output,
            }
        }
    }
}