logviewer/src/tui/widgets/fieldtree.rs

use std::borrow::Cow;

use dumpster::sync::Gc;
use itertools::Itertools;
use logparse::{self as lp, Config, SpanKind, into_spans, parse_input};
use ratatui::{
    buffer::Buffer,
    layout::{Constraint, Layout, Rect, Spacing},
    text::{Line, Span},
    widgets::{Paragraph, StatefulWidget, Widget, Wrap},
};
use ratatui_themes::Style;
use tui_widget_list::{ListBuilder, ListState, ListView};

use crate::tui::{
    block_around,
    log_viewer::LogViewer,
    model::{LogEntry, LogFields, SpanDescriptor},
    widgets::{line_text::style_span, styled::Styled},
};

#[derive(Default)]
pub struct State {
    spans_focussed: bool,

    current_span: ListState,
    current_field: ListState,
}

impl State {
    pub fn home(&mut self) {
        if self.spans_focussed {
            self.current_span.select(Some(0));
        } else {
            self.current_field.select(Some(0));
        }
    }

    pub fn down(&mut self) {
        if self.spans_focussed {
            self.current_span.next();
        } else {
            self.current_field.next();
        }
    }
    pub fn up(&mut self) {
        if self.spans_focussed {
            self.current_span.previous();
        } else {
            self.current_field.previous();
        }
    }
    pub fn right(&mut self) {
        self.spans_focussed = false;
    }
    pub fn left(&mut self) {
        self.spans_focussed = true;
    }

    pub fn get(&self) -> Option<(usize, usize)> {
        self.current_span.selected.zip(self.current_field.selected)
    }
}

pub struct FieldTree<'a> {
    focussed: bool,
    lv: &'a mut LogViewer,
}

impl<'a> FieldTree<'a> {
    pub fn new(lv: &'a mut LogViewer, focussed: bool) -> Self {
        if lv.field_state.current_span.selected.is_none() {
            lv.field_state.current_span.select(Some(0));
        }
        if lv.field_state.current_field.selected.is_none() {
            lv.field_state.current_field.select(Some(0));
        }

        Self { focussed, lv }
    }

    fn state(&self) -> &State {
        &self.lv.field_state
    }

    fn state_mut(&mut self) -> &mut State {
        &mut self.lv.field_state
    }

    fn current_entry(&self) -> Option<Gc<LogEntry>> {
        self.lv.selected().map(|(s, _)| s)
    }

    fn current_span_fields(
        &mut self,
        spans: &[(SpanDescriptor, &LogFields)],
    ) -> Vec<(String, String)> {
        let mut selected = self.state().current_span.selected.unwrap_or(0);
        if selected > spans.len() {
            selected = spans.len().saturating_sub(1);
            self.state_mut().current_span.select(Some(selected));
        }

        spans
            .get(selected)
            .map(|(_, fields)| {
                fields
                    .fields
                    .iter()
                    .map(|(k, v)| (k.clone(), v.clone()))
                    .collect_vec()
            })
            .unwrap_or_default()
    }

    fn spans_focussed(&self) -> bool {
        self.focussed && self.state().spans_focussed
    }

    fn fields_focussed(&self) -> bool {
        self.focussed && !self.state().spans_focussed
    }
}

impl Styled<'_, &mut FieldTree<'_>> {
    fn areas(
        &self,
        area: Rect,
        buf: &mut Buffer,
        name: Option<&SpanDescriptor>,
    ) -> (Rect, Rect, Rect) {
        let [spans_area, fields_area] = Layout::horizontal([
            Constraint::Length(if self.spans_focussed() || !self.focussed {
                25
            } else {
                12
            }),
            Constraint::Fill(1),
        ])
        .spacing(Spacing::Overlap(1))
        .areas(area);

        let spans_area = block_around(
            spans_area,
            buf,
            self.spans_focussed(),
            self.styles,
            Some("spans"),
        );
        let fields_area = block_around(
            fields_area,
            buf,
            self.fields_focussed(),
            self.styles,
            Some(if let Some(i) = name {
                match i {
                    SpanDescriptor::Main => "own fields".to_string(),
                    SpanDescriptor::Span(s) => format!("fields of {s}"),
                    SpanDescriptor::Numbered(idx) => format!("fields of span #{idx}"),
                }
            } else {
                "fields".to_string()
            }),
        );

        let [field_names, field_values] =
            Layout::horizontal([Constraint::Length(15), Constraint::Fill(1)]).areas(fields_area);

        (spans_area, field_names, field_values)
    }
}

