1use std::fmt::Display;
2use std::fmt::Formatter;
3use std::ops::Deref;
4use std::path::Path;
5
6use owo_colors::OwoColorize;
7use rootcause::prelude::ResultExt;
8use rootcause::report;
9use ytil_git::branch::Branch;
10use ytil_sys::cli::Args;
11
12const ZSHRC_INSTALL_LINE: &str = r#"(( $+commands[gbm] )) && eval "$(gbm init zsh)""#;
13const ZSH_WRAPPER: &str = r#"gbm() {
14 if (( $# == 0 )); then
15 local branch
16 branch="$(command gbm --pick)" || return
17 [[ -n "$branch" ]] || return
18 print -z -- "gbm ${(q)branch}"
19 return
20 fi
21
22 command gbm "$@"
23}
24"#;
25
26#[ytil_sys::main]
28fn main() -> rootcause::Result<()> {
29 let args = ytil_sys::cli::get();
30 if args.has_help() {
31 println!("{}", include_str!("../help.txt"));
32 return Ok(());
33 }
34
35 let args: Vec<_> = args.iter().map(String::as_str).collect();
36 match args.as_slice() {
37 [] => Err(report!("gbm shell wrapper is not installed or loaded")
38 .attach("run `gbm install` first, then restart zsh or source ~/.zshrc")),
39 ["--pick"] => pick_git_branch(),
40 ["install"] => install_zsh_wrapper(),
41 ["init", "zsh"] => {
42 print!("{ZSH_WRAPPER}");
43 Ok(())
44 }
45 [branch_name] => rename_current_branch(branch_name),
46 _ => rename_current_branch(&args.join("-")),
47 }
48}
49
50fn pick_git_branch() -> rootcause::Result<()> {
51 let Some(branch) = select_branch_with_current_first()? else {
52 return Ok(());
53 };
54
55 println!("{}", branch.name_no_origin());
56 Ok(())
57}
58
59fn install_zsh_wrapper() -> rootcause::Result<()> {
60 let zshrc = std::env::var("HOME")
61 .context("error missing HOME environment variable")
62 .map(|home| Path::new(&home).join(".zshrc"))?;
63
64 install_zsh_wrapper_at(&zshrc)?;
65 println!("{} gbm in {}", "Patched".green().bold(), zshrc.display());
66
67 Ok(())
68}
69
70fn install_zsh_wrapper_at(path: &Path) -> rootcause::Result<bool> {
71 let content = std::fs::read_to_string(path)
72 .context("error reading zshrc")
73 .attach_with(|| format!("path={}", path.display()))?;
74
75 if content.lines().any(|line| line.trim() == ZSHRC_INSTALL_LINE) {
76 return Ok(false);
77 }
78
79 let mut updated = content;
80 if !updated.is_empty() && !updated.ends_with('\n') {
81 updated.push('\n');
82 }
83 updated.push_str(ZSHRC_INSTALL_LINE);
84 updated.push('\n');
85
86 std::fs::write(path, updated)
87 .context("error installing zshrc")
88 .attach_with(|| format!("path={}", path.display()))?;
89
90 Ok(true)
91}
92
93fn select_branch_with_current_first() -> rootcause::Result<Option<Branch>> {
94 let repo = ytil_git::repo::discover(Path::new(".")).context("error discovering repo for branch selection")?;
95 let branches = prioritize_current_branch_first(
96 ytil_git::branch::get_all_no_redundant(&repo)?,
97 ytil_git::branch::get_current()?.as_str(),
98 ytil_git::branch::get_previous(&repo).as_deref(),
99 ytil_git::branch::get_user_email(&repo)?.as_deref(),
100 );
101
102 let Some(branch) = ytil_tui::minimal_select(branches.into_iter().map(RenderableBranch).collect())? else {
103 return Ok(None);
104 };
105
106 Ok(Some(branch.0))
107}
108
109fn rename_current_branch(branch_name: &str) -> rootcause::Result<()> {
110 ytil_git::branch::rename_current(branch_name, None)?;
111 println!("{} {}", ">".magenta().bold(), branch_name.bold());
112 Ok(())
113}
114
115fn prioritize_current_branch_first(
116 branches: Vec<Branch>,
117 current_branch: &str,
118 previous_branch: Option<&str>,
119 user_email: Option<&str>,
120) -> Vec<Branch> {
121 let branches = prioritize_recent_branches(branches, previous_branch, user_email);
122 let mut current = None;
123 let mut rest = Vec::with_capacity(branches.len());
124
125 for branch in branches {
126 if current.is_none() && branch.name_no_origin() == current_branch {
127 current = Some(branch);
128 } else {
129 rest.push(branch);
130 }
131 }
132
133 current.into_iter().chain(rest).collect()
134}
135
136fn prioritize_recent_branches(
137 branches: Vec<Branch>,
138 previous_branch: Option<&str>,
139 user_email: Option<&str>,
140) -> Vec<Branch> {
141 const MINE_DESIRED_COUNT: usize = 5;
142
143 let branches_len = branches.len();
144 let mut previous = None;
145 let mut mine = Vec::new();
146 let mut rest = Vec::new();
147
148 for branch in branches {
149 if previous.is_none() && previous_branch.is_some_and(|prev| branch.name_no_origin() == prev) {
150 previous = Some(branch);
151 } else if mine.len() < MINE_DESIRED_COUNT && user_email.is_some_and(|email| branch.committer_email() == email) {
152 mine.push(branch);
153 } else {
154 rest.push(branch);
155 }
156 }
157
158 let mut prioritized = Vec::with_capacity(branches_len);
159 prioritized.extend(previous);
160 prioritized.extend(mine);
161 prioritized.extend(rest);
162 prioritized
163}
164
165struct RenderableBranch(pub Branch);
166
167impl Deref for RenderableBranch {
168 type Target = Branch;
169
170 fn deref(&self) -> &Self::Target {
171 &self.0
172 }
173}
174
175impl Display for RenderableBranch {
176 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
177 let styled_date_time = format!("({})", self.committer_date_time());
178 let styled_email = format!("<{}>", self.committer_email());
179 write!(
180 f,
181 "{} {} {}",
182 self.name(),
183 styled_date_time.green(),
184 styled_email.blue().bold(),
185 )
186 }
187}
188
189#[cfg(test)]
190mod tests {
191 use chrono::DateTime;
192 use chrono::Utc;
193 use rstest::rstest;
194
195 use super::*;
196
197 #[rstest]
198 #[case(
199 vec![branch("main", 30), branch("feature-a", 20), branch("feature-b", 10)],
200 "feature-b",
201 vec![branch("feature-b", 10), branch("main", 30), branch("feature-a", 20)]
202 )]
203 #[case(
204 vec![remote_branch("origin/feature-a", 30), branch("main", 20)],
205 "feature-a",
206 vec![remote_branch("origin/feature-a", 30), branch("main", 20)]
207 )]
208 #[case(
209 vec![branch("main", 30), branch("feature-a", 20)],
210 "missing",
211 vec![branch("main", 30), branch("feature-a", 20)]
212 )]
213 fn test_prioritize_current_branch_first_when_current_branch_varies_orders_expected_branches(
214 #[case] branches: Vec<Branch>,
215 #[case] current_branch: &str,
216 #[case] expected: Vec<Branch>,
217 ) {
218 pretty_assertions::assert_eq!(
219 prioritize_current_branch_first(branches, current_branch, None, None),
220 expected
221 );
222 }
223
224 #[test]
225 fn test_prioritize_current_branch_first_preserves_gcu_recent_order_after_current() {
226 let branches = vec![
227 branch_with_email("other-1", "other@example.com", 100),
228 branch_with_email("mine-1", "me@example.com", 99),
229 branch_with_email("previous", "other@example.com", 98),
230 branch_with_email("current", "me@example.com", 97),
231 branch_with_email("mine-2", "me@example.com", 96),
232 ];
233
234 pretty_assertions::assert_eq!(
235 prioritize_current_branch_first(branches, "current", Some("previous"), Some("me@example.com")),
236 vec![
237 branch_with_email("current", "me@example.com", 97),
238 branch_with_email("previous", "other@example.com", 98),
239 branch_with_email("mine-1", "me@example.com", 99),
240 branch_with_email("mine-2", "me@example.com", 96),
241 branch_with_email("other-1", "other@example.com", 100),
242 ],
243 );
244 }
245
246 #[test]
247 fn test_install_zsh_wrapper_at_appends_guarded_line_and_is_idempotent() {
248 let dir = tempfile::tempdir().unwrap();
249 let zshrc = dir.path().join(".zshrc");
250 std::fs::write(&zshrc, "source ~/.zshrc.local\n").unwrap();
251
252 assert2::assert!(let Ok(true) = install_zsh_wrapper_at(&zshrc));
253 let first = std::fs::read_to_string(&zshrc).unwrap();
254 assert2::assert!(let Ok(false) = install_zsh_wrapper_at(&zshrc));
255 let second = std::fs::read_to_string(&zshrc).unwrap();
256
257 pretty_assertions::assert_eq!(first, second);
258 pretty_assertions::assert_eq!(first, format!("source ~/.zshrc.local\n{ZSHRC_INSTALL_LINE}\n"));
259 pretty_assertions::assert_eq!(first.matches(ZSHRC_INSTALL_LINE).count(), 1);
260 }
261
262 #[test]
263 fn test_install_zsh_wrapper_at_fails_when_zshrc_is_missing() {
264 let dir = tempfile::tempdir().unwrap();
265 let zshrc = dir.path().join(".zshrc");
266
267 assert2::assert!(let Err(err) = install_zsh_wrapper_at(&zshrc));
268
269 assert!(err.to_string().contains("error reading zshrc"));
270 }
271
272 fn branch(name: &str, timestamp: i64) -> Branch {
273 branch_with_email(name, "me@example.com", timestamp)
274 }
275
276 fn branch_with_email(name: &str, email: &str, timestamp: i64) -> Branch {
277 Branch::Local {
278 name: name.to_string(),
279 committer_email: email.to_string(),
280 committer_date_time: DateTime::<Utc>::from_timestamp(timestamp, 0).unwrap(),
281 }
282 }
283
284 fn remote_branch(name: &str, timestamp: i64) -> Branch {
285 Branch::Remote {
286 name: name.to_string(),
287 committer_email: "me@example.com".to_string(),
288 committer_date_time: DateTime::<Utc>::from_timestamp(timestamp, 0).unwrap(),
289 }
290 }
291}