novem/src/screen/dashboard/pull_request_diff_view.rs

use std::{sync::Arc, usize};

use gpui::{
    AppContext, FocusHandle, InteractiveElement, IntoElement, ParentElement, Styled, div,
    prelude::FluentBuilder,
};

use crate::{
    api::{self},
    app,
    component::{
        code_view::symbol_highlight_style,
        diff_view::{DiffViewContent, DiffViewState, diff_view},
        text_input::{self, TextInput, text_input},
    },
    query::{self, QueryStatus, read_query, use_query, watch_query},
    util::{self, diff::DiffLineIndex},
};

pub(crate) struct PullRequestDiffView {
    pr_query: query::Entity<api::issues::FetchPullRequest>,
    content_diff_query: Option<query::Entity<api::repo::FetchFileDiff>>,

    diff_view_state: DiffViewState,
    diff_view_state_in_search_mode: Option<DiffViewState>,
    diff_view_content: Option<DiffViewContent>,
    diff_search_result: Option<DiffSearchResult>,
    diff_search_result_cursor: Option<DiffSearchResultCursor>,

    current_file_path: Option<Arc<str>>,

    focus_handle: gpui::FocusHandle,
    search_input: gpui::Entity<TextInput>,
}

struct DiffSearchHit {
    diff_line_i: DiffLineIndex,
    src_byte_range: std::ops::Range<usize>,
}

struct DiffSearchResult {
    cursor: DiffSearchResultCursor,
    old_side: Vec<DiffSearchHit>,
    new_side: Vec<DiffSearchHit>,
}

struct DiffSearchResultCursor {
    side: util::diff::DiffSide,
    index: usize,
}

gpui::actions!([Search, Escape]);

const KEY_CONTEXT: &'static str = "PullRequestDiffView";

impl PullRequestDiffView {
    pub(crate) fn new(pr_id: api::issues::Id, cx: &mut gpui::Context<Self>) -> Self {
        let mut s = Self {
            pr_query: use_query(api::issues::FetchPullRequest { id: pr_id }, cx),
            content_diff_query: None,
            diff_view_state: DiffViewState::new(),
            diff_view_content: None,
            diff_view_state_in_search_mode: None,
            diff_search_result: None,
            diff_search_result_cursor: None,
            current_file_path: None,

            focus_handle: cx.focus_handle(),
            search_input: cx.new(|cx| TextInput::with_placeholder("Search", cx)),
        };
        s.on_create(cx);
        s
    }

    pub(crate) fn key_bindings() -> [gpui::KeyBinding; 2] {
        [
            gpui::KeyBinding::new("cmd-f", Search, Some(KEY_CONTEXT)),
            gpui::KeyBinding::new("escape", Escape, Some(KEY_CONTEXT)),
        ]
    }

    pub(crate) fn show_diff_for_file(
        &mut self,
        file_path: &Arc<str>,
        cx: &mut gpui::Context<Self>,
    ) {
        self.current_file_path = Some(Arc::clone(file_path));
        self.start_content_queries(cx);
    }

    fn start_content_queries(&mut self, cx: &mut gpui::Context<Self>) {
        if self.content_diff_query.is_some() {
            return;
        }
        let Some(selected_file_path) = self.current_file_path.as_deref() else {
            return;
        };

        let Some((old_file_ref, new_file_ref)) = ({
            if let QueryStatus::Loaded(pr) = read_query(&self.pr_query, cx) {
                Some((
                    api::repo::FileRef {
                        repo_slug: pr.base_repo_slug.clone(),
                        path: Arc::from(selected_file_path),
                        reff: Some(pr.base_ref.clone()),
                    },
                    api::repo::FileRef {
                        repo_slug: pr.head_repo_slug.clone(),
                        path: Arc::from(selected_file_path),
                        reff: Some(pr.head_ref.clone()),
                    },
                ))
            } else {
                None
            }
        }) else {
            return;
        };

        let content_diff_query = use_query(
            api::repo::FetchFileDiff {
                base: old_file_ref,
                head: new_file_ref,
            },
            cx,
        );
        _ = watch_query(&content_diff_query, Self::sync_content_diff_query, cx).detach();

        self.content_diff_query = Some(content_diff_query);
    }

