securegit/ops/

worktree.rs

use crate::cli::UI;
use crate::ops::oplog;
use crate::ops::utils::short_oid;
use anyhow::{bail, Result};
use git2::{Repository, WorktreeAddOptions, WorktreeLockStatus, WorktreePruneOptions};
use std::path::{Path, PathBuf};

/// Information about a single worktree for structured output.
#[derive(Debug)]
struct WorktreeInfo {
    name: String,
    path: PathBuf,
    head_oid: String,
    branch: String,
    is_main: bool,
    locked: Option<String>,
}

/// Gather information about all worktrees (main + secondary).
fn gather_worktrees(repo: &Repository) -> Result<Vec<WorktreeInfo>> {
    let mut entries = Vec::new();

    // Main worktree
    let main_path = repo
        .workdir()
        .map(|p| p.to_path_buf())
        .unwrap_or_else(|| repo.path().to_path_buf());

    let head_oid = repo
        .head()
        .ok()
        .and_then(|h| h.target())
        .map(|oid| short_oid(&oid))
        .unwrap_or_else(|| "???????".to_string());

    let branch = repo
        .head()
        .ok()
        .and_then(|h| {
            if h.is_branch() {
                h.shorthand().map(|s| s.to_string())
            } else {
                Some("(detached)".to_string())
            }
        })
        .unwrap_or_else(|| "(detached)".to_string());

    entries.push(WorktreeInfo {
        name: "(main)".to_string(),
        path: main_path,
        head_oid,
        branch,
        is_main: true,
        locked: None,
    });

    // Secondary worktrees
    let worktree_names = repo.worktrees()?;
    for name in worktree_names.iter().flatten() {
        if let Ok(wt) = repo.find_worktree(name) {
            let wt_path = wt.path().to_path_buf();

            // Get HEAD info by opening the worktree as a repo
            let (wt_oid, wt_branch) = if let Ok(wt_repo) = Repository::open(&wt_path) {
                let oid = wt_repo
                    .head()
                    .ok()
                    .and_then(|h| h.target())
                    .map(|oid| short_oid(&oid))
                    .unwrap_or_else(|| "???????".to_string());
                let br = wt_repo
                    .head()
                    .ok()
                    .and_then(|h| {
                        if h.is_branch() {
                            h.shorthand().map(|s| s.to_string())
                        } else {
                            Some("(detached)".to_string())
                        }
                    })
                    .unwrap_or_else(|| "???????".to_string());
                (oid, br)
            } else {
                ("???????".to_string(), "???????".to_string())
            };

            let locked = match wt.is_locked() {
                Ok(WorktreeLockStatus::Locked(reason)) => Some(
                    reason.unwrap_or_else(|| "(no reason)".to_string()),
                ),
                _ => None,
            };

            entries.push(WorktreeInfo {
                name: name.to_string(),
                path: wt_path,
                head_oid: wt_oid,
                branch: wt_branch,
                is_main: false,
                locked,
            });
        }
    }

    Ok(entries)
}

pub fn list(path: &Path, ui: &UI) -> Result<()> {
    let repo = Repository::open(path)?;
    let entries = gather_worktrees(&repo)?;

    for entry in &entries {
        let lock_indicator = if let Some(ref reason) = entry.locked {
            format!(" [locked: {}]", reason)
        } else {
            String::new()
        };

        let label = if entry.is_main {
            format!("{} (main)", entry.path.display())
        } else {
            format!("{}", entry.path.display())
        };

        ui.branch_item(
            &label,
            &entry.head_oid,
            &format!("{}{}", entry.branch, lock_indicator),
            entry.is_main,
        );
    }

    ui.info(format!("{} worktree(s)", entries.len()));
    Ok(())
}

pub fn list_compact(path: &Path) -> Result<String> {
    let repo = Repository::open(path)?;
    let entries = gather_worktrees(&repo)?;

    let mut lines = Vec::new();
    for entry in &entries {
        let marker = if entry.is_main { "*" } else { " " };
        let lock = if entry.locked.is_some() { " [locked]" } else { "" };
        lines.push(format!(
            "{} {} {} {}{}",
            marker,
            entry.path.display(),
            entry.head_oid,
            entry.branch,
            lock
        ));
    }

    Ok(lines.join("\n"))
}

/// List worktrees as JSON (for MCP).
pub fn list_json(path: &Path) -> Result<String> {
    let repo = Repository::open(path)?;
    let entries = gather_worktrees(&repo)?;

    let json_entries: Vec<serde_json::Value> = entries
        .iter()
        .map(|e| {
            serde_json::json!({
                "name": e.name,
                "path": e.path.display().to_string(),
                "head": e.head_oid,
                "branch": e.branch,
                "is_main": e.is_main,
                "locked": e.locked,
            })
        })
        .collect();

    Ok(serde_json::to_string_pretty(&json_entries)?)
}

pub fn add(
    path: &Path,
    name: &str,
    worktree_path: &Path,
    branch: Option<&str>,
    ui: &UI,
) -> Result<()> {
    let desc = format!(
        "worktree add '{}' at '{}'",
        name,
        worktree_path.display()
    );
    oplog::with_oplog(path, "worktree", &desc, || {
        add_inner(path, name, worktree_path, branch, ui)
    })
}

