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 escalator = crate::pkg::utils::get_privilege_escalator()
201 .ok_or_else(|| anyhow!("Root privileges required for system-wide setup, but neither 'sudo' nor 'doas' was found."))?;
202
203 let status = Command::new(escalator)
204 .arg(&exe)
205 .args(&args[1..])
206 .status()
207 .map_err(|e| anyhow!("Failed to elevate privileges: {}", e))?;
208 std::process::exit(status.code().unwrap_or(1));
209 }
210
211 println!(
212 "{} Setting up shell: {}...",
213 "::".bold().blue(),
214 shell.to_string().cyan()
215 );
216
217 let mut cmd = Cli::command();
218 install_completions(shell, scope, &mut cmd)?;
219
220 install_package_completions(shell, scope)?;
221
222 println!();
223
224 let scope_to_pass = match scope {
225 SetupScope::User => types::Scope::User,
226 SetupScope::System => types::Scope::System,
227 };
228 utils::setup_path(scope_to_pass)?;
229 Ok(())
230}
231
232fn get_completions_dir(scope: SetupScope, shell: &str) -> Result<PathBuf> {
233 match scope {
234 SetupScope::User => {
235 let home = dirs::home_dir().ok_or_else(|| anyhow!("Home directory not found"))?;
236 Ok(home.join(".zoi/pkgs/shell").join(shell))
237 }
238 SetupScope::System => {
239 if cfg!(target_os = "windows") {
240 Ok(PathBuf::from(format!(
241 "C:\\ProgramData\\zoi\\pkgs\\shell\\{}",
242 shell
243 )))
244 } else {
245 let base = match shell {
246 "bash" => "/usr/share/bash-completion/completions",
247 "zsh" => "/usr/share/zsh/site-functions",
248 "fish" => "/usr/share/fish/vendor_completions.d",
249 "elvish" => "/usr/share/elvish/lib",
250 _ => "/usr/local/share/zoi/completions",
251 };
252 Ok(PathBuf::from(base))
253 }
254 }
255 }
256}
257
258fn install_package_completions(shell: Shell, scope: SetupScope) -> Result<()> {
259 let shell_name = match shell {
260 Shell::Bash => "bash",
261 Shell::Zsh => "zsh",
262 Shell::Fish => "fish",
263 Shell::Elvish => "elvish",
264 _ => return Ok(()),
265 };
266
267 let completions_dir = get_completions_dir(scope, shell_name)?;
268 fs::create_dir_all(&completions_dir)?;
269
270 match shell {
271 Shell::Zsh => {
272 let fpath_entry = completions_dir.to_string_lossy().to_string();
273 println!(
274 "{} Add this to your .zshrc to load package completions:",
275 "::".bold().blue()
276 );
277 println!(" fpath=({:?} $fpath)", fpath_entry);
278 println!(" autoload -Uz compinit && compinit");
279 }
280 Shell::Bash => {
281 let bash_completion_dir = match scope {
282 SetupScope::User => {
283 let home =
284 dirs::home_dir().ok_or_else(|| anyhow!("Home directory not found"))?;
285 home.join(".local/share/bash-completion/completions")
286 }
287 SetupScope::System => PathBuf::from("/usr/share/bash-completion/completions"),
288 };
289 if bash_completion_dir.exists() {
290 println!(
291 "{} Package completions directory: {:?}",
292 "::".bold().blue(),
293 completions_dir
294 );
295 println!(" Completions from installed packages will be available automatically.");
296 }
297 }
298 Shell::Fish => {
299 println!(
300 "{} Package completions directory: {:?}",
301 "::".bold().blue(),
302 completions_dir
303 );
304 println!(" Completions from installed packages will be available automatically.");
305 }
306 _ => {}
307 }
308
309 Ok(())
310}
311
312pub fn print_hook(shell: Shell) -> Result<()> {
313 match shell {
314 Shell::Bash => {
315 println!(
316 r#"
317_zoi_hook() {{
318 local previous_exit_status=$?;
319 eval "$(zoi env --export-shell bash)";
320 return $previous_exit_status;
321}};
322if [[ ";${{PROMPT_COMMAND[*]:-}};" != *";_zoi_hook;"* ]]; then
323 if [[ "$(declare -p PROMPT_COMMAND 2>/dev/null)" == "declare -a"* ]]; then
324 PROMPT_COMMAND=(_zoi_hook "${{PROMPT_COMMAND[@]}}")
325 else
326 PROMPT_COMMAND="_zoi_hook${{PROMPT_COMMAND:+;$PROMPT_COMMAND}}"
327 fi
328fi
329"#
330 );
331 }
332 Shell::Zsh => {
333 println!(
334 r#"
335_zoi_hook() {{
336 eval "$(zoi env --export-shell zsh)";
337}};
338typeset -ag precmd_functions;
339if [[ -z "${{precmd_functions[(r)_zoi_hook]}}" ]]; then
340 precmd_functions+=(_zoi_hook);
341fi
342"#
343 );
344 }
345 Shell::Fish => {
346 println!(
347 r#"
348function _zoi_hook --on-variable PWD
349 zoi env --export-shell fish | source
350end
351"#
352 );
353 }
354 _ => return Err(anyhow!("Shell hook not supported for {:?}", shell)),
355 }
356 Ok(())
357}
358
359pub fn enter_ephemeral_shell(
360 package_sources: &[String],
361 run_cmd: Option<String>,
362 verbose: bool,
363 _plugin_manager: Option<&plugin::PluginManager>,
364) -> Result<()> {
365 if verbose {
366 println!("{} Resolving ephemeral environment...", "::".bold().blue());
367 }
368
369 let installed_before: HashSet<String> = local::get_installed_packages()?
370 .into_iter()
371 .map(|m| local::installed_manifest_source(&m))
372 .collect();
373
374 let (graph, _non_zoi_deps) = install::resolver::resolve_dependency_graph(
375 package_sources,
376 None,
377 false,
378 true,
379 true,
380 None,
381 !verbose,
382 None,
383 )?;
384
385 let install_plan = install::plan::create_install_plan(&graph.nodes, None, false)?;
386 let stages = graph.toposort()?;
387
388 let mut session_installed = Vec::new();
389
390 if !install_plan.is_empty() {
391 if verbose {
392 println!(
393 "{} Preparing {} ephemeral dependencies...",
394 "::".bold().blue(),
395 install_plan.len()
396 );
397 }
398 let m = indicatif::MultiProgress::new();
399 if !verbose {
400 m.set_draw_target(indicatif::ProgressDrawTarget::hidden());
401 }
402 let session_installed_mutex = Mutex::new(Vec::new());
403
404 for stage in stages {
405 use rayon::prelude::*;
406 stage.into_par_iter().try_for_each(|pkg_id| -> Result<()> {
407 let node = graph
408 .nodes
409 .get(&pkg_id)
410 .ok_or_else(|| anyhow!("Package node missing from graph for '{}'", pkg_id))?;
411 let action = install_plan
412 .get(&pkg_id)
413 .ok_or_else(|| anyhow!("Install action missing for package '{}'", pkg_id))?;
414
415 let manifest = install::installer::install_node(
416 node,
417 action,
418 Some(&m),
419 None,
420 true,
421 false,
422 false,
423 verbose,
424 )?;
425
426 let mut session_lock = session_installed_mutex.lock().unwrap();
427 session_lock.push(manifest);
428 Ok(())
429 })?;
430 }
431 session_installed = session_installed_mutex.into_inner().unwrap();
432 }
433
434 let temp_dir = tempfile::Builder::new().prefix("zoi-shell-").tempdir()?;
435 let temp_bin_dir = temp_dir.path().join("bin");
436 fs::create_dir_all(&temp_bin_dir)?;
437
438 for node in graph.nodes.values() {
439 let handle = &node.registry_handle;
440 let pkg = &node.pkg;
441
442 let package_dir = local::get_package_dir(pkg.scope, handle, &pkg.repo, &pkg.name)?;
443 let version_dir = package_dir.join(&node.version);
444 let bin_dir = version_dir.join("bin");
445
446 if bin_dir.exists() {
447 for entry in fs::read_dir(bin_dir)? {
448 let entry = entry?;
449 let path = entry.path();
450 if path.is_file() || path.is_symlink() {
451 let file_name = path
452 .file_name()
453 .ok_or_else(|| anyhow!("Path has no file name: {:?}", path))?;
454 let dest = temp_bin_dir.join(file_name);
455 utils::symlink_file(&path, &dest)?;
456 }
457 }
458 }
459 }
460
461 let sep = if cfg!(windows) { ";" } else { ":" };
462 let mut new_path = temp_bin_dir.to_string_lossy().to_string();
463 if let Ok(old_path) = std::env::var("PATH") {
464 new_path = format!("{}{}{}", new_path, sep, old_path);
465 }
466
467 let package_list = package_sources.join(",");
468
469 let shell_bin = std::env::var("SHELL").unwrap_or_else(|_| {
470 if cfg!(windows) {
471 "pwsh".to_string()
472 } else {
473 "bash".to_string()
474 }
475 });
476
477 let mut envs = HashMap::new();
478 envs.insert("PATH".to_string(), new_path);
479 envs.insert("ZOI_SHELL".to_string(), "ephemeral".to_string());
480 envs.insert("IN_ZOI_SHELL".to_string(), "ephemeral".to_string());
481 envs.insert("ZOI_SHELL_PACKAGES".to_string(), package_list);
482
483 #[cfg(target_os = "linux")]
484 let mut shell_command = {
485 let sysroot = zoi_core::sysroot::get_sysroot();
486 if let Some(root) = sysroot {
487 if verbose {
488 println!(
489 "{} Entering shell within sysroot: {}",
490 "::".bold().yellow(),
491 root.display()
492 );
493 }
494
495 let extra_binds = vec![(temp_dir.path().to_path_buf(), temp_dir.path().to_path_buf())];
496
497 if let Some(cmd_str) = run_cmd {
498 let args = vec!["-c".to_string(), cmd_str];
499 crate::sandbox::wrap_command_in_root(
500 &root,
501 Path::new(&shell_bin),
502 &args,
503 &envs,
504 &extra_binds,
505 false,
506 )?
507 } else {
508 crate::sandbox::wrap_command_in_root(
509 &root,
510 Path::new(&shell_bin),
511 &[],
512 &envs,
513 &extra_binds,
514 false,
515 )?
516 }
517 } else if let Some(cmd_str) = run_cmd {
518 if verbose {
519 println!("{} Running: {}", "::".bold().blue(), cmd_str.cyan());
520 }
521 let mut c = if cfg!(windows) {
522 Command::new("pwsh")
523 } else {
524 Command::new("bash")
525 };
526 if !cfg!(windows) {
527 c.arg("-c");
528 } else {
529 c.arg("-Command");
530 }
531 c.arg(&cmd_str);
532 c.envs(&envs);
533 c
534 } else {
535 if verbose {
536 println!(
537 "{} Entering ephemeral shell (type 'exit' to leave)...",
538 "::".bold().green()
539 );
540 }
541 let mut c = Command::new(&shell_bin);
542 c.envs(&envs);
543 c
544 }
545 };
546
547 #[cfg(not(target_os = "linux"))]
548 let mut shell_command = {
549 if let Some(cmd_str) = run_cmd {
550 if verbose {
551 println!("{} Running: {}", "::".bold().blue(), cmd_str.cyan());
552 }
553 let mut c = if cfg!(windows) {
554 Command::new("pwsh")
555 } else {
556 Command::new("bash")
557 };
558 if !cfg!(windows) {
559 c.arg("-c");
560 } else {
561 c.arg("-Command");
562 }
563 c.arg(&cmd_str);
564 c.envs(&envs);
565 c
566 } else {
567 if verbose {
568 println!(
569 "{} Entering ephemeral shell (type 'exit' to leave)...",
570 "::".bold().green()
571 );
572 }
573 let mut c = Command::new(&shell_bin);
574 c.envs(&envs);
575 c
576 }
577 };
578
579 let status = shell_command.status()?;
580
581 if !session_installed.is_empty() {
582 if verbose {
583 println!("{} Cleaning up ephemeral packages...", "::".bold().blue());
584 }
585 for manifest in session_installed {
586 let ident = local::installed_manifest_source(&manifest);
587 if installed_before.contains(&ident) {
588 continue;
589 }
590 let version_dir = match get_version_dir_from_manifest(&manifest) {
591 Ok(d) => d,
592 Err(e) => {
593 eprintln!("Warning: failed to resolve path for {}: {}", ident, e);
594 continue;
595 }
596 };
597 if version_dir.exists()
598 && let Err(e) = fs::remove_dir_all(&version_dir)
599 {
600 eprintln!(
601 "Warning: failed to cleanup ephemeral package {}: {}",
602 ident, e
603 );
604 }
605 let package_dir = version_dir.parent().unwrap().to_path_buf();
606 if let Ok(mut entries) = fs::read_dir(&package_dir) {
607 let has_other_entries = entries.any(|e| {
608 e.as_ref()
609 .is_ok_and(|e| e.file_name() != "latest" && e.file_name() != "dependents")
610 });
611 if !has_other_entries {
612 let _ = fs::remove_dir_all(&package_dir);
613 }
614 }
615 }
616 }
617
618 if !status.success() {
619 std::process::exit(status.code().unwrap_or(1));
620 }
621
622 Ok(())
623}
624
625fn get_version_dir_from_manifest(manifest: &zoi_core::types::InstallManifest) -> Result<PathBuf> {
626 local::get_package_version_dir(
627 manifest.scope,
628 &manifest.registry_handle,
629 &manifest.repo,
630 &manifest.name,
631 &manifest.version,
632 )
633}