    fn on_create(&mut self, cx: &mut gpui::Context<Self>) {
        _ = cx
            .observe(&self.pr_query, |this, _, cx| {
                this.start_content_queries(cx);
            })
            .detach();

        _ = cx
            .subscribe(&self.search_input, |this, _, event, cx| match event {
            | text_input::Event::Changed(value) => {
                this.search_in_diff(&value, cx);
            }

            | text_input::Event::Submitted(_query) => {
                this.advance_search_result_cursor(cx);
            }
            })
            .detach();
    }

    fn sync_content_diff_query(
        &mut self,
        query: &query::Entity<api::repo::FetchFileDiff>,
        cx: &mut gpui::Context<Self>,
    ) {
        if let Some(diff) = {
            match read_query(query, cx) {
            | QueryStatus::Loaded(diff) => Some(Arc::clone(diff)),
            | _ => None,
            }
        } {
            self.load_diff_view(diff, cx);
            cx.notify();
        }
    }

    fn load_diff_view(
        &mut self,
        content_diff: Arc<util::diff::ContentDiff>,
        cx: &mut gpui::Context<Self>,
    ) {
        let theme = app::current_theme(cx);
        let old_content = content_diff.old_content.clone();
        let new_content = content_diff.new_content.clone();

        self.diff_view_state.reset(content_diff.len());
        self.diff_view_content = Some(content_diff.into());

        let theme_syntax = theme.syntax;

        if let Some(path) = &self.current_file_path {
            let path = Arc::clone(&path);
            let file_type = util::file::file_type_from_path(&path);

            let t1 = cx.background_spawn(async move {
                util::syntax_highlight::highlight_content(old_content, file_type, &theme_syntax)
            });
            let t2 = cx.background_spawn(async move {
                util::syntax_highlight::highlight_content(new_content, file_type, &theme_syntax)
            });

            _ = cx
                .spawn(async move |weak, cx| match tokio::join!(t1, t2) {
                | (Some(old_side_highlights), Some(new_side_highlights)) => {
                    _ = weak.update(cx, |this, cx| {
                        this.diff_view_state
                            .set_old_side_highlights(old_side_highlights);
                        this.diff_view_state
                            .set_new_side_highlights(new_side_highlights);
                        cx.notify();
                    });
                }
                | _ => {}
                })
                .detach();
        }
    }

