vkit 0.1.4

Fast Rust dev CLI: manage git worktrees, Node ports, run scripts, install & sync VS Code / Cursor extensions.
//! `vkit wt switch`:按分支 / 快捷符快速进入 worktree。

use std::path::{Path, PathBuf};
use std::process::Command;

use anyhow::{anyhow, bail, Context, Result};
use clap::Parser;

use super::config;
use super::git::{self, AddExistingBranch, AddNewBranch, RepoContext};
use super::name;
use super::placement::{self, Placement};
use super::shell;
use super::state;

#[derive(Debug, Clone, Parser)]
#[command(about = "切换到 worktree;不存在时可创建")]
pub struct SwitchArgs {
    /// 分支名,或快捷符:`-` 上一个 · `^` Main · `@` 当前 · `mr:N` GitLab MR
    pub target: Option<String>,

    /// 新建分支并创建 worktree
    #[arg(short = 'c', long)]
    pub create: bool,

    /// 新分支的 base(默认主分支)
    #[arg(short = 'b', long)]
    pub base: Option<String>,

    /// 切换后执行的命令(有 shell 集成时在 cd 后 source;否则直接跑)
    #[arg(short = 'x', long)]
    pub execute: Option<String>,

    /// 创建时使用 Global Placement(默认 In-repo)
    #[arg(long)]
    pub global: bool,
}

pub fn run(args: SwitchArgs) -> Result<()> {
    let cfg = config::load()?;
    let cwd = std::env::current_dir()?;
    let ctx = RepoContext::discover(&cwd, cfg.global_root)?;

    let Some(raw_target) = args.target.as_deref() else {
        bail!("请指定目标,例如:vkit wt switch feat  或  vkit wt switch -");
    };

    let current = git::current_toplevel(&cwd).unwrap_or_else(|_| cwd.clone());
    let resolved = resolve_target(&ctx, &cwd, raw_target, &args)?;
    let path = match resolved {
        Resolved::Existing(path) => path,
        Resolved::Create(kind) => {
            let branch = kind.branch().to_string();
            create_worktree(&ctx, kind, args.global)?;
            git::worktree_for_branch(&ctx.main_path, &branch)?
                .ok_or_else(|| anyhow!("创建后未找到 worktree:{branch}"))?
        }
    };

    if path != current {
        let _ = state::save_previous(&ctx.main_path, &current);
    }

    enter(&path, args.execute.as_deref())
}

enum Resolved {
    Existing(PathBuf),
    Create(CreateKind),
}

enum CreateKind {
    NewBranch { branch: String, base: String },
    CheckoutExisting { branch: String },
    TrackRemote { branch: String, start: String },
}

impl CreateKind {
    fn branch(&self) -> &str {
        match self {
            Self::NewBranch { branch, .. }
            | Self::CheckoutExisting { branch }
            | Self::TrackRemote { branch, .. } => branch,
        }
    }
}

fn resolve_target(
    ctx: &RepoContext,
    cwd: &Path,
    raw: &str,
    args: &SwitchArgs,
) -> Result<Resolved> {
    match raw {
        "-" => {
            let prev = state::load_previous(&ctx.main_path)
                .ok_or_else(|| anyhow!("没有上一个 worktree(尚未 switch 过)"))?;
            Ok(Resolved::Existing(prev))
        }
        "^" => Ok(Resolved::Existing(ctx.main_path.clone())),
        "@" => Ok(Resolved::Existing(git::current_toplevel(cwd)?)),
        other if other.to_ascii_lowercase().starts_with("mr:") => {
            let iid = other[3..]
                .parse::<u64>()
                .map_err(|_| anyhow!("无效的 MR 编号:{other}"))?;
            if args.create {
                bail!("mr:{{N}} 已有分支,不能与 --create 同用");
            }
            let branch = mr_source_branch(&ctx.main_path, iid)?;
            if let Some(path) = git::worktree_for_branch(&ctx.main_path, &branch)? {
                return Ok(Resolved::Existing(path));
            }
            Ok(Resolved::Create(CreateKind::CheckoutExisting { branch }))
        }
        branch if args.create => {
            if git::worktree_for_branch(&ctx.main_path, branch)?.is_some() {
                bail!("分支 {branch} 已有 worktree,去掉 --create 即可 switch");
            }
            let base = args
                .base
                .clone()
                .or_else(|| git::default_base_branch(&ctx.main_path).ok())
                .unwrap_or_else(|| "main".into());
            Ok(Resolved::Create(CreateKind::NewBranch {
                branch: branch.to_string(),
                base,
            }))
        }
        branch => {
            if let Some(path) = git::worktree_for_branch(&ctx.main_path, branch)? {
                return Ok(Resolved::Existing(path));
            }
            let locals = git::list_local_branches(&ctx.main_path).unwrap_or_default();
            if locals.iter().any(|b| b == branch) {
                return Ok(Resolved::Create(CreateKind::CheckoutExisting {
                    branch: branch.to_string(),
                }));
            }
            let remotes = git::list_remote_branches(&ctx.main_path).unwrap_or_default();
            if let Some(start) = remotes.iter().find(|r| {
                *r == &format!("origin/{branch}") || r.ends_with(&format!("/{branch}"))
            }) {
                return Ok(Resolved::Create(CreateKind::TrackRemote {
                    branch: branch.to_string(),
                    start: start.clone(),
                }));
            }
            bail!("没有分支 `{branch}` 的 worktree。创建:vkit wt switch -c {branch}");
        }
    }
}

