ZIm/crates/assistant/src/slash_command/docs_command.rs

use std::path::Path;
use std::sync::atomic::AtomicBool;
use std::sync::Arc;
use std::time::Duration;

use anyhow::{anyhow, bail, Result};
use assistant_slash_command::{
    ArgumentCompletion, SlashCommand, SlashCommandOutput, SlashCommandOutputSection,
};
use gpui::{AppContext, BackgroundExecutor, Model, Task, WeakView};
use indexed_docs::{
    DocsDotRsProvider, IndexedDocsRegistry, IndexedDocsStore, LocalRustdocProvider, PackageName,
    ProviderId,
};
use language::LspAdapterDelegate;
use project::{Project, ProjectPath};
use ui::prelude::*;
use util::{maybe, ResultExt};
use workspace::Workspace;

pub(crate) struct DocsSlashCommand;

impl DocsSlashCommand {
    pub const NAME: &'static str = "docs";

    fn path_to_cargo_toml(project: Model<Project>, cx: &mut AppContext) -> Option<Arc<Path>> {
        let worktree = project.read(cx).worktrees(cx).next()?;
        let worktree = worktree.read(cx);
        let entry = worktree.entry_for_path("Cargo.toml")?;
        let path = ProjectPath {
            worktree_id: worktree.id(),
            path: entry.path.clone(),
        };
        Some(Arc::from(
            project.read(cx).absolute_path(&path, cx)?.as_path(),
        ))
    }

    /// Ensures that the indexed doc providers for Rust are registered.
    ///
    /// Ideally we would do this sooner, but we need to wait until we're able to
    /// access the workspace so we can read the project.
    fn ensure_rust_doc_providers_are_registered(
        &self,
        workspace: Option<WeakView<Workspace>>,
        cx: &mut AppContext,
    ) {
        let indexed_docs_registry = IndexedDocsRegistry::global(cx);
        if indexed_docs_registry
            .get_provider_store(LocalRustdocProvider::id())
            .is_none()
        {
            let index_provider_deps = maybe!({
                let workspace = workspace.clone().ok_or_else(|| anyhow!("no workspace"))?;
                let workspace = workspace
                    .upgrade()
                    .ok_or_else(|| anyhow!("workspace was dropped"))?;
                let project = workspace.read(cx).project().clone();
                let fs = project.read(cx).fs().clone();
                let cargo_workspace_root = Self::path_to_cargo_toml(project, cx)
                    .and_then(|path| path.parent().map(|path| path.to_path_buf()))
                    .ok_or_else(|| anyhow!("no Cargo workspace root found"))?;

                anyhow::Ok((fs, cargo_workspace_root))
            });

            if let Some((fs, cargo_workspace_root)) = index_provider_deps.log_err() {
                indexed_docs_registry.register_provider(Box::new(LocalRustdocProvider::new(
                    fs,
                    cargo_workspace_root,
                )));
            }
        }

        if indexed_docs_registry
            .get_provider_store(DocsDotRsProvider::id())
            .is_none()
        {
            let http_client = maybe!({
                let workspace = workspace.ok_or_else(|| anyhow!("no workspace"))?;
                let workspace = workspace
                    .upgrade()
                    .ok_or_else(|| anyhow!("workspace was dropped"))?;
                let project = workspace.read(cx).project().clone();
                anyhow::Ok(project.read(cx).client().http_client().clone())
            });

            if let Some(http_client) = http_client.log_err() {
                indexed_docs_registry
                    .register_provider(Box::new(DocsDotRsProvider::new(http_client)));
            }
        }
    }

    /// Runs just-in-time indexing for a given package, in case the slash command
    /// is run without any entries existing in the index.
    fn run_just_in_time_indexing(
        store: Arc<IndexedDocsStore>,
        key: String,
        package: PackageName,
        executor: BackgroundExecutor,
    ) -> Task<()> {
        executor.clone().spawn(async move {
            let (prefix, needs_full_index) = if let Some((prefix, _)) = key.split_once('*') {
                // If we have a wildcard in the search, we want to wait until
                // we've completely finished indexing so we get a full set of
                // results for the wildcard.
                (prefix.to_string(), true)
            } else {
                (key, false)
            };

            // If we already have some entries, we assume that we've indexed the package before
            // and don't need to do it again.
            let has_any_entries = store
                .any_with_prefix(prefix.clone())
                .await
                .unwrap_or_default();
            if has_any_entries {
                return ();
            };

            let index_task = store.clone().index(package.clone());

            if needs_full_index {
                _ = index_task.await;
            } else {
                loop {
                    executor.timer(Duration::from_millis(200)).await;

                    if store
                        .any_with_prefix(prefix.clone())
                        .await
                        .unwrap_or_default()
                        || !store.is_indexing(&package)
                    {
                        break;
                    }
                }
            }
        })
    }
}

