intelli_shell/service/

tldr.rs

use color_eyre::{
    Report,
    eyre::{Context, OptionExt, eyre},
};
use futures_util::{FutureExt, StreamExt, TryStreamExt, stream};
use git2::{
    FetchOptions, ProxyOptions, RemoteCallbacks, Repository,
    build::{CheckoutBuilder, RepoBuilder},
};
use tokio::{fs::File, sync::mpsc};
use tokio_util::sync::CancellationToken;
use tracing::instrument;
use walkdir::WalkDir;

use super::{IntelliShellService, import::parse_import_items};
use crate::{
    errors::{Result, UserFacingError},
    model::{ImportStats, SOURCE_TLDR, TldrConnectionMode},
};

/// Progress events for the `tldr fetch` operation
#[derive(Debug)]
pub enum TldrFetchProgress {
    /// Indicates the status of the tldr git repository
    Repository(RepoStatus),
    /// Indicates that the tldr command files are being located
    LocatingFiles,
    /// Indicates that the tldr command files have been located
    FilesLocated(u64),
    /// Indicates the start of the file processing stage
    ProcessingStart(u64),
    /// Indicates that a single file is being processed
    ProcessingFile(String),
    /// Indicates that a single file has been processed
    FileProcessed(String),
}

/// The status of the tldr git repository
#[derive(Debug)]
pub enum RepoStatus {
    /// Cloning the repository for the first time
    Cloning,
    /// The repository has been successfully cloned
    DoneCloning,
    /// Fetching latest changes
    Fetching,
    /// The repository is already up-to-date
    UpToDate,
    /// Updating the local repository
    Updating,
    /// The repository has been successfully updated
    DoneUpdating,
}

impl IntelliShellService {
    /// Removes tldr commands matching the given criteria.
    ///
    /// Returns the number of commands removed
    #[instrument(skip_all)]
    pub async fn clear_tldr_commands(&self, category: Option<String>) -> Result<u64> {
        self.storage.delete_tldr_commands(category).await
    }