    fn search_in_diff(&mut self, search_str: &str, cx: &mut gpui::Context<Self>) {
        let Some(diff_view_content) = &self.diff_view_content else {
            return;
        };

        let diff_view_state_in_search_mode = self
            .diff_view_state_in_search_mode
            .get_or_insert(DiffViewState::fork_from(&self.diff_view_state));

        let Some(query) = &self.content_diff_query else {
            return;
        };
        let QueryStatus::Loaded(diff) = read_query(query, cx) else {
            return;
        };

        let theme = app::current_theme(cx);

        let search_str_len = search_str.len();

        let old_search_result = memchr::memmem::find_iter(&diff.old_content, search_str.as_bytes())
            .filter_map(|idx| {
                let range = idx..idx + search_str_len;
                let line_idx = self
                    .diff_view_state
                    .old_side_highlights()?
                    .line_index_of_range(&range)?;
                let diff_line_i = diff_view_content
                    .diff
                    .diff_line_index_for_line(util::diff::DiffSide::Old, line_idx);
                Some(DiffSearchHit {
                    diff_line_i,
                    src_byte_range: range,
                })
            })
            .collect::<Vec<_>>();
        let new_search_result = memchr::memmem::find_iter(&diff.new_content, search_str.as_bytes())
            .filter_map(|idx| {
                let range = idx..idx + search_str_len;
                let line_idx = self
                    .diff_view_state
                    .new_side_highlights()?
                    .line_index_of_range(&range)?;
                let diff_line_i = diff_view_content
                    .diff
                    .diff_line_index_for_line(util::diff::DiffSide::New, line_idx);
                Some(DiffSearchHit {
                    diff_line_i,
                    src_byte_range: range,
                })
            })
            .collect::<Vec<_>>();

        let old_side_highlights =
            old_search_result
                .iter()
                .map(|hit| -> util::syntax_highlight::HighlightedRange {
                    (hit.src_byte_range.clone(), symbol_highlight_style(theme))
                });
        let new_side_highlights =
            new_search_result
                .iter()
                .map(|hit| -> util::syntax_highlight::HighlightedRange {
                    (hit.src_byte_range.clone(), symbol_highlight_style(theme))
                });

        if let Some(h) = self
            .diff_view_state
            .old_side_highlights()
            .map(|it| it.clone())
        {
            diff_view_state_in_search_mode
                .set_old_side_highlights(h.extended_with_highlights(old_side_highlights));
        }

        if let Some(h) = self
            .diff_view_state
            .new_side_highlights()
            .map(|it| it.clone())
        {
            diff_view_state_in_search_mode
                .set_new_side_highlights(h.extended_with_highlights(new_side_highlights));
        }

        if old_search_result.is_empty() && new_search_result.is_empty() {
            self.diff_search_result = None;
        } else {
            let cursor = if !old_search_result.is_empty() {
                DiffSearchResultCursor {
                    side: util::diff::DiffSide::Old,
                    index: 0,
                }
            } else {
                DiffSearchResultCursor {
                    side: util::diff::DiffSide::New,
                    index: 0,
                }
            };
            self.diff_search_result = Some(DiffSearchResult {
                cursor,
                old_side: old_search_result,
                new_side: new_search_result,
            });
        }

        cx.notify();
    }

    fn advance_search_result_cursor(&mut self, cx: &mut gpui::Context<Self>) {
        let (Some(search_result), Some(diff_view_state)) = (
            &mut self.diff_search_result,
            &mut self.diff_view_state_in_search_mode,
        ) else {
            return;
        };

        let (current_side, other_side) = match search_result.cursor.side {
        | util::diff::DiffSide::Old => (&search_result.old_side, &search_result.new_side),
        | util::diff::DiffSide::New => (&search_result.new_side, &search_result.old_side),
        };

        let theme = app::current_theme(cx);

        let current_search_hit = &current_side[search_result.cursor.index];
        let highlight_range = match current_side.get(search_result.cursor.index + 1) {
        | Some(DiffSearchHit {
            diff_line_i,
            src_byte_range,
        }) if *diff_line_i == current_search_hit.diff_line_i => {
            // go to next search result on same side & same line
            search_result.cursor.index += 1;
            Some((
                src_byte_range.clone(),
                gpui::HighlightStyle {
                    background_color: Some(gpui::red()),
                    ..symbol_highlight_style(theme)
                },
            ))
        }
        | next => {
            let next_highlight_line = next
                .map(|hit| hit.diff_line_i)
                .unwrap_or(DiffLineIndex::MAX);

            let mut i = 0;
            let mut other_side_result: Option<(DiffLineIndex, usize, &std::ops::Range<usize>)> =
                None;
            while let Some(DiffSearchHit {
                diff_line_i,
                src_byte_range,
            }) = &other_side.get(i)
            {
                if *diff_line_i == current_search_hit.diff_line_i
                    && matches!(search_result.cursor.side, util::diff::DiffSide::Old)
                {
                    // found other side highlight on the current line
                    // we are on old side, so jump to new side on same line
                    other_side_result = Some((*diff_line_i, i, src_byte_range));
                    break;
                }
                if *diff_line_i > current_search_hit.diff_line_i
                    && *diff_line_i <= next_highlight_line
                {
                    // found other side highlight in between current side highlight and next side highlight
                    other_side_result = Some((*diff_line_i, i, src_byte_range));
                    break;
                }
                if *diff_line_i > next_highlight_line {
                    break;
                }
                i += 1;
            }

            if let Some((_, other_side_i, r)) = other_side_result {
                // next cursor should be on other side with the found index
                search_result.cursor.side = search_result.cursor.side.flipped();
                search_result.cursor.index = other_side_i;
                Some((
                    r.clone(),
                    gpui::HighlightStyle {
                        background_color: Some(gpui::red()),
                        ..symbol_highlight_style(theme)
                    },
                ))
            } else if let Some(next) = next {
                // stay on old side, point to next old side highlight
                search_result.cursor.index = search_result.cursor.index + 1;
                Some((
                    next.src_byte_range.clone(),
                    gpui::HighlightStyle {
                        background_color: Some(gpui::red()),
                        ..symbol_highlight_style(theme)
                    },
                ))
            } else {
                None
            }
        }
        };

        if let Some(highlight_range) = highlight_range {
            let content = match search_result.cursor.side {
            | util::diff::DiffSide::Old => diff_view_state.old_side_highlights_mut(),
            | util::diff::DiffSide::New => diff_view_state.new_side_highlights_mut(),
            };
            if let Some(mut content) = content {
                content.replace_highlight_range(&highlight_range);
            }
        }

        cx.notify();
    }

