1use std::collections::HashMap;
82
83use git2::{FetchOptions, Repository};
84use log::debug;
85
86use crate::{
87 error::{PrError, Result},
88 get_remote_callbacks,
89};
90
91#[derive(Debug, Clone, PartialEq, Eq)]
93pub struct PullRequest {
94 pub number: u32,
96 pub remote: Option<String>,
98}
99
100#[derive(Debug, Clone)]
102pub struct PrMetadata {
103 pub number: u32,
105 pub title: String,
107 pub author: String,
109 pub head_ref: String,
111 pub base_ref: String,
113 pub is_fork: bool,
115 pub fork_owner: Option<String>,
117 pub fork_url: Option<String>,
119}
120
121pub fn parse_pr_reference(input: &str) -> Result<Option<PullRequest>> {
133 if let Some(num_str) = input.strip_prefix('#') {
135 return parse_number(num_str, input).map(|num| {
136 Some(PullRequest {
137 number: num,
138 remote: None,
139 })
140 });
141 }
142
143 if let Some(num_str) = input.strip_prefix("pr#") {
145 return parse_number(num_str, input).map(|num| {
146 Some(PullRequest {
147 number: num,
148 remote: None,
149 })
150 });
151 }
152
153 if let Some(num_str) = input.strip_prefix("pr-") {
155 return parse_number(num_str, input).map(|num| {
156 Some(PullRequest {
157 number: num,
158 remote: None,
159 })
160 });
161 }
162
163 if input.contains("github.com") && input.contains("/pull/") {
165 return parse_github_url(input);
166 }
167
168 if input.contains("/pull/") && input.ends_with("/head") {
170 return parse_remote_ref(input);
171 }
172
173 Ok(None)
175}
176
177fn parse_number(num_str: &str, original_input: &str) -> Result<u32> {
179 num_str.parse::<u32>().map_err(|_| {
180 PrError::InvalidReference {
181 input: original_input.to_string(),
182 }
183 .into()
184 })
185}
186
187fn parse_github_url(url: &str) -> Result<Option<PullRequest>> {
189 let parts: Vec<&str> = url.split('/').collect();
191
192 for (i, &part) in parts.iter().enumerate() {
194 if part == "pull" && i + 1 < parts.len() {
195 let num_str = parts[i + 1];
196 let number = parse_number(num_str, url)?;
197 return Ok(Some(PullRequest {
198 number,
199 remote: None,
200 }));
201 }
202 }
203
204 Err(PrError::InvalidReference {
205 input: url.to_string(),
206 }
207 .into())
208}
209
210fn parse_remote_ref(ref_str: &str) -> Result<Option<PullRequest>> {
212 let parts: Vec<&str> = ref_str.split('/').collect();
214
215 if parts.len() >= 4 && parts[parts.len() - 3] == "pull" && parts[parts.len() - 1] == "head" {
216 let num_str = parts[parts.len() - 2];
217 let number = parse_number(num_str, ref_str)?;
218 return Ok(Some(PullRequest {
219 number,
220 remote: None,
221 }));
222 }
223
224 Err(PrError::InvalidReference {
225 input: ref_str.to_string(),
226 }
227 .into())
228}
229
230pub fn check_gh_available() -> Result<()> {
234 std::process::Command::new("gh")
235 .arg("--version")
236 .output()
237 .map_err(|_| PrError::GhNotInstalled)?;
238 Ok(())
239}
240
241pub fn fetch_pr_metadata(pr_number: u32) -> Result<PrMetadata> {
246 check_gh_available()?;
248
249 let output = std::process::Command::new("gh")
251 .args([
252 "pr",
253 "view",
254 &pr_number.to_string(),
255 "--json",
256 "number,title,author,headRefName,baseRefName,isCrossRepository,headRepository",
257 ])
258 .output()
259 .map_err(|e| PrError::GhFetchFailed {
260 message: e.to_string(),
261 })?;
262
263 if !output.status.success() {
264 let stderr = String::from_utf8_lossy(&output.stderr);
265 return Err(PrError::GhFetchFailed {
266 message: stderr.to_string(),
267 }
268 .into());
269 }
270
271 let json_str = String::from_utf8_lossy(&output.stdout);
273 let json: serde_json::Value =
274 serde_json::from_str(&json_str).map_err(|e| PrError::GhJsonParseFailed {
275 message: e.to_string(),
276 })?;
277
278 let number = json["number"]
280 .as_u64()
281 .ok_or_else(|| PrError::GhJsonParseFailed {
282 message: "Missing 'number' field".to_string(),
283 })? as u32;
284
285 let title = json["title"]
286 .as_str()
287 .ok_or_else(|| PrError::GhJsonParseFailed {
288 message: "Missing 'title' field".to_string(),
289 })?
290 .to_string();
291
292 let author = json["author"]["login"]
293 .as_str()
294 .ok_or_else(|| PrError::GhJsonParseFailed {
295 message: "Missing 'author.login' field".to_string(),
296 })?
297 .to_string();
298
299 let head_ref = json["headRefName"]
300 .as_str()
301 .ok_or_else(|| PrError::GhJsonParseFailed {
302 message: "Missing 'headRefName' field".to_string(),
303 })?
304 .to_string();
305
306 let base_ref = json["baseRefName"]
307 .as_str()
308 .ok_or_else(|| PrError::GhJsonParseFailed {
309 message: "Missing 'baseRefName' field".to_string(),
310 })?
311 .to_string();
312
313 let is_fork = json["isCrossRepository"].as_bool().unwrap_or(false);
314
315 let (fork_owner, fork_url) = if is_fork {
316 let owner = json["headRepository"]["owner"]["login"]
317 .as_str()
318 .ok_or(PrError::MissingForkOwner)?
319 .to_string();
320 let url = json["headRepository"]["url"]
321 .as_str()
322 .map(|s| s.to_string());
323 (Some(owner), url)
324 } else {
325 (None, None)
326 };
327
328 Ok(PrMetadata {
329 number,
330 title,
331 author,
332 head_ref,
333 base_ref,
334 is_fork,
335 fork_owner,
336 fork_url,
337 })
338}
339
340#[derive(Debug, Clone, PartialEq, Eq)]
342pub struct MergedPr {
343 pub number: u32,
344 pub head_oid: String,
345}
346
347const MERGED_PR_BATCH_SIZE: usize = 100;
349
350pub fn find_merged_prs(branches: &[&str]) -> Result<HashMap<String, MergedPr>> {
354 let mut found = HashMap::new();
355 for batch in branches.chunks(MERGED_PR_BATCH_SIZE) {
356 let mut cmd = std::process::Command::new("gh");
357 cmd.args(["api", "graphql", "-F", "owner={owner}", "-F", "name={repo}"]);
358 for (i, branch) in batch.iter().enumerate() {
359 cmd.arg("-f").arg(format!("h{i}={branch}"));
360 }
361 cmd.arg("-f")
362 .arg(format!("query={}", build_merged_pr_query(batch.len())));
363
364 let output = cmd.output().map_err(|e| PrError::GhFetchFailed {
365 message: e.to_string(),
366 })?;
367
368 if !output.status.success() {
369 let stderr = String::from_utf8_lossy(&output.stderr);
370 return Err(PrError::GhFetchFailed {
371 message: stderr.to_string(),
372 }
373 .into());
374 }
375
376 let json_str = String::from_utf8_lossy(&output.stdout);
377 found.extend(parse_merged_prs(&json_str, batch));
378 }
379 Ok(found)
380}
381
382fn build_merged_pr_query(count: usize) -> String {
387 let mut vars = String::from("$owner:String!,$name:String!");
388 let mut aliases = String::new();
389 for i in 0..count {
390 vars.push_str(&format!(",$h{i}:String!"));
391 aliases.push_str(&format!(
392 " b{i}: pullRequests(headRefName:$h{i}, states:MERGED, first:1, \
393 orderBy:{{field:CREATED_AT, direction:DESC}}){{ nodes {{ number headRefOid }} }}"
394 ));
395 }
396 format!("query({vars}){{ repository(owner:$owner,name:$name){{{aliases} }} }}")
397}
398
399fn parse_merged_prs(json: &str, branches: &[&str]) -> HashMap<String, MergedPr> {
405 let mut found = HashMap::new();
406 let Ok(json) = serde_json::from_str::<serde_json::Value>(json) else {
407 return found;
408 };
409 let Some(repository) = json.pointer("/data/repository").filter(|r| r.is_object()) else {
410 return found;
411 };
412 for (i, branch) in branches.iter().enumerate() {
413 let Some(node) = repository.pointer(&format!("/b{i}/nodes/0")) else {
414 continue;
415 };
416 let Some(number) = node.get("number").and_then(|n| n.as_u64()) else {
417 continue;
418 };
419 let Some(head_oid) = node.get("headRefOid").and_then(|o| o.as_str()) else {
420 continue;
421 };
422 found.insert(
423 (*branch).to_string(),
424 MergedPr {
425 number: number as u32,
426 head_oid: head_oid.to_string(),
427 },
428 );
429 }
430 found
431}
432
433fn sanitize_for_branch_name(s: &str) -> String {
435 let sanitized = s
436 .chars()
437 .map(|c| match c {
438 'a'..='z' | 'A'..='Z' | '0'..='9' | '-' | '_' => c,
439 ' ' | '/' => '-',
440 _ => '-',
441 })
442 .collect::<String>()
443 .to_lowercase();
444
445 let mut result = String::new();
447 let mut last_was_dash = false;
448 for c in sanitized.chars() {
449 if c == '-' {
450 if !last_was_dash {
451 result.push(c);
452 }
453 last_was_dash = true;
454 } else {
455 result.push(c);
456 last_was_dash = false;
457 }
458 }
459
460 result.trim_matches(|c| c == '-' || c == '_').to_string()
461}
462
463pub fn format_pr_name_with_metadata(format: &str, metadata: &PrMetadata) -> String {
468 format
469 .replace("{number}", &metadata.number.to_string())
470 .replace("{title}", &sanitize_for_branch_name(&metadata.title))
471 .replace("{author}", &sanitize_for_branch_name(&metadata.author))
472 .replace("{branch}", &sanitize_for_branch_name(&metadata.head_ref))
473}
474
475pub fn is_pr_reference(input: &str) -> bool {
479 parse_pr_reference(input).ok().flatten().is_some()
480}
481
482pub fn remote_priority(remote: &str) -> usize {
489 match remote {
490 "upstream" => 0,
491 "origin" => 1,
492 _ => 2,
493 }
494}
495
496pub fn preferred_remote_order(repo: &Repository) -> Vec<String> {
499 let Ok(remotes) = repo.remotes() else {
500 return vec![];
501 };
502 let mut all: Vec<String> = remotes
503 .iter()
504 .flatten()
505 .flatten()
506 .map(str::to_string)
507 .collect();
508 all.sort_by_key(|r| remote_priority(r));
509 all
510}
511
512pub fn has_github_remote(repo: &Repository) -> bool {
519 let Ok(remotes) = repo.remotes() else {
520 return false;
521 };
522 remotes.iter().flatten().flatten().any(|name| {
523 repo.find_remote(name)
524 .ok()
525 .and_then(|r| r.url().ok().map(str::to_string))
526 .and_then(|url| crate::ssh_config::extract_host_from_url(&url))
527 .is_some_and(|host| host.eq_ignore_ascii_case("github.com"))
528 })
529}
530
531pub fn detect_pr_remote(repo: &Repository) -> Result<String> {
536 preferred_remote_order(repo)
537 .into_iter()
538 .next()
539 .ok_or_else(|| PrError::NoRemoteConfigured.into())
540}
541
542pub fn setup_fork_remote(repo: &Repository, metadata: &PrMetadata) -> Result<String> {
548 if !metadata.is_fork {
549 return detect_pr_remote(repo);
551 }
552
553 let _fork_owner = metadata
555 .fork_owner
556 .as_ref()
557 .ok_or(PrError::MissingForkOwner)?;
558
559 let fork_url = metadata
560 .fork_url
561 .as_ref()
562 .ok_or(PrError::MissingForkOwner)?;
563
564 let fork_remote_name = format!("pr-{}-fork", metadata.number);
566
567 if repo.find_remote(&fork_remote_name).is_ok() {
568 debug!("Fork remote {} already exists", fork_remote_name);
569 return Ok(fork_remote_name);
570 }
571
572 debug!("Adding fork remote: {} -> {}", fork_remote_name, fork_url);
574 repo.remote(&fork_remote_name, fork_url)
575 .map_err(|e| PrError::FetchFailed {
576 remote: fork_remote_name.clone(),
577 message: format!("Failed to add fork remote: {}", e),
578 })?;
579
580 Ok(fork_remote_name)
581}
582
583pub fn fetch_branch(repo: &Repository, remote_name: &str, branch: &str) -> Result<()> {
590 let branch_ref = format!("refs/remotes/{}/{}", remote_name, branch);
592 if repo.find_reference(&branch_ref).is_ok() {
593 debug!("Branch ref {} already exists", branch_ref);
594 return Ok(());
595 }
596
597 debug!("Fetching branch {} from remote {}", branch, remote_name);
598
599 let refspec = format!(
600 "+refs/heads/{}:refs/remotes/{}/{}",
601 branch, remote_name, branch
602 );
603
604 let remote_url = repo
605 .find_remote(remote_name)
606 .ok()
607 .and_then(|r| r.url().ok().map(str::to_string));
608 let auth = get_remote_callbacks(repo, remote_url.as_deref())?;
609 let mut fetch_options = FetchOptions::new();
610 fetch_options.remote_callbacks(auth.callbacks());
611
612 repo.find_remote(remote_name)?
613 .fetch(
614 &[refspec.as_str()],
615 Some(&mut fetch_options),
616 Some("Fetching PR branch"),
617 )
618 .map_err(|e| PrError::FetchFailed {
619 remote: remote_name.to_string(),
620 message: e.message().to_string(),
621 })?;
622
623 debug!("Successfully fetched branch {}", branch);
624 Ok(())
625}
626
627pub fn format_pr_name(format: &str, pr_number: u32) -> String {
631 format.replace("{number}", &pr_number.to_string())
632}
633
634pub fn prepare_pr_worktree(
645 repo: &Repository,
646 pr_number: u32,
647 pr_format: &str,
648) -> Result<(String, String, String)> {
649 debug!("Preparing PR worktree for PR #{}", pr_number);
650
651 let metadata = fetch_pr_metadata(pr_number)?;
653 debug!(
654 "Fetched metadata: title='{}', author='{}', is_fork={}",
655 metadata.title, metadata.author, metadata.is_fork
656 );
657
658 let remote_name = if metadata.is_fork {
662 setup_fork_remote(repo, &metadata)?
663 } else {
664 detect_pr_remote(repo)?
665 };
666
667 fetch_branch(repo, &remote_name, &metadata.head_ref)?;
669
670 let worktree_name = format_pr_name_with_metadata(pr_format, &metadata);
672 debug!("Worktree name: {}", worktree_name);
673
674 let remote_ref = format!("{}/{}", remote_name, metadata.head_ref);
676 debug!("Remote ref: {}", remote_ref);
677
678 Ok((worktree_name, remote_ref, metadata.base_ref))
679}
680
681#[cfg(test)]
682mod tests {
683 use super::*;
684
685 #[test]
686 fn test_parse_hash_number() {
687 let pr = parse_pr_reference("#123").unwrap().unwrap();
688 assert_eq!(pr.number, 123);
689 assert_eq!(pr.remote, None);
690 }
691
692 #[test]
693 fn test_parse_pr_hash_number() {
694 let pr = parse_pr_reference("pr#456").unwrap().unwrap();
695 assert_eq!(pr.number, 456);
696 assert_eq!(pr.remote, None);
697 }
698
699 #[test]
700 fn test_parse_pr_dash_number() {
701 let pr = parse_pr_reference("pr-789").unwrap().unwrap();
702 assert_eq!(pr.number, 789);
703 assert_eq!(pr.remote, None);
704 }
705
706 #[test]
707 fn test_parse_github_url() {
708 let pr = parse_pr_reference("https://github.com/owner/repo/pull/999")
709 .unwrap()
710 .unwrap();
711 assert_eq!(pr.number, 999);
712 assert_eq!(pr.remote, None);
713 }
714
715 #[test]
716 fn test_parse_remote_ref() {
717 let pr = parse_pr_reference("origin/pull/111/head").unwrap().unwrap();
718 assert_eq!(pr.number, 111);
719 assert_eq!(pr.remote, None);
720 }
721
722 #[test]
723 fn test_parse_regular_branch_name() {
724 let result = parse_pr_reference("my-feature-branch").unwrap();
725 assert!(result.is_none());
726 }
727
728 #[test]
729 fn test_parse_invalid_number() {
730 let result = parse_pr_reference("#abc");
731 assert!(result.is_err());
732 }
733
734 #[test]
735 fn test_is_pr_reference_true() {
736 assert!(is_pr_reference("#123"));
737 assert!(is_pr_reference("pr#456"));
738 assert!(is_pr_reference("pr-789"));
739 assert!(is_pr_reference("https://github.com/owner/repo/pull/999"));
740 }
741
742 #[test]
743 fn test_is_pr_reference_false() {
744 assert!(!is_pr_reference("my-branch"));
745 assert!(!is_pr_reference("feature"));
746 }
747
748 #[test]
749 fn test_format_pr_name() {
750 assert_eq!(format_pr_name("pr-{number}", 123), "pr-123");
751 assert_eq!(format_pr_name("review-{number}", 456), "review-456");
752 assert_eq!(format_pr_name("{number}-test", 789), "789-test");
753 }
754
755 #[test]
756 fn test_sanitize_branch_name() {
757 assert_eq!(sanitize_for_branch_name("Fix Bug #123"), "fix-bug-123");
758 assert_eq!(
759 sanitize_for_branch_name("Add Feature (v2)"),
760 "add-feature-v2"
761 );
762 assert_eq!(sanitize_for_branch_name("john-smith"), "john-smith");
763 assert_eq!(
764 sanitize_for_branch_name("Fix: Authentication Issue"),
765 "fix-authentication-issue"
766 );
767 assert_eq!(sanitize_for_branch_name("Test@#$%"), "test");
768 }
769
770 #[test]
771 fn test_format_with_metadata() {
772 let metadata = PrMetadata {
773 number: 123,
774 title: "Fix Authentication Bug".to_string(),
775 author: "john-smith".to_string(),
776 head_ref: "feature/fix-auth".to_string(),
777 base_ref: "main".to_string(),
778 is_fork: false,
779 fork_owner: None,
780 fork_url: None,
781 };
782
783 assert_eq!(
784 format_pr_name_with_metadata("pr-{number}", &metadata),
785 "pr-123"
786 );
787 assert_eq!(
788 format_pr_name_with_metadata("{number}-{title}", &metadata),
789 "123-fix-authentication-bug"
790 );
791 assert_eq!(
792 format_pr_name_with_metadata("{author}/pr-{number}", &metadata),
793 "john-smith/pr-123"
794 );
795 assert_eq!(
796 format_pr_name_with_metadata("{branch}-{number}", &metadata),
797 "feature-fix-auth-123"
798 );
799 }
800
801 fn graphql(members: &str) -> String {
802 format!(r#"{{"data":{{"repository":{{{members}}}}}}}"#)
803 }
804
805 #[test]
806 fn test_parse_merged_prs_empty_nodes() {
807 let json = graphql(r#""b0":{"nodes":[]}"#);
808 assert!(parse_merged_prs(&json, &["a"]).is_empty());
809 }
810
811 #[test]
812 fn test_parse_merged_prs_populated() {
813 let json = graphql(r#""b0":{"nodes":[{"number":66,"headRefOid":"abc123"}]}"#);
814 let map = parse_merged_prs(&json, &["a"]);
815 assert_eq!(
816 map.get("a"),
817 Some(&MergedPr {
818 number: 66,
819 head_oid: "abc123".to_string()
820 })
821 );
822 }
823
824 #[test]
825 fn test_parse_merged_prs_two_aliases_one_empty() {
826 let json = graphql(r#""b0":{"nodes":[]},"b1":{"nodes":[{"number":7,"headRefOid":"def"}]}"#);
827 let map = parse_merged_prs(&json, &["a", "b"]);
828 assert_eq!(map.len(), 1);
829 assert_eq!(map["b"].number, 7);
830 }
831
832 #[test]
833 fn test_parse_merged_prs_number_missing() {
834 let json = graphql(r#""b0":{"nodes":[{"headRefOid":"abc123"}]}"#);
835 assert!(parse_merged_prs(&json, &["a"]).is_empty());
836 }
837
838 #[test]
839 fn test_parse_merged_prs_number_non_numeric() {
840 let json = graphql(r#""b0":{"nodes":[{"number":"x","headRefOid":"abc123"}]}"#);
841 assert!(parse_merged_prs(&json, &["a"]).is_empty());
842 }
843
844 #[test]
845 fn test_parse_merged_prs_head_oid_missing() {
846 let json = graphql(r#""b0":{"nodes":[{"number":66}]}"#);
847 assert!(parse_merged_prs(&json, &["a"]).is_empty());
848 }
849
850 #[test]
851 fn test_parse_merged_prs_missing_alias() {
852 let json = graphql(r#""b1":{"nodes":[{"number":66,"headRefOid":"abc"}]}"#);
853 let map = parse_merged_prs(&json, &["a", "b"]);
854 assert_eq!(map.len(), 1);
855 assert!(map.contains_key("b"));
856 }
857
858 #[test]
859 fn test_parse_merged_prs_null_repository() {
860 let json = r#"{"data":{"repository":null}}"#;
861 assert!(parse_merged_prs(json, &["a"]).is_empty());
862 }
863
864 #[test]
865 fn test_parse_merged_prs_malformed_json() {
866 assert!(parse_merged_prs("not json", &["a"]).is_empty());
867 }
868
869 #[test]
870 fn test_merged_pr_query_declares_variable_and_alias_per_branch() {
871 let query = build_merged_pr_query(3);
872 assert!(query.starts_with(
873 "query($owner:String!,$name:String!,$h0:String!,$h1:String!,$h2:String!){"
874 ));
875 for i in 0..3 {
876 assert!(query.contains(&format!("b{i}: pullRequests(headRefName:$h{i}")));
877 }
878 assert!(!query.contains("$h3"));
879 assert!(query.contains("direction:DESC"));
880 }
881
882 #[test]
884 #[ignore]
885 fn test_gh_cli_available() {
886 check_gh_available().expect("gh CLI should be installed");
887 }
888
889 #[test]
890 #[ignore]
891 fn test_fetch_real_pr_metadata() {
892 let metadata = fetch_pr_metadata(1).expect("Failed to fetch PR metadata");
896 assert_eq!(metadata.number, 1);
897 assert!(!metadata.title.is_empty());
898 assert!(!metadata.author.is_empty());
899 }
900}