fn add_inner(
    path: &Path,
    name: &str,
    worktree_path: &Path,
    branch: Option<&str>,
    ui: &UI,
) -> Result<()> {
    let repo = Repository::open(path)?;

    // Ensure parent directory exists
    if let Some(parent) = worktree_path.parent() {
        std::fs::create_dir_all(parent)?;
    }

    if worktree_path.exists() {
        bail!(
            "Directory already exists: {}",
            worktree_path.display()
        );
    }

    let mut opts = WorktreeAddOptions::new();

    if let Some(branch_name) = branch {
        // Check if the branch already exists
        if let Ok(branch_ref) = repo.find_branch(branch_name, git2::BranchType::Local) {
            let reference = branch_ref.into_reference();
            opts.reference(Some(&reference));
            opts.checkout_existing(true);
            repo.worktree(name, worktree_path, Some(&opts))?;
            ui.success(format!(
                "Created worktree '{}' at '{}' (existing branch '{}')",
                name,
                worktree_path.display(),
                branch_name
            ));
        } else {
            // Create the branch first, then use it
            let head_commit = repo.head()?.peel_to_commit()?;
            let new_branch = repo.branch(branch_name, &head_commit, false)?;
            let reference = new_branch.into_reference();
            opts.reference(Some(&reference));
            repo.worktree(name, worktree_path, Some(&opts))?;
            ui.success(format!(
                "Created worktree '{}' at '{}' (new branch '{}')",
                name,
                worktree_path.display(),
                branch_name
            ));
        }
    } else {
        // No branch specified — git2 creates a new branch named after the worktree
        repo.worktree(name, worktree_path, Some(&opts))?;
        ui.success(format!(
            "Created worktree '{}' at '{}' (new branch '{}')",
            name,
            worktree_path.display(),
            name
        ));
    }

    Ok(())
}

pub fn remove(path: &Path, name: &str, force: bool, ui: &UI) -> Result<()> {
    let desc = format!("worktree remove '{}'", name);
    oplog::with_oplog(path, "worktree", &desc, || {
        remove_inner(path, name, force, ui)
    })
}

fn remove_inner(path: &Path, name: &str, force: bool, ui: &UI) -> Result<()> {
    let repo = Repository::open(path)?;
    let wt = repo.find_worktree(name)?;
    let wt_path = wt.path().to_path_buf();

    // Check lock status
    if let Ok(WorktreeLockStatus::Locked(reason)) = wt.is_locked() {
        if !force {
            let reason_msg = reason.unwrap_or_else(|| "no reason given".to_string());
            bail!(
                "Worktree '{}' is locked ({}). Use --force to remove anyway.",
                name,
                reason_msg
            );
        }
        // Unlock before removing
        wt.unlock()?;
    }

    // Check for uncommitted changes if not forcing
    if !force && wt_path.exists() {
        if let Ok(wt_repo) = Repository::open(&wt_path) {
            let mut opts = git2::StatusOptions::new();
            opts.include_untracked(true);
            if let Ok(statuses) = wt_repo.statuses(Some(&mut opts)) {
                if !statuses.is_empty() {
                    bail!(
                        "Worktree '{}' has uncommitted changes. Use --force to remove anyway.",
                        name
                    );
                }
            }
        }
    }

    // Prune the worktree (removes admin files + working tree)
    let mut prune_opts = WorktreePruneOptions::new();
    prune_opts.valid(true).working_tree(true);
    if force {
        prune_opts.locked(true);
    }
    wt.prune(Some(&mut prune_opts))?;

    // Remove the directory if it still exists (prune with working_tree should handle this)
    if wt_path.exists() {
        std::fs::remove_dir_all(&wt_path)?;
    }

    ui.success(format!("Removed worktree '{}'", name));
    Ok(())
}

pub fn lock(path: &Path, name: &str, reason: Option<&str>, ui: &UI) -> Result<()> {
    let repo = Repository::open(path)?;
    let wt = repo.find_worktree(name)?;

    if let Ok(WorktreeLockStatus::Locked(_)) = wt.is_locked() {
        bail!("Worktree '{}' is already locked", name);
    }

    wt.lock(reason)?;

    if let Some(r) = reason {
        ui.success(format!("Locked worktree '{}' ({})", name, r));
    } else {
        ui.success(format!("Locked worktree '{}'", name));
    }
    Ok(())
}

pub fn unlock(path: &Path, name: &str, ui: &UI) -> Result<()> {
    let repo = Repository::open(path)?;
    let wt = repo.find_worktree(name)?;

    if let Ok(WorktreeLockStatus::Unlocked) = wt.is_locked() {
        bail!("Worktree '{}' is not locked", name);
    }

    wt.unlock()?;
    ui.success(format!("Unlocked worktree '{}'", name));
    Ok(())
}

pub fn prune(path: &Path, dry_run: bool, ui: &UI) -> Result<()> {
    let repo = Repository::open(path)?;
    let worktree_names = repo.worktrees()?;

    let mut pruned = 0;

    for name in worktree_names.iter().flatten() {
        if let Ok(wt) = repo.find_worktree(name) {
            let mut opts = WorktreePruneOptions::new();
            if let Ok(true) = wt.is_prunable(Some(&mut opts)) {
                if dry_run {
                    ui.info(format!("Would prune: {} ({})", name, wt.path().display()));
                } else {
                    let mut prune_opts = WorktreePruneOptions::new();
                    prune_opts.working_tree(true);
                    match wt.prune(Some(&mut prune_opts)) {
                        Ok(()) => {
                            ui.success(format!("Pruned: {}", name));
                            pruned += 1;
                        }
                        Err(e) => {
                            ui.error(format!("Failed to prune '{}': {}", name, e));
                        }
                    }
                }
            }
        }
    }

    if pruned == 0 && !dry_run {
        ui.info("No stale worktrees to prune.".to_string());
    }

    Ok(())
}