1use std::collections::HashMap;
2use std::path::{Path, PathBuf};
3
4use serde::{Deserialize, Serialize};
5
6use crate::core::bm25_index::{BM25Index, SearchResult};
7
8const DEFAULT_RRF_K: f64 = 60.0;
10
11const MAX_ROOTS: usize = 16;
13
14#[derive(Debug, Clone)]
16pub struct RepoSearchResult {
17 pub repo_alias: String,
18 pub repo_path: String,
19 pub file_path: String,
20 pub symbol_name: String,
21 pub content: String,
22 pub start_line: usize,
23 pub end_line: usize,
24 pub score: f64,
25}
26
27#[derive(Debug, Clone)]
29pub struct FusedSearchResult {
30 pub repo_alias: String,
31 pub repo_path: String,
32 pub file_path: String,
33 pub symbol_name: String,
34 pub content: String,
35 pub start_line: usize,
36 pub end_line: usize,
37 pub rrf_score: f64,
38}
39
40#[derive(Debug, Clone, Serialize, Deserialize)]
42pub struct RepoRootConfig {
43 pub path: String,
44 #[serde(default)]
45 pub alias: Option<String>,
46}
47
48impl RepoRootConfig {
49 pub fn effective_alias(&self) -> String {
50 self.alias.clone().unwrap_or_else(|| {
51 Path::new(&self.path)
52 .file_name()
53 .and_then(|n| n.to_str())
54 .unwrap_or("unknown")
55 .to_string()
56 })
57 }
58}
59
60#[derive(Debug, Clone, Serialize, Deserialize, Default)]
62pub struct MultiRepoConfig {
63 #[serde(default)]
64 pub repos: Vec<RepoRootConfig>,
65 #[serde(default)]
66 pub rrf_k: Option<f64>,
67}
68
69impl MultiRepoConfig {
70 pub fn load() -> Self {
71 let config_path = config_file_path();
72 if !config_path.exists() {
73 return Self::default();
74 }
75 match std::fs::read_to_string(&config_path) {
76 Ok(content) => toml::from_str(&content).unwrap_or_default(),
77 Err(_) => Self::default(),
78 }
79 }
80
81 pub fn save(&self) -> Result<(), String> {
82 let config_path = config_file_path();
83 if let Some(parent) = config_path.parent() {
84 std::fs::create_dir_all(parent)
85 .map_err(|e| format!("Failed to create config dir: {e}"))?;
86 }
87 let content =
88 toml::to_string_pretty(self).map_err(|e| format!("Failed to serialize config: {e}"))?;
89 let defaults = toml::to_string_pretty(&Self::default())
90 .map_err(|e| format!("Failed to serialize defaults: {e}"))?;
91 crate::config_io::write_toml_preserving_minimal(&config_path, &content, &defaults)
92 .map_err(|e| format!("Failed to write config: {e}"))?;
93 Ok(())
94 }
95}
96
97pub struct ActiveRepoRoot {
99 pub config: RepoRootConfig,
100 pub path: PathBuf,
101 index: Option<BM25Index>,
102}
103
104impl std::fmt::Debug for ActiveRepoRoot {
105 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
106 f.debug_struct("ActiveRepoRoot")
107 .field("config", &self.config)
108 .field("path", &self.path)
109 .field("has_index", &self.index.is_some())
110 .finish()
111 }
112}
113
114impl ActiveRepoRoot {
115 fn new(config: RepoRootConfig) -> Result<Self, String> {
116 let path = PathBuf::from(&config.path);
117 if !path.is_dir() {
118 return Err(format!(
119 "Path does not exist or is not a directory: {}",
120 config.path
121 ));
122 }
123 let path = path
124 .canonicalize()
125 .map_err(|e| format!("Cannot canonicalize {}: {e}", config.path))?;
126 Ok(Self {
127 config,
128 path,
129 index: None,
130 })
131 }
132
133 fn ensure_index(&mut self) {
134 if self.index.is_some() {
135 return;
136 }
137 self.index = Some(BM25Index::load_or_build(&self.path));
138 }
139
140 pub fn alias(&self) -> String {
141 self.config.effective_alias()
142 }
143
144 pub fn search(&mut self, query: &str, max_results: usize) -> Vec<RepoSearchResult> {
145 self.ensure_index();
146 let Some(ref index) = self.index else {
147 return Vec::new();
148 };
149
150 let results: Vec<SearchResult> = index.search(query, max_results);
151 let alias = self.alias();
152 let repo_path = self.path.to_string_lossy().to_string();
153
154 results
155 .into_iter()
156 .enumerate()
157 .map(|(rank, sr)| RepoSearchResult {
158 repo_alias: alias.clone(),
159 repo_path: repo_path.clone(),
160 file_path: sr.file_path,
161 symbol_name: sr.symbol_name,
162 content: sr.snippet,
163 start_line: sr.start_line,
164 end_line: sr.end_line,
165 score: 1.0 / (rank as f64 + 1.0),
166 })
167 .collect()
168 }
169}
170
171pub struct MultiRepoManager {
173 roots: Vec<ActiveRepoRoot>,
174 rrf_k: f64,
175}
176
177impl MultiRepoManager {
178 pub fn new() -> Self {
179 Self {
180 roots: Vec::new(),
181 rrf_k: DEFAULT_RRF_K,
182 }
183 }
184
185 pub fn with_rrf_k(mut self, k: f64) -> Self {
186 self.rrf_k = k;
187 self
188 }
189
190 pub fn from_config(config: &MultiRepoConfig) -> Result<Self, String> {
191 let mut manager = Self::new();
192 if let Some(k) = config.rrf_k {
193 manager.rrf_k = k;
194 }
195 for repo_config in &config.repos {
196 manager.add_root_config(repo_config.clone())?;
197 }
198 Ok(manager)
199 }
200
201 pub fn add_root(&mut self, path: &str, alias: Option<&str>) -> Result<(), String> {
202 if self.roots.len() >= MAX_ROOTS {
203 return Err(format!("Maximum number of roots ({MAX_ROOTS}) reached"));
204 }
205 let config = RepoRootConfig {
206 path: path.to_string(),
207 alias: alias.map(String::from),
208 };
209 let root = ActiveRepoRoot::new(config)?;
210 if self.roots.iter().any(|r| r.path == root.path) {
211 return Err(format!("Root already exists: {path}"));
212 }
213 self.roots.push(root);
214 Ok(())
215 }
216
217 fn add_root_config(&mut self, config: RepoRootConfig) -> Result<(), String> {
218 if self.roots.len() >= MAX_ROOTS {
219 return Err(format!("Maximum number of roots ({MAX_ROOTS}) reached"));
220 }
221 let root = ActiveRepoRoot::new(config)?;
222 if self.roots.iter().any(|r| r.path == root.path) {
223 return Err(format!(
224 "Root already exists: {}",
225 root.path.to_string_lossy()
226 ));
227 }
228 self.roots.push(root);
229 Ok(())
230 }
231
232 pub fn remove_root(&mut self, path: &str) -> Result<(), String> {
233 let normalized = PathBuf::from(path)
234 .canonicalize()
235 .unwrap_or_else(|_| PathBuf::from(path));
236 let before = self.roots.len();
237 self.roots
238 .retain(|r| r.path != normalized && r.config.path != path);
239 if self.roots.len() == before {
240 return Err(format!("Root not found: {path}"));
241 }
242 Ok(())
243 }
244
245 pub fn list_roots(&self) -> Vec<RootInfo> {
246 self.roots
247 .iter()
248 .map(|r| RootInfo {
249 path: r.path.to_string_lossy().to_string(),
250 alias: r.alias(),
251 has_index: r.index.is_some(),
252 })
253 .collect()
254 }
255
256 pub fn root_count(&self) -> usize {
257 self.roots.len()
258 }
259
260 pub fn is_active(&self) -> bool {
261 self.roots.len() > 1
262 }
263
264 pub fn resolve_root(&self, repo: &str) -> Option<usize> {
266 self.roots.iter().position(|r| {
267 r.alias() == repo || r.config.path == repo || r.path.to_string_lossy() == repo
268 })
269 }
270
271 pub fn search(
273 &mut self,
274 query: &str,
275 max_results: usize,
276 filter_roots: Option<&[String]>,
277 ) -> Vec<FusedSearchResult> {
278 let per_root_max = (max_results * 2).max(20);
279
280 let mut all_results: HashMap<String, FusedSearchResult> = HashMap::new();
281
282 for root in &mut self.roots {
283 if let Some(filter) = filter_roots {
284 let alias = root.alias();
285 let path = root.path.to_string_lossy().to_string();
286 if !filter.iter().any(|f| f == &alias || f == &path) {
287 continue;
288 }
289 }
290
291 let results = root.search(query, per_root_max);
292
293 for (rank, result) in results.iter().enumerate() {
294 let rrf_contribution = 1.0 / (self.rrf_k + rank as f64 + 1.0);
295 let key = format!(
296 "{}:{}:{}",
297 result.repo_alias, result.file_path, result.start_line
298 );
299
300 all_results
301 .entry(key)
302 .and_modify(|existing| {
303 existing.rrf_score += rrf_contribution;
304 })
305 .or_insert_with(|| FusedSearchResult {
306 repo_alias: result.repo_alias.clone(),
307 repo_path: result.repo_path.clone(),
308 file_path: result.file_path.clone(),
309 symbol_name: result.symbol_name.clone(),
310 content: result.content.clone(),
311 start_line: result.start_line,
312 end_line: result.end_line,
313 rrf_score: rrf_contribution,
314 });
315 }
316 }
317
318 let mut fused: Vec<FusedSearchResult> = all_results.into_values().collect();
319 fused.sort_by(|a, b| {
320 b.rrf_score
321 .partial_cmp(&a.rrf_score)
322 .unwrap_or(std::cmp::Ordering::Equal)
323 });
324 fused.truncate(max_results);
325 fused
326 }
327
328 pub fn search_single_repo(
330 &mut self,
331 repo: &str,
332 query: &str,
333 max_results: usize,
334 ) -> Result<Vec<RepoSearchResult>, String> {
335 let idx = self
336 .resolve_root(repo)
337 .ok_or_else(|| format!("Unknown repo: {repo}"))?;
338 Ok(self.roots[idx].search(query, max_results))
339 }
340}
341
342impl Default for MultiRepoManager {
343 fn default() -> Self {
344 Self::new()
345 }
346}
347
348#[derive(Debug, Clone, Serialize)]
350pub struct RootInfo {
351 pub path: String,
352 pub alias: String,
353 pub has_index: bool,
354}
355
356pub fn config_file_path() -> PathBuf {
362 crate::core::paths::config_dir_member("multi-repo.toml")
363 .unwrap_or_else(|_| PathBuf::from("~/.config/lean-ctx/multi-repo.toml"))
364}
365
366static GLOBAL_MANAGER: std::sync::OnceLock<std::sync::Mutex<MultiRepoManager>> =
368 std::sync::OnceLock::new();
369
370pub fn global_manager() -> &'static std::sync::Mutex<MultiRepoManager> {
371 GLOBAL_MANAGER.get_or_init(|| {
372 let config = MultiRepoConfig::load();
373 let manager = MultiRepoManager::from_config(&config).unwrap_or_default();
374 std::sync::Mutex::new(manager)
375 })
376}
377
378pub fn init_with_roots(
380 roots: &[(String, Option<String>)],
381 rrf_k: Option<f64>,
382) -> Result<(), String> {
383 let mut manager = MultiRepoManager::new();
384 if let Some(k) = rrf_k {
385 manager.rrf_k = k;
386 }
387 for (path, alias) in roots {
388 manager.add_root(path, alias.as_deref())?;
389 }
390 GLOBAL_MANAGER
391 .set(std::sync::Mutex::new(manager))
392 .map_err(|_| "Multi-repo manager already initialized".to_string())
393}
394
395pub fn resolve_repo_root(repo: &str) -> Option<String> {
399 let manager = global_manager();
400 let mgr = manager.lock().ok()?;
401 let idx = mgr.resolve_root(repo)?;
402 Some(mgr.roots[idx].path.to_string_lossy().to_string())
403}
404
405pub fn is_multi_repo_active() -> bool {
407 let manager = global_manager();
408 manager.lock().is_ok_and(|mgr| mgr.is_active())
409}
410
411pub fn all_root_paths() -> Vec<String> {
413 let manager = global_manager();
414 let Ok(mgr) = manager.lock() else {
415 return Vec::new();
416 };
417 mgr.roots
418 .iter()
419 .map(|r| r.path.to_string_lossy().to_string())
420 .collect()
421}
422
423pub fn format_fused_results(results: &[FusedSearchResult]) -> String {
425 if results.is_empty() {
426 return "No results found across repos.".to_string();
427 }
428
429 let mut out = String::with_capacity(results.len() * 200);
430 out.push_str(&format!(
431 "Cross-repo results ({} matches):\n\n",
432 results.len()
433 ));
434
435 for (i, result) in results.iter().enumerate() {
436 out.push_str(&format!(
437 "{}. [{}] {}:{}-{} ({})\n RRF: {:.4}\n",
438 i + 1,
439 result.repo_alias,
440 result.file_path,
441 result.start_line,
442 result.end_line,
443 result.symbol_name,
444 result.rrf_score,
445 ));
446 let preview: String = result
447 .content
448 .lines()
449 .take(3)
450 .collect::<Vec<_>>()
451 .join("\n");
452 if !preview.is_empty() {
453 out.push_str(&format!(" {}\n", preview.replace('\n', "\n ")));
454 }
455 out.push('\n');
456 }
457 out
458}
459
460#[cfg(test)]
461mod tests {
462 use super::*;
463
464 #[test]
465 fn repo_root_config_effective_alias() {
466 let cfg = RepoRootConfig {
467 path: "/home/user/projects/backend".to_string(),
468 alias: None,
469 };
470 assert_eq!(cfg.effective_alias(), "backend");
471
472 let cfg_with_alias = RepoRootConfig {
473 path: "/home/user/projects/backend".to_string(),
474 alias: Some("api".to_string()),
475 };
476 assert_eq!(cfg_with_alias.effective_alias(), "api");
477 }
478
479 #[test]
480 fn multi_repo_config_default_is_empty() {
481 let cfg = MultiRepoConfig::default();
482 assert!(cfg.repos.is_empty());
483 assert!(cfg.rrf_k.is_none());
484 }
485
486 #[test]
487 fn multi_repo_config_deserialize() {
488 let toml_str = r#"
489rrf_k = 45.0
490
491[[repos]]
492path = "/home/user/backend"
493alias = "backend"
494
495[[repos]]
496path = "/home/user/frontend"
497"#;
498 let cfg: MultiRepoConfig = toml::from_str(toml_str).unwrap();
499 assert_eq!(cfg.repos.len(), 2);
500 assert_eq!(cfg.rrf_k, Some(45.0));
501 assert_eq!(cfg.repos[0].alias, Some("backend".to_string()));
502 assert_eq!(cfg.repos[1].alias, None);
503 }
504
505 #[test]
506 fn manager_max_roots_enforced() {
507 let mut manager = MultiRepoManager::new();
508 for i in 0..MAX_ROOTS {
509 let dir = std::env::temp_dir().join(format!("multi_repo_test_{i}"));
510 let _ = std::fs::create_dir_all(&dir);
511 let _ = manager.add_root(&dir.to_string_lossy(), Some(&format!("repo{i}")));
512 }
513 let extra = std::env::temp_dir().join("multi_repo_test_extra");
514 let _ = std::fs::create_dir_all(&extra);
515 let result = manager.add_root(&extra.to_string_lossy(), None);
516 assert!(result.is_err());
517
518 for i in 0..=MAX_ROOTS {
519 let dir = std::env::temp_dir().join(format!("multi_repo_test_{i}"));
520 let _ = std::fs::remove_dir_all(&dir);
521 }
522 let _ = std::fs::remove_dir_all(&extra);
523 }
524
525 #[test]
526 fn manager_duplicate_root_rejected() {
527 let dir = std::env::temp_dir().join("multi_repo_dup_test");
528 let _ = std::fs::create_dir_all(&dir);
529 let mut manager = MultiRepoManager::new();
530 assert!(
531 manager
532 .add_root(&dir.to_string_lossy(), Some("first"))
533 .is_ok()
534 );
535 assert!(
536 manager
537 .add_root(&dir.to_string_lossy(), Some("second"))
538 .is_err()
539 );
540 let _ = std::fs::remove_dir_all(&dir);
541 }
542
543 #[test]
544 fn rrf_fusion_basic() {
545 let manager = MultiRepoManager::new().with_rrf_k(60.0);
546 let score: f64 = 1.0 / (60.0 + 0.0 + 1.0);
548 assert!((score - 0.01639).abs() < 0.001);
549
550 assert_eq!(manager.rrf_k, 60.0);
551 }
552
553 #[test]
554 fn remove_root_works() {
555 let dir = std::env::temp_dir().join("multi_repo_remove_test");
556 let _ = std::fs::create_dir_all(&dir);
557 let mut manager = MultiRepoManager::new();
558 manager
559 .add_root(&dir.to_string_lossy(), Some("removable"))
560 .unwrap();
561 assert_eq!(manager.root_count(), 1);
562 manager.remove_root(&dir.to_string_lossy()).unwrap();
563 assert_eq!(manager.root_count(), 0);
564 let _ = std::fs::remove_dir_all(&dir);
565 }
566
567 #[test]
568 fn list_roots_returns_info() {
569 let dir = std::env::temp_dir().join("multi_repo_list_test");
570 let _ = std::fs::create_dir_all(&dir);
571 let mut manager = MultiRepoManager::new();
572 manager
573 .add_root(&dir.to_string_lossy(), Some("myrepo"))
574 .unwrap();
575 let roots = manager.list_roots();
576 assert_eq!(roots.len(), 1);
577 assert_eq!(roots[0].alias, "myrepo");
578 let _ = std::fs::remove_dir_all(&dir);
579 }
580
581 #[test]
582 fn format_empty_results() {
583 let results: Vec<FusedSearchResult> = Vec::new();
584 let output = format_fused_results(&results);
585 assert!(output.contains("No results"));
586 }
587}