impl SlashCommand for DocsSlashCommand {
    fn name(&self) -> String {
        Self::NAME.into()
    }

    fn description(&self) -> String {
        "insert docs".into()
    }

    fn menu_text(&self) -> String {
        "Insert Documentation".into()
    }

    fn requires_argument(&self) -> bool {
        true
    }

    fn complete_argument(
        self: Arc<Self>,
        query: String,
        _cancel: Arc<AtomicBool>,
        workspace: Option<WeakView<Workspace>>,
        cx: &mut AppContext,
    ) -> Task<Result<Vec<ArgumentCompletion>>> {
        self.ensure_rust_doc_providers_are_registered(workspace, cx);

        let indexed_docs_registry = IndexedDocsRegistry::global(cx);
        let args = DocsSlashCommandArgs::parse(&query);
        let store = args
            .provider()
            .ok_or_else(|| anyhow!("no docs provider specified"))
            .and_then(|provider| IndexedDocsStore::try_global(provider, cx));
        cx.background_executor().spawn(async move {
            fn build_completions(
                provider: ProviderId,
                items: Vec<String>,
            ) -> Vec<ArgumentCompletion> {
                items
                    .into_iter()
                    .map(|item| ArgumentCompletion {
                        label: item.clone(),
                        new_text: format!("{provider} {item}"),
                        run_command: true,
                    })
                    .collect()
            }

            match args {
                DocsSlashCommandArgs::NoProvider => {
                    let providers = indexed_docs_registry.list_providers();
                    if providers.is_empty() {
                        return Ok(vec![ArgumentCompletion {
                            label: "No available docs providers.".to_string(),
                            new_text: String::new(),
                            run_command: false,
                        }]);
                    }

                    Ok(providers
                        .into_iter()
                        .map(|provider| ArgumentCompletion {
                            label: provider.to_string(),
                            new_text: provider.to_string(),
                            run_command: false,
                        })
                        .collect())
                }
                DocsSlashCommandArgs::SearchPackageDocs {
                    provider,
                    package,
                    index,
                } => {
                    let store = store?;

                    if index {
                        // We don't need to hold onto this task, as the `IndexedDocsStore` will hold it
                        // until it completes.
                        drop(store.clone().index(package.as_str().into()));
                    }

                    let items = store.search(package).await;
                    Ok(build_completions(provider, items))
                }
                DocsSlashCommandArgs::SearchItemDocs {
                    provider,
                    item_path,
                    ..
                } => {
                    let store = store?;
                    let items = store.search(item_path).await;
                    Ok(build_completions(provider, items))
                }
            }
        })
    }

    fn run(
        self: Arc<Self>,
        argument: Option<&str>,
        _workspace: WeakView<Workspace>,
        _delegate: Option<Arc<dyn LspAdapterDelegate>>,
        cx: &mut WindowContext,
    ) -> Task<Result<SlashCommandOutput>> {
        let Some(argument) = argument else {
            return Task::ready(Err(anyhow!("missing argument")));
        };

        let args = DocsSlashCommandArgs::parse(argument);
        let executor = cx.background_executor().clone();
        let task = cx.background_executor().spawn({
            let store = args
                .provider()
                .ok_or_else(|| anyhow!("no docs provider specified"))
                .and_then(|provider| IndexedDocsStore::try_global(provider, cx));
            async move {
                let (provider, key) = match args.clone() {
                    DocsSlashCommandArgs::NoProvider => bail!("no docs provider specified"),
                    DocsSlashCommandArgs::SearchPackageDocs {
                        provider, package, ..
                    } => (provider, package),
                    DocsSlashCommandArgs::SearchItemDocs {
                        provider,
                        item_path,
                        ..
                    } => (provider, item_path),
                };

                let store = store?;

                if let Some(package) = args.package() {
                    Self::run_just_in_time_indexing(store.clone(), key.clone(), package, executor)
                        .await;
                }

                let (text, ranges) = if let Some((prefix, _)) = key.split_once('*') {
                    let docs = store.load_many_by_prefix(prefix.to_string()).await?;

                    let mut text = String::new();
                    let mut ranges = Vec::new();

                    for (key, docs) in docs {
                        let prev_len = text.len();

                        text.push_str(&docs.0);
                        text.push_str("\n");
                        ranges.push((key, prev_len..text.len()));
                        text.push_str("\n");
                    }

                    (text, ranges)
                } else {
                    let item_docs = store.load(key.clone()).await?;
                    let text = item_docs.to_string();
                    let range = 0..text.len();

                    (text, vec![(key, range)])
                };

                anyhow::Ok((provider, text, ranges))
            }
        });

        cx.foreground_executor().spawn(async move {
            let (provider, text, ranges) = task.await?;
            Ok(SlashCommandOutput {
                text,
                sections: ranges
                    .into_iter()
                    .map(|(key, range)| SlashCommandOutputSection {
                        range,
                        icon: IconName::FileDoc,
                        label: format!("docs ({provider}): {key}",).into(),
                    })
                    .collect(),
                run_commands_in_text: false,
            })
        })
    }
}

