From e873ec96023dcaee87ee520da44df314650ea160 Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Sun, 30 Nov 2025 20:46:02 -0800
Subject: [PATCH] Refactor WIM setup-image logic

---
 win_installer/src/app.rs   | 26 ++++++++++++--------------
 win_installer/src/state.rs |  9 +++------
 win_installer/src/wim.rs   |  6 +++---
 3 files changed, 18 insertions(+), 23 deletions(-)

diff --git a/win_installer/src/app.rs b/win_installer/src/app.rs
index c15f3b1..2c9b953 100644
--- a/win_installer/src/app.rs
+++ b/win_installer/src/app.rs
@@ -109,19 +109,7 @@ impl App {
             Mode::SelectTableType => Mode::ScanWinSources,
             Mode::ScanWinSources => Mode::SelectWinSource,
             Mode::SelectWinSource => Mode::SelectWinImage,
-            Mode::SelectWinImage => {
-                let mut next_mode = Mode::SetUserName;
-                // TODO: FIXME - Race condition?
-                // if let Ok(wim_sources) = self.state.wim_sources.lock()
-                //     && let Some(index) = self.state.wim_image_index
-                // {
-                //     let image = wim_sources.get_file(index);
-                //     if !image.is_installer {
-                //         next_mode = Mode::Confirm;
-                //     }
-                // }
-                next_mode
-            }
+            Mode::SelectWinImage => Mode::SetUserName,
             Mode::SetUserName => Mode::Confirm,
             Mode::Confirm => Mode::Process, // i.e. format, apply, etc
             Mode::Process | Mode::Done => Mode::Done,
@@ -162,6 +150,16 @@ impl App {
                 ))?;
             }
             Mode::ScanWinSources => self.state.scan_wim_local(ScanType::WindowsInstallers),
+            Mode::SetUserName => {
+                if let Ok(wim_sources) = self.state.wim_sources.lock()
+                    && let Some(index) = self.state.wim_image_index
+                {
+                    let image = wim_sources.get_file(index);
+                    if !image.is_setup {
+                        self.action_tx.send(Action::NextScreen)?;
+                    }
+                }
+            }
             Mode::Done => {
                 self.action_tx
                     .send(Action::DisplayPopup(popup::Type::Success, popup::fortune()))?;
@@ -725,7 +723,7 @@ fn build_right_items(app: &App) -> Action {
                             DVLine::blank(),
                         ]);
                     }
-                    if wim_file.is_installer
+                    if wim_file.is_setup
                         && let Some(username) = &app.state.username
                     {
                         label_dv_lines.append(&mut vec![DVLine {
diff --git a/win_installer/src/state.rs b/win_installer/src/state.rs
index 98428c3..4c948a8 100644
--- a/win_installer/src/state.rs
+++ b/win_installer/src/state.rs
@@ -27,8 +27,6 @@ use core::{
     },
 };
 
-use tracing::info;
-
 use crate::{
     net::connect_network_share,
     wim::{WimFile, WimSources, parse_wim_file},
@@ -123,7 +121,7 @@ pub fn scan_local_drives(
     wim_sources_arc: Arc<Mutex<WimSources>>,
     scan_type: ScanType,
 ) {
-    let mut to_check = vec![String::from(".")];
+    let mut to_check: Vec<String> = Vec::new();
     let mut wim_files: Vec<WimFile> = Vec::new();
 
     // Get drive letters
@@ -146,14 +144,13 @@ pub fn scan_local_drives(
 
     // Scan drives
     to_check.iter().for_each(|scan_path| {
-        let installer = scan_path.ends_with("\\Images");
-        info!("Scanning: {}", &scan_path);
+        let is_setup = scan_path.ends_with("\\Images");
         if let Ok(read_dir) = read_dir(scan_path) {
             read_dir.for_each(|item| {
                 if let Ok(item) = item
                     && item.file_name().to_string_lossy().ends_with(".wim")
                     && let Some(path_str) = item.path().to_str()
-                    && let Ok(new_source) = parse_wim_file(path_str, installer)
+                    && let Ok(new_source) = parse_wim_file(path_str, is_setup)
                 {
                     wim_files.push(new_source);
                 }
diff --git a/win_installer/src/wim.rs b/win_installer/src/wim.rs
index 032ff86..c5e7541 100644
--- a/win_installer/src/wim.rs
+++ b/win_installer/src/wim.rs
@@ -66,7 +66,7 @@ static WIN_BUILDS: LazyLock<HashMap<&str, &str>> = LazyLock::new(|| {
 pub struct WimFile {
     pub path: String,
     pub images: Vec<WimImage>,
-    pub is_installer: bool,
+    pub is_setup: bool,
 }
 
 impl WimFile {
@@ -206,7 +206,7 @@ fn get_wim_xml(wim_file: &str) -> std::io::Result<File> {
     Ok(file)
 }
 
-pub fn parse_wim_file(wim_file: &str, installer: bool) -> std::io::Result<WimFile> {
+pub fn parse_wim_file(wim_file: &str, is_setup: bool) -> std::io::Result<WimFile> {
     let mut wim_images: Vec<WimImage> = Vec::new();
     if !Path::new(wim_file).exists() {
         return Err(std::io::Error::new(
@@ -271,7 +271,7 @@ pub fn parse_wim_file(wim_file: &str, installer: bool) -> std::io::Result<WimFil
     let wim_file = WimFile {
         path: wim_file.to_string(),
         images: wim_images,
-        is_installer: installer,
+        is_setup,
     };
 
     Ok(wim_file)