    /// Fetches and imports tldr commands matching the given criteria
    #[instrument(skip_all)]
    pub async fn fetch_tldr_commands(
        &self,
        category: Option<String>,
        connection_mode: TldrConnectionMode,
        commands: Vec<String>,
        progress: mpsc::Sender<TldrFetchProgress>,
        cancellation_token: CancellationToken,
    ) -> Result<ImportStats> {
        // Check for cancellation at the beginning
        if cancellation_token.is_cancelled() {
            tracing::info!("TLDR fetch cancelled before starting");
            return Err(UserFacingError::Cancelled.into());
        }

        // Setup repository
        self.setup_tldr_repo(connection_mode, progress.clone(), cancellation_token.clone())
            .await?;

        // Determine which categories to import
        let categories = if let Some(cat) = category {
            vec![cat]
        } else {
            vec![
                "common".to_owned(),
                #[cfg(target_os = "windows")]
                "windows".to_owned(),
                #[cfg(target_os = "android")]
                "android".to_owned(),
                #[cfg(target_os = "macos")]
                "osx".to_owned(),
                #[cfg(target_os = "freebsd")]
                "freebsd".to_owned(),
                #[cfg(target_os = "openbsd")]
                "openbsd".to_owned(),
                #[cfg(target_os = "netbsd")]
                "netbsd".to_owned(),
                #[cfg(any(
                    target_os = "linux",
                    target_os = "freebsd",
                    target_os = "openbsd",
                    target_os = "netbsd",
                    target_os = "dragonfly",
                ))]
                "linux".to_owned(),
            ]
        };

        // Construct the path to the tldr pages directory
        let pages_path = self.tldr_repo_path.join("pages");

        tracing::info!("Locating files for categories: {}", categories.join(", "));
        progress.send(TldrFetchProgress::LocatingFiles).await.ok();

        // Iterate over directory entries
        let mut command_files = Vec::new();
        let mut iter = WalkDir::new(&pages_path).max_depth(2).into_iter();
        while let Some(result) = iter.next() {
            // Check for cancellation within the file discovery loop
            if cancellation_token.is_cancelled() {
                tracing::info!("TLDR fetch cancelled during file discovery");
                return Err(UserFacingError::Cancelled.into());
            }

            let entry = result.wrap_err("Couldn't read tldr repository files")?;
            let path = entry.path();

            // Skip base path
            if path == pages_path {
                continue;
            }

            // Skip non-included categories
            let file_name = entry.file_name().to_str().ok_or_eyre("Non valid file name")?;
            if entry.file_type().is_dir() {
                if !categories.iter().any(|c| c == file_name) {
                    tracing::trace!("Skipped directory: {file_name}");
                    iter.skip_current_dir();
                    continue;
                } else {
                    // The directory entry itself must be skipped as well, we only care about files
                    continue;
                }
            }

            // We only care about markdown files
            let Some(file_name_no_ext) = file_name.strip_suffix(".md") else {
                tracing::warn!("Unexpected file found: {}", path.display());
                continue;
            };

            // Skip non-included commands
            if !commands.is_empty() {
                if !commands.iter().any(|c| c == file_name_no_ext) {
                    continue;
                } else {
                    tracing::trace!("Included command: {file_name_no_ext}");
                }
            }

            // Retrieve the category
            let category = path
                .parent()
                .and_then(|p| p.file_name())
                .and_then(|p| p.to_str())
                .ok_or_eyre("Couldn't read tldr category")?
                .to_owned();

            // Include the command file
            command_files.push((path.to_path_buf(), category, file_name_no_ext.to_owned()));
        }

        progress
            .send(TldrFetchProgress::FilesLocated(command_files.len() as u64))
            .await
            .ok();

        tracing::info!("Found {} files to be processed", command_files.len());

        progress
            .send(TldrFetchProgress::ProcessingStart(command_files.len() as u64))
            .await
            .ok();

        // Create a stream that reads and parses each command file concurrently
        let items_stream = stream::iter(command_files)
            .map(move |(path, category, command)| {
                let progress = progress.clone();
                async move {
                    progress
                        .send(TldrFetchProgress::ProcessingFile(command.clone()))
                        .await
                        .ok();

                    // Open and parse the file
                    let file = File::open(&path)
                        .await
                        .wrap_err_with(|| format!("Failed to open tldr file: {}", path.display()))?;
                    let stream = parse_import_items(file, vec![], category, SOURCE_TLDR);

                    progress.send(TldrFetchProgress::FileProcessed(command)).await.ok();
                    Ok::<_, Report>(stream)
                }
            })
            .buffered(5)
            .try_flatten();

        // Import items while the token is not cancelled
        let stats = self
            .storage
            .import_items(
                items_stream.take_until(cancellation_token.clone().cancelled_owned().fuse()),
                true,
                false,
            )
            .await?;

        // After processing, check if cancellation was the reason the stream ended
        if cancellation_token.is_cancelled() {
            tracing::info!("TLDR fetch cancelled during command processing");
            return Err(UserFacingError::Cancelled.into());
        }

        Ok(stats)
    }

    #[instrument(skip_all)]
    async fn setup_tldr_repo(
        &self,
        connection_mode: TldrConnectionMode,
        progress: mpsc::Sender<TldrFetchProgress>,
        cancellation_token: CancellationToken,
    ) -> Result<bool> {
        const BRANCH: &str = "main";
        const HTTPS_URL: &str = "https://github.com/tldr-pages/tldr.git";
        const SSH_URL: &str = "git@github.com:tldr-pages/tldr.git";

        let tldr_repo_path = self.tldr_repo_path.clone();

        tokio::task::spawn_blocking(move || {
            // Helper to send progress to the channel
            let send_progress = |status| {
                // Use blocking_send as we are in a sync context
                progress.blocking_send(TldrFetchProgress::Repository(status)).ok();
            };

            // Open the repository upfront when it already exists, so the connection mode can be resolved from it
            let repo = if tldr_repo_path.exists() {
                Some(Repository::open(&tldr_repo_path).wrap_err("Failed to open existing tldr repository")?)
            } else {
                None
            };

            // Resolve the effective repository URL from the requested connection mode:
            // - `Https`/`Ssh` force the matching transport
            // - `Auto` reuses the transport of an existing clone (this also honors git `insteadOf` rewrites), falling
            //   back to HTTPS for a fresh clone
            let repo_url = match connection_mode {
                TldrConnectionMode::Https => HTTPS_URL.to_owned(),
                TldrConnectionMode::Ssh => SSH_URL.to_owned(),
                TldrConnectionMode::Auto => repo
                    .as_ref()
                    .and_then(|repo| repo.find_remote("origin").ok())
                    .and_then(|remote| remote.url().ok().map(ToOwned::to_owned))
                    .unwrap_or_else(|| HTTPS_URL.to_owned()),
            };
            // Setup git callbacks to enable cancellation
            let mut callbacks = RemoteCallbacks::new();
            callbacks.transfer_progress(move |_| !cancellation_token.is_cancelled());
            // Always register the SSH credentials callback: libgit2 only invokes it when the resolved transport
            // is SSH, which can happen via an explicit `ssh` mode, a reused SSH remote, or a git `insteadOf`
            // rewrite of an HTTPS URL (so sniffing the configured scheme isn't reliable). SSH auth is negotiated
            // in steps: libgit2 first asks for the username (when it isn't part of the URL), then for the
            // credentials, so the returned credential type must match the requested `allowed_types`, otherwise
            // git2 passes through and libgit2 reports it as a missing callback. The agent is only offered once to
            // avoid looping forever when authentication fails.
            let mut agent_offered = false;
            callbacks.credentials(move |_url, username_from_url, allowed_types| {
                let username = username_from_url.unwrap_or("git");
                if allowed_types.contains(git2::CredentialType::USERNAME) {
                    git2::Cred::username(username)
                } else if allowed_types.contains(git2::CredentialType::SSH_KEY) && !agent_offered {
                    agent_offered = true;
                    git2::Cred::ssh_key_from_agent(username)
                } else {
                    Err(git2::Error::from_str(
                        "SSH authentication failed: no usable identity found in the SSH agent",
                    ))
                }
            });
            // Setup git fetch options for a swallow copy with auto proxy config
            let mut proxy_opts = ProxyOptions::new();
            proxy_opts.auto();
            let mut fetch_options = FetchOptions::new();
            fetch_options.proxy_options(proxy_opts);
            fetch_options.remote_callbacks(callbacks);
            fetch_options.depth(1);
            // Fetch latest repo changes or clone it if it doesn't exist yet
            if let Some(repo) = repo {
                tracing::info!("Fetching latest tldr changes ...");
                send_progress(RepoStatus::Fetching);

                // Update 'origin' URL if it doesn't match `repo_url` (no-op in `Auto` mode, which already
                // derives the URL from the existing remote)
                let current_url = repo.find_remote("origin")?.url().ok().map(|s| s.to_owned());
                if current_url.as_deref() != Some(repo_url.as_str()) {
                    repo.remote_set_url("origin", &repo_url)
                        .wrap_err("Failed to update remote URL")?;
                }

                // Get the 'origin' remote
                let mut remote = repo.find_remote("origin")?;

                // Fetch the latest changes from the remote 'main' branch
                let refspec = format!("refs/heads/{BRANCH}:refs/remotes/origin/{BRANCH}");
                if let Err(err) = remote.fetch(&[refspec], Some(&mut fetch_options), None) {
                    // Check if the error was a user-initiated cancellation
                    if err.code() == git2::ErrorCode::User && err.class() == git2::ErrorClass::Callback {
                        return Err(UserFacingError::Cancelled.into());
                    }
                    return Err(Report::from(err).wrap_err("Failed to fetch from tldr remote").into());
                }

                // Get the commit OID from the fetched data (FETCH_HEAD)
                let fetch_head = repo.find_reference("FETCH_HEAD")?;
                let fetch_commit_oid = fetch_head
                    .target()
                    .ok_or_else(|| eyre!("FETCH_HEAD is not a direct reference"))?;

                // Get the OID of the current commit on the local branch
                let local_ref_name = format!("refs/heads/{BRANCH}");
                let local_commit_oid = local_branch_target(&repo, &local_ref_name)?;

                // If the commit OIDs are the same, the repo is already up-to-date
                if Some(fetch_commit_oid) == local_commit_oid {
                    tracing::info!("Repository is already up-to-date");
                    send_progress(RepoStatus::UpToDate);
                    return Ok(false);
                }

                tracing::info!("Updating to the latest version ...");
                send_progress(RepoStatus::Updating);

                // Create the local branch when an older or incomplete clone has none; otherwise
                // advance it to the fetched commit.
                let msg = format!("Resetting to latest commit {fetch_commit_oid}");
                update_local_branch(&repo, &local_ref_name, fetch_commit_oid, &msg)?;

                // Point HEAD to the updated local branch
                repo.set_head(&local_ref_name)?;

                // Checkout the new HEAD to update the files in the working directory
                let mut checkout_builder = CheckoutBuilder::new();
                checkout_builder.force();
                repo.checkout_head(Some(&mut checkout_builder))?;

                tracing::info!("Repository successfully updated");
                send_progress(RepoStatus::DoneUpdating);
                Ok(true)
            } else {
                tracing::info!("Performing a shallow clone of '{repo_url}' ...");
                send_progress(RepoStatus::Cloning);

                // Clone the repository
                if let Err(err) = RepoBuilder::new()
                    .branch(BRANCH)
                    .fetch_options(fetch_options)
                    .clone(&repo_url, &tldr_repo_path)
                {
                    if err.code() == git2::ErrorCode::User && err.class() == git2::ErrorClass::Callback {
                        return Err(UserFacingError::Cancelled.into());
                    }
                    return Err(Report::from(err).wrap_err("Failed to clone tldr repository").into());
                }

                tracing::info!("Repository successfully cloned");
                send_progress(RepoStatus::DoneCloning);
                Ok(true)
            }
        })
        .await
        .wrap_err("tldr repository task failed")?
    }
}