impl Widget for Styled<'_, &mut FieldTree<'_>> {
    fn render(mut self, area: Rect, buf: &mut Buffer)
    where
        Self: Sized,
    {
        let self_focussed = self.focussed;

        let entry = self.current_entry();
        let spans = entry
            .as_ref()
            .map(|i| i.spans().named().collect_vec())
            .unwrap_or_default();

        let current_span_name = self
            .state()
            .current_span
            .selected
            .and_then(|i| spans.get(i))
            .map(|(i, _)| i);
        let (spans_area, field_names_area, field_values_area) =
            self.areas(area, buf, current_span_name);

        let spans_list = {
            let builder = ListBuilder::new(|cx| {
                let Some((span_name, _)) = &spans.get(cx.index) else {
                    return (Line::from(""), 1);
                };

                let mut item = Line::from(match span_name {
                    SpanDescriptor::Main => "own fields".to_string(),
                    SpanDescriptor::Span(s) => s.to_string(),
                    SpanDescriptor::Numbered(idx) => format!("span #{idx}"),
                });
                if cx.is_selected && self_focussed {
                    item = item.style(self.styles.highlighted);
                } else {
                    item = item.style(self.styles.default);
                }
                (item, 1)
            });

            ListView::new(builder, spans.len())
        };

        let (values_list, names_list) = {
            let values_width = field_values_area.width;
            let styles = self.styles;

            let current_span_fields = self.current_span_fields(&spans);

            let paragraph = move |field_value: &str, base_style: Style| {
                // TODO: cache this parse
                let spans = parse_input(field_value)
                    .map(|i| {
                        into_spans(
                            i,
                            Config {
                                collapse_space: true,
                                ..Default::default()
                            },
                        )
                    })
                    .unwrap_or_else(|_| {
                        vec![lp::Span {
                            text: Cow::Borrowed(field_value),
                            kind: SpanKind::Text,
                        }]
                    });

                let spans = spans
                    .into_iter()
                    .map(|lp::Span { text, kind }| {
                        let span = Span::from(text.into_owned());

                        let style = style_span(kind, base_style, styles);
                        span.style(style)
                    })
                    .collect_vec();

                Paragraph::new(Line::from(spans)).wrap(Wrap { trim: true })
            };

            let num_fields = current_span_fields.len();

            let values_builder = ListBuilder::new({
                let current_span_fields = current_span_fields.clone();
                move |cx| {
                    let Some((_, field_value)) = &current_span_fields.get(cx.index) else {
                        return (Paragraph::new(Span::from("")), 1);
                    };

                    let base_style = if cx.is_selected && self_focussed {
                        self.styles.highlighted
                    } else {
                        self.styles.default
                    };

                    let item = paragraph(field_value, base_style);
                    assert_eq!(cx.cross_axis_size, values_width);
                    let height = item.line_count(values_width);

                    (item, height as u16)
                }
            });

            let names_builder = ListBuilder::new({
                move |cx| {
                    let Some((name, value)) = &current_span_fields.get(cx.index) else {
                        return (Line::from(""), 1);
                    };

                    let item = paragraph(value, self.styles.default);
                    let height = item.line_count(values_width);

                    let mut item = Line::from(name.clone());
                    if cx.is_selected && self_focussed {
                        item = item.style(self.styles.highlighted);
                    } else {
                        item = item.style(self.styles.default);
                    }
                    (item, height as u16)
                }
            });

            (
                ListView::new(values_builder, num_fields),
                ListView::new(names_builder, num_fields),
            )
        };

        spans_list.render(spans_area, buf, &mut self.state_mut().current_span);
        values_list.render(field_values_area, buf, &mut self.state_mut().current_field);
        names_list.render(field_names_area, buf, &mut self.state_mut().current_field);
    }
}