    fn open_search_box(&mut self, window: &mut gpui::Window, cx: &mut gpui::Context<Self>) {
        self.diff_view_state_in_search_mode = Some(DiffViewState::fork_from(&self.diff_view_state));
        cx.focus_view(&self.search_input, window);
        cx.notify();
    }

    fn close_search_box(&mut self, window: &mut gpui::Window, cx: &mut gpui::Context<Self>) {
        self.diff_view_state_in_search_mode = None;
        self.focus_handle.focus(window);
        cx.notify();
    }
}

impl gpui::Focusable for PullRequestDiffView {
    fn focus_handle(&self, _: &gpui::App) -> FocusHandle {
        self.focus_handle.clone()
    }
}

impl gpui::Render for PullRequestDiffView {
    fn render(
        &mut self,
        _window: &mut gpui::Window,
        cx: &mut gpui::prelude::Context<Self>,
    ) -> impl gpui::IntoElement {
        let theme = app::current_theme(cx);
        let content_diff = self
            .content_diff_query
            .as_ref()
            .map(|q| read_query(q, cx))
            .unwrap_or(QueryStatus::Loading);

        let is_search_input_visible = self.diff_view_state_in_search_mode.is_some();

        div()
            .id(KEY_CONTEXT)
            .key_context(KEY_CONTEXT)
            .track_focus(&self.focus_handle)
            .flex()
            .flex_col()
            .size_full()
            .when(is_search_input_visible, |it| {
                it.child(
                    text_input(self.search_input.clone())
                        .w_full()
                        .text_color(theme.colors.text)
                        .text_xs()
                        .px_3()
                        .bg(theme.colors.surface_chrome)
                        .border_0()
                        .border_b_1()
                        .border_color(theme.colors.border_muted),
                )
            })
            .when_some(
                match (content_diff, &self.diff_view_content) {
                | (QueryStatus::Loaded(_), Some(content)) => Some(content.clone()),
                | (_, _) => None,
                },
                |it, content| {
                    it.child(
                        diff_view(
                            self.diff_view_state_in_search_mode
                                .as_ref()
                                .unwrap_or_else(|| &self.diff_view_state)
                                .clone(),
                            content.clone(),
                        )
                        .into_any_element(),
                    )
                },
            )
            .on_action(cx.listener(|this, _: &Search, window, cx| {
                this.open_search_box(window, cx);
            }))
            .on_action(cx.listener(|this, _: &Escape, window, cx| {
                this.close_search_box(window, cx);
            }))
    }
}