fn is_item_path_delimiter(char: char) -> bool {
    !char.is_alphanumeric() && char != '-' && char != '_'
}

#[derive(Debug, PartialEq, Clone)]
pub(crate) enum DocsSlashCommandArgs {
    NoProvider,
    SearchPackageDocs {
        provider: ProviderId,
        package: String,
        index: bool,
    },
    SearchItemDocs {
        provider: ProviderId,
        package: String,
        item_path: String,
    },
}

impl DocsSlashCommandArgs {
    pub fn parse(argument: &str) -> Self {
        let Some((provider, argument)) = argument.split_once(' ') else {
            return Self::NoProvider;
        };

        let provider = ProviderId(provider.into());

        if let Some((package, rest)) = argument.split_once(is_item_path_delimiter) {
            if rest.trim().is_empty() {
                Self::SearchPackageDocs {
                    provider,
                    package: package.to_owned(),
                    index: true,
                }
            } else {
                Self::SearchItemDocs {
                    provider,
                    package: package.to_owned(),
                    item_path: argument.to_owned(),
                }
            }
        } else {
            Self::SearchPackageDocs {
                provider,
                package: argument.to_owned(),
                index: false,
            }
        }
    }

    pub fn provider(&self) -> Option<ProviderId> {
        match self {
            Self::NoProvider => None,
            Self::SearchPackageDocs { provider, .. } | Self::SearchItemDocs { provider, .. } => {
                Some(provider.clone())
            }
        }
    }

    pub fn package(&self) -> Option<PackageName> {
        match self {
            Self::NoProvider => None,
            Self::SearchPackageDocs { package, .. } | Self::SearchItemDocs { package, .. } => {
                Some(package.as_str().into())
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_parse_docs_slash_command_args() {
        assert_eq!(
            DocsSlashCommandArgs::parse(""),
            DocsSlashCommandArgs::NoProvider
        );
        assert_eq!(
            DocsSlashCommandArgs::parse("rustdoc"),
            DocsSlashCommandArgs::NoProvider
        );

        assert_eq!(
            DocsSlashCommandArgs::parse("rustdoc "),
            DocsSlashCommandArgs::SearchPackageDocs {
                provider: ProviderId("rustdoc".into()),
                package: "".into(),
                index: false
            }
        );
        assert_eq!(
            DocsSlashCommandArgs::parse("gleam "),
            DocsSlashCommandArgs::SearchPackageDocs {
                provider: ProviderId("gleam".into()),
                package: "".into(),
                index: false
            }
        );

        assert_eq!(
            DocsSlashCommandArgs::parse("rustdoc gpui"),
            DocsSlashCommandArgs::SearchPackageDocs {
                provider: ProviderId("rustdoc".into()),
                package: "gpui".into(),
                index: false,
            }
        );
        assert_eq!(
            DocsSlashCommandArgs::parse("gleam gleam_stdlib"),
            DocsSlashCommandArgs::SearchPackageDocs {
                provider: ProviderId("gleam".into()),
                package: "gleam_stdlib".into(),
                index: false
            }
        );

        // Adding an item path delimiter indicates we can start indexing.
        assert_eq!(
            DocsSlashCommandArgs::parse("rustdoc gpui:"),
            DocsSlashCommandArgs::SearchPackageDocs {
                provider: ProviderId("rustdoc".into()),
                package: "gpui".into(),
                index: true,
            }
        );
        assert_eq!(
            DocsSlashCommandArgs::parse("gleam gleam_stdlib/"),
            DocsSlashCommandArgs::SearchPackageDocs {
                provider: ProviderId("gleam".into()),
                package: "gleam_stdlib".into(),
                index: true
            }
        );

        assert_eq!(
            DocsSlashCommandArgs::parse("rustdoc gpui::foo::bar::Baz"),
            DocsSlashCommandArgs::SearchItemDocs {
                provider: ProviderId("rustdoc".into()),
                package: "gpui".into(),
                item_path: "gpui::foo::bar::Baz".into()
            }
        );
        assert_eq!(
            DocsSlashCommandArgs::parse("gleam gleam_stdlib/gleam/int"),
            DocsSlashCommandArgs::SearchItemDocs {
                provider: ProviderId("gleam".into()),
                package: "gleam_stdlib".into(),
                item_path: "gleam_stdlib/gleam/int".into()
            }
        );
    }
}