1use crate::cli::{Cli, SetupScope};
2use crate::pkg::{install, local, plugin, types};
3use crate::utils;
4use anyhow::{Result, anyhow};
5use clap::CommandFactory;
6use clap_complete::{Shell, generate};
7use colored::*;
8use std::collections::{HashMap, HashSet};
9use std::fs;
10use std::io::Write;
11use std::path::{Path, PathBuf};
12use std::process::Command;
13use std::sync::Mutex;
14
15fn get_completion_path(shell: Shell, scope: SetupScope) -> Result<PathBuf> {
16 if scope == SetupScope::System {
17 if !utils::is_admin() {
18 return Err(anyhow!(
19 "System-wide installation requires root privileges. Please run with sudo or as an administrator.",
20 ));
21 }
22 Ok(match shell {
23 Shell::Bash => PathBuf::from("/usr/share/bash-completion/completions/zoi"),
24 Shell::Elvish => PathBuf::from("/usr/share/elvish/lib/zoi.elv"),
25 Shell::Fish => PathBuf::from("/usr/share/fish/vendor_completions.d/zoi.fish"),
26 Shell::Zsh => PathBuf::from("/usr/share/zsh/site-functions/_zoi"),
27 _ => {
28 return Err(anyhow!(
29 "System-wide completion installation not supported for this shell.",
30 ));
31 }
32 })
33 } else {
34 let home = dirs::home_dir().ok_or_else(|| anyhow!("Home directory not found"))?;
35 Ok(match shell {
36 Shell::Bash => home.join(".local/share/bash-completion/completions/zoi"),
37 Shell::Zsh => home.join(".zsh/completions/_zoi"),
38 Shell::Fish => home.join(".config/fish/completions/zoi.fish"),
39 Shell::Elvish => home.join(".config/elvish/completions/zoi.elv"),
40 Shell::PowerShell => {
41 if cfg!(windows) {
42 home.join("Documents/PowerShell/Microsoft.PowerShell_profile.ps1")
43 } else {
44 home.join(".config/powershell/Microsoft.PowerShell_profile.ps1")
45 }
46 }
47 _ => {
48 return Err(anyhow!(
49 "User-level completion installation not supported for this shell.",
50 ));
51 }
52 })
53 }
54}
55
56fn install_completions(shell: Shell, scope: SetupScope, cmd: &mut clap::Command) -> Result<()> {
57 if cfg!(windows) && scope == SetupScope::System {
58 return Err(anyhow!(
59 "System-wide shell setup is not supported on Windows.",
60 ));
61 }
62
63 let path = get_completion_path(shell, scope)?;
64 if let Some(parent) = path.parent() {
65 fs::create_dir_all(parent)?;
66 }
67
68 if shell == Shell::PowerShell {
69 let mut file = fs::OpenOptions::new()
70 .append(true)
71 .create(true)
72 .open(&path)?;
73 writeln!(file)?;
74 let mut script_buf = Vec::new();
75 generate(shell, cmd, "zoi", &mut script_buf);
76 let script =
77 post_process_completions(shell, String::from_utf8_lossy(&script_buf).to_string());
78 file.write_all(script.as_bytes())?;
79 println!(
80 "PowerShell completion script appended to your profile: {:?}",
81 path
82 );
83 println!("Please restart your shell or run '. $PROFILE' to activate it.");
84 } else {
85 let mut script_buf = Vec::new();
86 generate(shell, cmd, "zoi", &mut script_buf);
87 let script =
88 post_process_completions(shell, String::from_utf8_lossy(&script_buf).to_string());
89 let mut file = fs::File::create(&path)?;
90 file.write_all(script.as_bytes())?;
91 println!("{} completions installed in: {:?}", shell, path);
92 }
93
94 if shell == Shell::Zsh && scope == SetupScope::User {
95 println!("Ensure the directory is in your $fpath. Add this to your .zshrc if it's not:");
96 println!(
97 " fpath=({:?} $fpath)",
98 path.parent()
99 .ok_or_else(|| anyhow!("Path should have a parent directory: {:?}", path))?
100 );
101 }
102
103 Ok(())
104}
105
106fn post_process_completions(shell: Shell, mut script: String) -> String {
107 match shell {
108 Shell::Zsh => {
109 let helper = r#"
110_zoi_packages() {
111 local -a entries
112 local line
113 while IFS= read -r line; do
114 [[ -z "$line" ]] && continue
115 entries+=("$line")
116 done < <(zoi complete zsh $CURRENT "${words[@]}" 2>/dev/null)
117 _describe -t packages 'packages' entries
118}
119
120_zoi_all_packages() {
121 _zoi_packages
122}
123
124_zoi_installed_packages() {
125 _zoi_packages
126}
127"#;
128 let mut parts = script.splitn(2, '\n');
129 let header = parts.next().unwrap_or("");
130 let body = parts.next().unwrap_or("");
131
132 script = format!("{}\n{}\n{}", header, helper, body);
133
134 script = script.replace("':ALL_SOURCES: '", "':package:(_zoi_packages)'");
135 script = script.replace("':ALL_PACKAGES: '", "':package:(_zoi_packages)'");
136 script = script.replace("':INST_PACKAGES: '", "':package:(_zoi_packages)'");
137
138 let desc_marker = " -- Package identifier (e.g. @repo/name, path, or URL):";
139 let mut search_start = 0;
140 while let Some(pos) = script[search_start..].find(desc_marker) {
141 let abs_pos = search_start + pos;
142 let after_colon = abs_pos + desc_marker.len();
143 if let Some(quote_pos) = script[after_colon..].find('\'') {
144 let action_end = after_colon + quote_pos;
145 script.replace_range(after_colon..action_end, ":_zoi_packages");
146 }
147 search_start = after_colon;
148 }
149 }
150 Shell::Bash => {
151 let helpers = r#"
152_zoi_all_packages_comp() {
153 local cur=${COMP_WORDS[COMP_CWORD]}
154 local pkgs=$(zoi list -a --names 2>/dev/null)
155 COMPREPLY=( $(compgen -W "${pkgs}" -- "$cur") )
156}
157
158_zoi_installed_packages_comp() {
159 local cur=${COMP_WORDS[COMP_CWORD]}
160 local pkgs=$(zoi list --names 2>/dev/null)
161 COMPREPLY=( $(compgen -W "${pkgs}" -- "$cur") )
162}
163
164_zoi_wrapper() {
165 local cur="${COMP_WORDS[COMP_CWORD]}"
166 local prev="${COMP_WORDS[COMP_CWORD-1]}"
167 local cmd="${COMP_WORDS[1]}"
168
169 if [[ "$prev" == -* ]]; then
170 _zoi
171 return 0
172 fi
173
174 case $cmd in
175 install|i|in|add|show|exec|x|create|clone|use|tree|man|shell)
176 _zoi_all_packages_comp
177 return 0
178 ;;
179 uninstall|un|rm|remove|mark|m|update|up|why|files|pin|unpin|downgrade|dg|rollback)
180 _zoi_installed_packages_comp
181 return 0
182 ;;
183 esac
184
185 _zoi
186}
187complete -F _zoi_wrapper zoi
188"#;
189 script = format!("{}\n{}", script, helpers);
190 }
191 _ => {}
192 }
193 script
194}
195
196pub fn run(shell: Shell, scope: SetupScope) -> Result<()> {
197 if scope == SetupScope::System && !utils::is_admin() {
198 let exe = std::env::current_exe()?;
199 let args: Vec<String> = std::env::args().collect();
200 let status = Command::new("sudo")
201 .arg(&exe)
202 .args(&args[1..])
203 .status()
204 .map_err(|e| anyhow!("Failed to elevate with sudo: {}", e))?;
205 std::process::exit(status.code().unwrap_or(1));
206 }
207
208 println!(
209 "{} Setting up shell: {}...",
210 "::".bold().blue(),
211 shell.to_string().cyan()
212 );
213
214 let mut cmd = Cli::command();
215 install_completions(shell, scope, &mut cmd)?;
216
217 install_package_completions(shell, scope)?;
218
219 println!();
220
221 let scope_to_pass = match scope {
222 SetupScope::User => types::Scope::User,
223 SetupScope::System => types::Scope::System,
224 };
225 utils::setup_path(scope_to_pass)?;
226 Ok(())
227}
228
229fn get_completions_dir(scope: SetupScope, shell: &str) -> Result<PathBuf> {
230 match scope {
231 SetupScope::User => {
232 let home = dirs::home_dir().ok_or_else(|| anyhow!("Home directory not found"))?;
233 Ok(home.join(".zoi/pkgs/shell").join(shell))
234 }
235 SetupScope::System => {
236 if cfg!(target_os = "windows") {
237 Ok(PathBuf::from(format!(
238 "C:\\ProgramData\\zoi\\pkgs\\shell\\{}",
239 shell
240 )))
241 } else {
242 let base = match shell {
243 "bash" => "/usr/share/bash-completion/completions",
244 "zsh" => "/usr/share/zsh/site-functions",
245 "fish" => "/usr/share/fish/vendor_completions.d",
246 "elvish" => "/usr/share/elvish/lib",
247 _ => "/usr/local/share/zoi/completions",
248 };
249 Ok(PathBuf::from(base))
250 }
251 }
252 }
253}
254
255fn install_package_completions(shell: Shell, scope: SetupScope) -> Result<()> {
256 let shell_name = match shell {
257 Shell::Bash => "bash",
258 Shell::Zsh => "zsh",
259 Shell::Fish => "fish",
260 Shell::Elvish => "elvish",
261 _ => return Ok(()),
262 };
263
264 let completions_dir = get_completions_dir(scope, shell_name)?;
265 fs::create_dir_all(&completions_dir)?;
266
267 match shell {
268 Shell::Zsh => {
269 let fpath_entry = completions_dir.to_string_lossy().to_string();
270 println!(
271 "{} Add this to your .zshrc to load package completions:",
272 "::".bold().blue()
273 );
274 println!(" fpath=({:?} $fpath)", fpath_entry);
275 println!(" autoload -Uz compinit && compinit");
276 }
277 Shell::Bash => {
278 let bash_completion_dir = match scope {
279 SetupScope::User => {
280 let home =
281 dirs::home_dir().ok_or_else(|| anyhow!("Home directory not found"))?;
282 home.join(".local/share/bash-completion/completions")
283 }
284 SetupScope::System => PathBuf::from("/usr/share/bash-completion/completions"),
285 };
286 if bash_completion_dir.exists() {
287 println!(
288 "{} Package completions directory: {:?}",
289 "::".bold().blue(),
290 completions_dir
291 );
292 println!(" Completions from installed packages will be available automatically.");
293 }
294 }
295 Shell::Fish => {
296 println!(
297 "{} Package completions directory: {:?}",
298 "::".bold().blue(),
299 completions_dir
300 );
301 println!(" Completions from installed packages will be available automatically.");
302 }
303 _ => {}
304 }
305
306 Ok(())
307}
308
309pub fn print_hook(shell: Shell) -> Result<()> {
310 match shell {
311 Shell::Bash => {
312 println!(
313 r#"
314_zoi_hook() {{
315 local previous_exit_status=$?;
316 eval "$(zoi env --export-shell bash)";
317 return $previous_exit_status;
318}};
319if [[ ";${{PROMPT_COMMAND[*]:-}};" != *";_zoi_hook;"* ]]; then
320 if [[ "$(declare -p PROMPT_COMMAND 2>/dev/null)" == "declare -a"* ]]; then
321 PROMPT_COMMAND=(_zoi_hook "${{PROMPT_COMMAND[@]}}")
322 else
323 PROMPT_COMMAND="_zoi_hook${{PROMPT_COMMAND:+;$PROMPT_COMMAND}}"
324 fi
325fi
326"#
327 );
328 }
329 Shell::Zsh => {
330 println!(
331 r#"
332_zoi_hook() {{
333 eval "$(zoi env --export-shell zsh)";
334}};
335typeset -ag precmd_functions;
336if [[ -z "${{precmd_functions[(r)_zoi_hook]}}" ]]; then
337 precmd_functions+=(_zoi_hook);
338fi
339"#
340 );
341 }
342 Shell::Fish => {
343 println!(
344 r#"
345function _zoi_hook --on-variable PWD
346 zoi env --export-shell fish | source
347end
348"#
349 );
350 }
351 _ => return Err(anyhow!("Shell hook not supported for {:?}", shell)),
352 }
353 Ok(())
354}
355
356pub fn enter_ephemeral_shell(
357 package_sources: &[String],
358 run_cmd: Option<String>,
359 verbose: bool,
360 _plugin_manager: Option<&plugin::PluginManager>,
361) -> Result<()> {
362 if verbose {
363 println!("{} Resolving ephemeral environment...", "::".bold().blue());
364 }
365
366 let installed_before: HashSet<String> = local::get_installed_packages()?
367 .into_iter()
368 .map(|m| local::installed_manifest_source(&m))
369 .collect();
370
371 let (graph, _non_zoi_deps) = install::resolver::resolve_dependency_graph(
372 package_sources,
373 None,
374 false,
375 true,
376 true,
377 None,
378 !verbose,
379 )?;
380
381 let install_plan = install::plan::create_install_plan(&graph.nodes, None, false)?;
382 let stages = graph.toposort()?;
383
384 let mut session_installed = Vec::new();
385
386 if !install_plan.is_empty() {
387 if verbose {
388 println!(
389 "{} Preparing {} ephemeral dependencies...",
390 "::".bold().blue(),
391 install_plan.len()
392 );
393 }
394 let m = indicatif::MultiProgress::new();
395 if !verbose {
396 m.set_draw_target(indicatif::ProgressDrawTarget::hidden());
397 }
398 let session_installed_mutex = Mutex::new(Vec::new());
399
400 for stage in stages {
401 use rayon::prelude::*;
402 stage.into_par_iter().try_for_each(|pkg_id| -> Result<()> {
403 let node = graph
404 .nodes
405 .get(&pkg_id)
406 .ok_or_else(|| anyhow!("Package node missing from graph for '{}'", pkg_id))?;
407 let action = install_plan
408 .get(&pkg_id)
409 .ok_or_else(|| anyhow!("Install action missing for package '{}'", pkg_id))?;
410
411 let manifest = install::installer::install_node(
412 node,
413 action,
414 Some(&m),
415 None,
416 true,
417 false,
418 false,
419 verbose,
420 )?;
421
422 let mut session_lock = session_installed_mutex.lock().unwrap();
423 session_lock.push(manifest);
424 Ok(())
425 })?;
426 }
427 session_installed = session_installed_mutex.into_inner().unwrap();
428 }
429
430 let temp_dir = tempfile::Builder::new().prefix("zoi-shell-").tempdir()?;
431 let temp_bin_dir = temp_dir.path().join("bin");
432 fs::create_dir_all(&temp_bin_dir)?;
433
434 for node in graph.nodes.values() {
435 let handle = &node.registry_handle;
436 let pkg = &node.pkg;
437
438 let package_dir = local::get_package_dir(pkg.scope, handle, &pkg.repo, &pkg.name)?;
439 let version_dir = package_dir.join(&node.version);
440 let bin_dir = version_dir.join("bin");
441
442 if bin_dir.exists() {
443 for entry in fs::read_dir(bin_dir)? {
444 let entry = entry?;
445 let path = entry.path();
446 if path.is_file() || path.is_symlink() {
447 let file_name = path
448 .file_name()
449 .ok_or_else(|| anyhow!("Path has no file name: {:?}", path))?;
450 let dest = temp_bin_dir.join(file_name);
451 utils::symlink_file(&path, &dest)?;
452 }
453 }
454 }
455 }
456
457 let sep = if cfg!(windows) { ";" } else { ":" };
458 let mut new_path = temp_bin_dir.to_string_lossy().to_string();
459 if let Ok(old_path) = std::env::var("PATH") {
460 new_path = format!("{}{}{}", new_path, sep, old_path);
461 }
462
463 let package_list = package_sources.join(",");
464
465 let shell_bin = std::env::var("SHELL").unwrap_or_else(|_| {
466 if cfg!(windows) {
467 "pwsh".to_string()
468 } else {
469 "bash".to_string()
470 }
471 });
472
473 let mut envs = HashMap::new();
474 envs.insert("PATH".to_string(), new_path);
475 envs.insert("ZOI_SHELL".to_string(), "ephemeral".to_string());
476 envs.insert("IN_ZOI_SHELL".to_string(), "ephemeral".to_string());
477 envs.insert("ZOI_SHELL_PACKAGES".to_string(), package_list);
478
479 #[cfg(target_os = "linux")]
480 let mut shell_command = {
481 let sysroot = zoi_core::sysroot::get_sysroot();
482 if let Some(root) = sysroot {
483 if verbose {
484 println!(
485 "{} Entering shell within sysroot: {}",
486 "::".bold().yellow(),
487 root.display()
488 );
489 }
490
491 let extra_binds = vec![(temp_dir.path().to_path_buf(), temp_dir.path().to_path_buf())];
492
493 if let Some(cmd_str) = run_cmd {
494 let args = vec!["-c".to_string(), cmd_str];
495 crate::sandbox::wrap_command_in_root(
496 &root,
497 Path::new(&shell_bin),
498 &args,
499 &envs,
500 &extra_binds,
501 false,
502 )?
503 } else {
504 crate::sandbox::wrap_command_in_root(
505 &root,
506 Path::new(&shell_bin),
507 &[],
508 &envs,
509 &extra_binds,
510 false,
511 )?
512 }
513 } else if let Some(cmd_str) = run_cmd {
514 if verbose {
515 println!("{} Running: {}", "::".bold().blue(), cmd_str.cyan());
516 }
517 let mut c = if cfg!(windows) {
518 Command::new("pwsh")
519 } else {
520 Command::new("bash")
521 };
522 if !cfg!(windows) {
523 c.arg("-c");
524 } else {
525 c.arg("-Command");
526 }
527 c.arg(&cmd_str);
528 c.envs(&envs);
529 c
530 } else {
531 if verbose {
532 println!(
533 "{} Entering ephemeral shell (type 'exit' to leave)...",
534 "::".bold().green()
535 );
536 }
537 let mut c = Command::new(&shell_bin);
538 c.envs(&envs);
539 c
540 }
541 };
542
543 #[cfg(not(target_os = "linux"))]
544 let mut shell_command = {
545 if let Some(cmd_str) = run_cmd {
546 if verbose {
547 println!("{} Running: {}", "::".bold().blue(), cmd_str.cyan());
548 }
549 let mut c = if cfg!(windows) {
550 Command::new("pwsh")
551 } else {
552 Command::new("bash")
553 };
554 if !cfg!(windows) {
555 c.arg("-c");
556 } else {
557 c.arg("-Command");
558 }
559 c.arg(&cmd_str);
560 c.envs(&envs);
561 c
562 } else {
563 if verbose {
564 println!(
565 "{} Entering ephemeral shell (type 'exit' to leave)...",
566 "::".bold().green()
567 );
568 }
569 let mut c = Command::new(&shell_bin);
570 c.envs(&envs);
571 c
572 }
573 };
574
575 let status = shell_command.status()?;
576
577 if !session_installed.is_empty() {
578 if verbose {
579 println!("{} Cleaning up ephemeral packages...", "::".bold().blue());
580 }
581 for manifest in session_installed {
582 let ident = local::installed_manifest_source(&manifest);
583 if installed_before.contains(&ident) {
584 continue;
585 }
586 let version_dir = match get_version_dir_from_manifest(&manifest) {
587 Ok(d) => d,
588 Err(e) => {
589 eprintln!("Warning: failed to resolve path for {}: {}", ident, e);
590 continue;
591 }
592 };
593 if version_dir.exists()
594 && let Err(e) = fs::remove_dir_all(&version_dir)
595 {
596 eprintln!(
597 "Warning: failed to cleanup ephemeral package {}: {}",
598 ident, e
599 );
600 }
601 let package_dir = version_dir.parent().unwrap().to_path_buf();
602 if let Ok(mut entries) = fs::read_dir(&package_dir) {
603 let has_other_entries = entries.any(|e| {
604 e.as_ref()
605 .is_ok_and(|e| e.file_name() != "latest" && e.file_name() != "dependents")
606 });
607 if !has_other_entries {
608 let _ = fs::remove_dir_all(&package_dir);
609 }
610 }
611 }
612 }
613
614 if !status.success() {
615 std::process::exit(status.code().unwrap_or(1));
616 }
617
618 Ok(())
619}
620
621fn get_version_dir_from_manifest(manifest: &zoi_core::types::InstallManifest) -> Result<PathBuf> {
622 local::get_package_version_dir(
623 manifest.scope,
624 &manifest.registry_handle,
625 &manifest.repo,
626 &manifest.name,
627 &manifest.version,
628 )
629}