1  //! Prerequisite installation.
  2  //!
  3  //! Installs system packages, Rust toolchain, and optional CUDA toolkit.
  4  
  5  use acdc_core::{Error, Result};
  6  use acdc_tui::output;
  7  use std::process::Command;
  8  use tracing::{debug, warn};
  9  
 10  /// Required apt packages for building from source.
 11  pub const APT_PACKAGES: &[&str] = &[
 12      "build-essential",
 13      "pkg-config",
 14      "libssl-dev",
 15      "git",
 16      "curl",
 17      "clang",
 18      "cmake",
 19      "libclang-dev",
 20      "llvm-dev",
 21  ];
 22  
 23  /// Optional packages for enhanced functionality.
 24  pub const APT_PACKAGES_OPTIONAL: &[&str] = &[
 25      "jq",    // JSON processing
 26      "htop",  // System monitoring
 27      "tmux",  // Terminal multiplexer
 28      "unzip", // Archive extraction
 29  ];
 30  
 31  /// CUDA apt packages for GPU provers.
 32  pub const CUDA_PACKAGES: &[&str] = &["nvidia-cuda-toolkit", "nvidia-cuda-dev"];
 33  
 34  /// Install required system prerequisites.
 35  pub async fn install(include_cuda: bool) -> Result<()> {
 36      output::section("Installing Prerequisites");
 37  
 38      // Check if we're on a Debian-based system
 39      if !is_debian_based() {
 40          output::warning("Non-Debian system detected. Manual package installation may be required.");
 41          output::info("Required packages", &APT_PACKAGES.join(", "));
 42      } else {
 43          // Update package lists
 44          update_apt().await?;
 45  
 46          // Install required packages
 47          install_apt_packages(APT_PACKAGES).await?;
 48  
 49          // Install optional packages (don't fail if these don't work)
 50          if let Err(e) = install_apt_packages(APT_PACKAGES_OPTIONAL).await {
 51              warn!("Optional packages failed to install: {}", e);
 52          }
 53  
 54          // Install CUDA if requested
 55          if include_cuda {
 56              output::subsection("CUDA Toolkit");
 57              if check_nvidia_gpu() {
 58                  install_cuda().await?;
 59              } else {
 60                  output::warning("No NVIDIA GPU detected, skipping CUDA installation");
 61              }
 62          }
 63      }
 64  
 65      // Install Rust if not present
 66      if !check_rust() {
 67          output::subsection("Rust Toolchain");
 68          install_rust().await?;
 69      } else {
 70          output::success("Rust toolchain already installed");
 71          // Ensure we have the right version
 72          update_rust().await?;
 73      }
 74  
 75      output::success("Prerequisites installation complete");
 76      Ok(())
 77  }
 78  
 79  /// Check if we're on a Debian-based system.
 80  fn is_debian_based() -> bool {
 81      std::path::Path::new("/etc/debian_version").exists()
 82          || Command::new("apt-get")
 83              .arg("--version")
 84              .output()
 85              .map(|o| o.status.success())
 86              .unwrap_or(false)
 87  }
 88  
 89  /// Update apt package lists.
 90  async fn update_apt() -> Result<()> {
 91      output::status("Updating package lists...");
 92  
 93      let output = Command::new("sudo")
 94          .args(["apt-get", "update", "-qq"])
 95          .output()
 96          .map_err(Error::Io)?;
 97  
 98      if !output.status.success() {
 99          let stderr = String::from_utf8_lossy(&output.stderr);
100          return Err(Error::Installation(format!(
101              "apt-get update failed: {}",
102              stderr
103          )));
104      }
105  
106      Ok(())
107  }
108  
109  /// Install apt packages.
110  async fn install_apt_packages(packages: &[&str]) -> Result<()> {
111      if packages.is_empty() {
112          return Ok(());
113      }
114  
115      output::status(&format!("Installing packages: {}", packages.join(", ")));
116  
117      let mut cmd = Command::new("sudo");
118      cmd.args(["apt-get", "install", "-y", "-qq"]);
119      cmd.args(packages);
120  
121      let output = cmd.output().map_err(Error::Io)?;
122  
123      if !output.status.success() {
124          let stderr = String::from_utf8_lossy(&output.stderr);
125          return Err(Error::Installation(format!(
126              "apt-get install failed: {}",
127              stderr
128          )));
129      }
130  
131      output::success(&format!("Installed {} packages", packages.len()));
132      Ok(())
133  }
134  
135  /// Check if an NVIDIA GPU is present.
136  fn check_nvidia_gpu() -> bool {
137      // Check lspci for NVIDIA
138      Command::new("lspci")
139          .output()
140          .map(|o| {
141              let stdout = String::from_utf8_lossy(&o.stdout);
142              stdout.to_lowercase().contains("nvidia")
143          })
144          .unwrap_or(false)
145  }
146  
147  /// Install CUDA toolkit.
148  async fn install_cuda() -> Result<()> {
149      output::status("Installing CUDA toolkit...");
150  
151      // First try the standard packages
152      let result = install_apt_packages(CUDA_PACKAGES).await;
153  
154      if result.is_err() {
155          output::warning("Standard CUDA packages not available");
156          output::info("Alternative", "Install CUDA from NVIDIA's repository");
157  
158          // Try to add NVIDIA repository and install
159          install_cuda_from_nvidia_repo().await?;
160      }
161  
162      // Verify CUDA installation
163      if check_cuda() {
164          output::success("CUDA toolkit installed successfully");
165      } else {
166          output::warning("CUDA installed but nvcc not in PATH. You may need to add it manually.");
167          output::info(
168              "Hint",
169              "Add to ~/.bashrc: export PATH=/usr/local/cuda/bin:$PATH",
170          );
171      }
172  
173      Ok(())
174  }
175  
176  /// Install CUDA from NVIDIA repository (fallback).
177  async fn install_cuda_from_nvidia_repo() -> Result<()> {
178      output::status("Setting up NVIDIA CUDA repository...");
179  
180      // Download and install the CUDA keyring
181      let keyring_url = "https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb";
182  
183      let download = Command::new("curl")
184          .args(["-fsSL", "-o", "/tmp/cuda-keyring.deb", keyring_url])
185          .output()
186          .map_err(Error::Io)?;
187  
188      if !download.status.success() {
189          return Err(Error::Installation(
190              "Failed to download CUDA keyring".to_string(),
191          ));
192      }
193  
194      let install = Command::new("sudo")
195          .args(["dpkg", "-i", "/tmp/cuda-keyring.deb"])
196          .output()
197          .map_err(Error::Io)?;
198  
199      if !install.status.success() {
200          let stderr = String::from_utf8_lossy(&install.stderr);
201          return Err(Error::Installation(format!(
202              "Failed to install CUDA keyring: {}",
203              stderr
204          )));
205      }
206  
207      // Update and install CUDA
208      update_apt().await?;
209  
210      let cuda_install = Command::new("sudo")
211          .args(["apt-get", "install", "-y", "cuda-toolkit"])
212          .output()
213          .map_err(Error::Io)?;
214  
215      if !cuda_install.status.success() {
216          let stderr = String::from_utf8_lossy(&cuda_install.stderr);
217          return Err(Error::Installation(format!(
218              "Failed to install CUDA toolkit: {}",
219              stderr
220          )));
221      }
222  
223      Ok(())
224  }
225  
226  /// Check if Rust is installed.
227  pub fn check_rust() -> bool {
228      Command::new("rustc")
229          .arg("--version")
230          .output()
231          .map(|o| o.status.success())
232          .unwrap_or(false)
233  }
234  
235  /// Get installed Rust version.
236  pub fn get_rust_version() -> Option<String> {
237      Command::new("rustc")
238          .arg("--version")
239          .output()
240          .ok()
241          .and_then(|o| {
242              if o.status.success() {
243                  String::from_utf8(o.stdout)
244                      .ok()
245                      .map(|s| s.trim().to_string())
246              } else {
247                  None
248              }
249          })
250  }
251  
252  /// Install Rust via rustup.
253  async fn install_rust() -> Result<()> {
254      output::status("Installing Rust via rustup...");
255  
256      // Download and run rustup installer
257      let curl = Command::new("curl")
258          .args([
259              "--proto",
260              "=https",
261              "--tlsv1.2",
262              "-sSf",
263              "https://sh.rustup.rs",
264          ])
265          .output()
266          .map_err(Error::Io)?;
267  
268      if !curl.status.success() {
269          return Err(Error::Installation(
270              "Failed to download rustup installer".to_string(),
271          ));
272      }
273  
274      // Run the installer non-interactively
275      let install = Command::new("sh")
276          .args([
277              "-c",
278              "curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y",
279          ])
280          .env("RUSTUP_INIT_SKIP_PATH_CHECK", "yes")
281          .output()
282          .map_err(Error::Io)?;
283  
284      if !install.status.success() {
285          let stderr = String::from_utf8_lossy(&install.stderr);
286          return Err(Error::Installation(format!(
287              "Rust installation failed: {}",
288              stderr
289          )));
290      }
291  
292      output::success("Rust installed successfully");
293      output::info("Note", "Run 'source ~/.cargo/env' or restart your shell");
294  
295      Ok(())
296  }
297  
298  /// Update Rust to latest stable.
299  async fn update_rust() -> Result<()> {
300      output::status("Updating Rust toolchain...");
301  
302      // Source cargo env if needed
303      let cargo_path = dirs::home_dir()
304          .map(|h| h.join(".cargo/bin/rustup"))
305          .unwrap_or_else(|| "rustup".into());
306  
307      let rustup = if cargo_path.exists() {
308          cargo_path.to_string_lossy().to_string()
309      } else {
310          "rustup".to_string()
311      };
312  
313      let output = Command::new(&rustup)
314          .args(["update", "stable"])
315          .output()
316          .map_err(Error::Io)?;
317  
318      if !output.status.success() {
319          // Non-fatal, just log
320          warn!("rustup update failed, continuing anyway");
321      } else {
322          debug!("Rust updated successfully");
323      }
324  
325      Ok(())
326  }
327  
328  /// Check if CUDA is installed.
329  pub fn check_cuda() -> bool {
330      Command::new("nvcc")
331          .arg("--version")
332          .output()
333          .map(|o| o.status.success())
334          .unwrap_or(false)
335  }
336  
337  /// Get installed CUDA version.
338  pub fn get_cuda_version() -> Option<String> {
339      Command::new("nvcc")
340          .arg("--version")
341          .output()
342          .ok()
343          .and_then(|o| {
344              if o.status.success() {
345                  let stdout = String::from_utf8_lossy(&o.stdout);
346                  // Parse version from output like "Cuda compilation tools, release 12.0, V12.0.140"
347                  stdout
348                      .lines()
349                      .find(|l| l.contains("release"))
350                      .and_then(|l| {
351                          l.split("release")
352                              .nth(1)
353                              .map(|s| s.trim().trim_matches(',').trim().to_string())
354                      })
355              } else {
356                  None
357              }
358          })
359  }
360  
361  /// Check if a specific command is available.
362  pub fn check_command(cmd: &str) -> bool {
363      Command::new("which")
364          .arg(cmd)
365          .output()
366          .map(|o| o.status.success())
367          .unwrap_or(false)
368  }
369  
370  /// Prerequisite check result.
371  #[derive(Debug, Clone)]
372  pub struct PrerequisiteStatus {
373      pub apt_packages: bool,
374      pub rust: bool,
375      pub rust_version: Option<String>,
376      pub cuda: bool,
377      pub cuda_version: Option<String>,
378      pub git: bool,
379      pub curl: bool,
380  }
381  
382  impl PrerequisiteStatus {
383      /// Check all prerequisites and return status.
384      pub fn check() -> Self {
385          Self {
386              apt_packages: check_command("gcc") && check_command("pkg-config"),
387              rust: check_rust(),
388              rust_version: get_rust_version(),
389              cuda: check_cuda(),
390              cuda_version: get_cuda_version(),
391              git: check_command("git"),
392              curl: check_command("curl"),
393          }
394      }
395  
396      /// Check if all required prerequisites are met.
397      pub fn is_ready(&self) -> bool {
398          self.rust && self.git && self.curl && self.apt_packages
399      }
400  
401      /// Check if CUDA prerequisites are met (for provers).
402      pub fn is_cuda_ready(&self) -> bool {
403          self.cuda
404      }
405  }
406  
407  /// Print prerequisite status.
408  pub fn print_status(status: &PrerequisiteStatus) {
409      output::subsection("Prerequisite Status");
410  
411      if status.apt_packages {
412          output::success("Build tools installed");
413      } else {
414          output::error("Build tools missing (run: sudo apt install build-essential pkg-config)");
415      }
416  
417      if status.git {
418          output::success("Git installed");
419      } else {
420          output::error("Git not found");
421      }
422  
423      if status.curl {
424          output::success("curl installed");
425      } else {
426          output::error("curl not found");
427      }
428  
429      if status.rust {
430          if let Some(ref version) = status.rust_version {
431              output::success(&format!("Rust: {}", version));
432          } else {
433              output::success("Rust installed");
434          }
435      } else {
436          output::error("Rust not installed");
437      }
438  
439      if status.cuda {
440          if let Some(ref version) = status.cuda_version {
441              output::success(&format!("CUDA: {}", version));
442          } else {
443              output::success("CUDA installed");
444          }
445      } else {
446          output::info("CUDA", "Not installed (optional, required for provers)");
447      }
448  }