fn create_worktree(ctx: &RepoContext, kind: CreateKind, global: bool) -> Result<()> {
    let placement = if global {
        Placement::Global
    } else {
        Placement::InRepo
    };
    let branch = kind.branch().to_string();
    let wt_name = name::from_branch(&branch);
    if wt_name.is_empty() {
        bail!("无法从分支名派生 Worktree Name");
    }
    let path = placement::worktree_path(
        placement,
        &ctx.main_path,
        &ctx.global_root,
        &ctx.repo_name,
        &wt_name,
    );
    if path.exists() {
        bail!("目标路径已存在:{}", path.display());
    }

    match kind {
        CreateKind::NewBranch { branch, base } => {
            git::add_new_branch(
                &ctx.main_path,
                &AddNewBranch {
                    path,
                    branch,
                    base,
                },
            )?;
        }
        CreateKind::CheckoutExisting { branch } => {
            git::add_existing_branch(
                &ctx.main_path,
                &AddExistingBranch {
                    path,
                    branch,
                    start_point: None,
                },
            )?;
        }
        CreateKind::TrackRemote { branch, start } => {
            git::add_existing_branch(
                &ctx.main_path,
                &AddExistingBranch {
                    path,
                    branch,
                    start_point: Some(start),
                },
            )?;
        }
    }
    Ok(())
}

fn mr_source_branch(cwd: &Path, iid: u64) -> Result<String> {
    let output = Command::new("glab")
        .args(["mr", "view", &iid.to_string(), "-F", "json"])
        .current_dir(cwd)
        .output()
        .context("执行 glab mr view 失败")?;
    if !output.status.success() {
        let stderr = String::from_utf8_lossy(&output.stderr);
        bail!("glab mr view {iid} 失败:{}", stderr.trim());
    }
    let v: serde_json::Value =
        serde_json::from_slice(&output.stdout).context("解析 glab mr view JSON 失败")?;
    v.get("source_branch")
        .and_then(|x| x.as_str())
        .map(str::to_string)
        .ok_or_else(|| anyhow!("MR !{iid} 响应缺少 source_branch"))
}

fn enter(path: &Path, execute: Option<&str>) -> Result<()> {
    let wrote_cd = shell::write_cd(path)?;

    if let Some(cmd) = execute {
        let script = format_exec_script(cmd, path);
        if !shell::write_exec(&script)? {
            let status = Command::new("sh")
                .arg("-c")
                .arg(format!(
                    "cd {} && {cmd}",
                    shell_quote(&path.display().to_string())
                ))
                .status()
                .with_context(|| format!("执行失败:{cmd}"))?;
            if !status.success() {
                bail!("命令退出码:{status}");
            }
        }
    }

    if wrote_cd {
        eprintln!("{}", path.display());
    } else if execute.is_none() {
        eprintln!(
            "{}\n提示:eval \"$(vkit wt shell-init zsh)\" 后用 vwt switch … 可自动 cd",
            path.display()
        );
    }
    Ok(())
}

fn format_exec_script(cmd: &str, path: &Path) -> String {
    let trimmed = cmd.trim();
    if matches!(trimmed, "cursor" | "code" | "nvim" | "vim") {
        format!("{trimmed} {}", shell_quote(&path.display().to_string()))
    } else {
        trimmed.to_string()
    }
}

fn shell_quote(s: &str) -> String {
    format!("'{}'", s.replace('\'', "'\\''"))
}

/// TUI「进入」:有 shell 集成则写 CD 并请求退出;否则开子 shell。
pub fn enter_from_tui(path: &Path) -> Result<bool> {
    if shell::write_cd(path)? {
        Ok(true) // 调用方应 quit
    } else {
        Ok(false)
    }
}

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

    #[test]
    fn shell_quote_escapes() {
        assert_eq!(shell_quote("a'b"), "'a'\\''b'");
    }
}