fn update_local_branch(
    repo: &Repository,
    local_ref_name: &str,
    commit_oid: git2::Oid,
    message: &str,
) -> std::result::Result<(), git2::Error> {
    repo.reference(local_ref_name, commit_oid, true, message)?;
    Ok(())
}

fn local_branch_target(repo: &Repository, local_ref_name: &str) -> std::result::Result<Option<git2::Oid>, git2::Error> {
    match repo.find_reference(local_ref_name) {
        Ok(local_ref) => Ok(local_ref.target()),
        Err(err) if err.code() == git2::ErrorCode::NotFound => Ok(None),
        Err(err) => Err(err),
    }
}

#[cfg(test)]
mod tests {
    use std::{error::Error, fs};

    use git2::Repository;
    use uuid::Uuid;

    use super::{local_branch_target, update_local_branch};

    #[test]
    fn setup_can_create_a_missing_local_main_branch() -> Result<(), Box<dyn Error>> {
        let path = std::env::temp_dir().join(format!("intelli-shell-tldr-{}", Uuid::now_v7()));
        let repo = Repository::init(&path)?;
        let blob_oid = repo.blob(b"tldr test")?;
        let mut tree_builder = repo.treebuilder(None)?;
        tree_builder.insert("test.md", blob_oid, 0o100644)?;
        let tree_oid = tree_builder.write()?;
        drop(tree_builder);
        let tree = repo.find_tree(tree_oid)?;
        let signature = git2::Signature::now("Test", "test@example.invalid")?;
        let commit_oid = repo.commit(None, &signature, &signature, "test commit", &tree, &[])?;
        drop(tree);
        let branch_name = "refs/heads/main";

        assert_eq!(local_branch_target(&repo, branch_name)?, None);

        update_local_branch(&repo, branch_name, commit_oid, "test branch creation")?;
        assert_eq!(local_branch_target(&repo, branch_name)?, Some(commit_oid));
        repo.set_head(branch_name)?;
        repo.checkout_head(Some(git2::build::CheckoutBuilder::new().force()))?;
        assert_eq!(fs::read(path.join("test.md"))?, b"tldr test");

        drop(repo);
        fs::remove_dir_all(path)?;
        Ok(())
    }
}