# Stage-Dispatcher.ps1 - Chain controller for the shopfloor imaging pipeline.
#
# Reads C:\Enrollment\setup-stage.txt to know which step of the imaging
# process we're in, launches the right script, advances the stage, and
# re-registers itself via RunOnce so the next logon picks up the next step.
#
# STAGE CHAIN:
#   "shopfloor-setup"  -> Run-ShopfloorSetup.ps1 (installs apps, deferred org)
#                         Run-ShopfloorSetup reboots via shutdown /r /t 10;
#                         dispatcher has ~10 seconds to write the next stage
#                         and re-register RunOnce before the machine goes down.
#
#   "sync-intune"      -> Monitor-IntuneProgress.ps1 -Unattended
#                         Monitor exits 2 = pre-reboot Y pressed (no reboot yet)
#                           -> dispatcher re-registers RunOnce, initiates reboot
#                         Monitor exits 0 = post-reboot install complete
#                           -> dispatcher advances to configure-pc
#
#   "configure-pc"     -> Configure-PC.ps1 (startup items, machine number)
#                         -> dispatcher deletes stage file. Chain complete.
#
#   (file missing)     -> exit 0. Setup is done, nothing to do.
#
# ENTRY POINT:
#   Registered as a RunOnce value by run-enrollment.ps1 after PPKG install.
#   Each stage re-registers RunOnce for the next boot. RunOnce entries are
#   consumed by Windows after execution, so no cleanup needed.
#
# LOGGING:
#   Appends to C:\Logs\SFLD\stage-dispatcher.log via Start-Transcript.

$ErrorActionPreference = 'Continue'

# --- Transcript ---
$logDir = 'C:\Logs\SFLD'
if (-not (Test-Path $logDir)) {
    try { New-Item -ItemType Directory -Path $logDir -Force | Out-Null } catch { $logDir = $env:TEMP }
}
$transcriptPath = Join-Path $logDir 'stage-dispatcher.log'
try { Start-Transcript -Path $transcriptPath -Append -Force | Out-Null } catch {}

Write-Host ""
Write-Host "================================================================"
Write-Host "=== Stage-Dispatcher.ps1 $(Get-Date -Format 'yyyy-MM-dd HH:mm:ss') ==="
Write-Host "================================================================"
Write-Host "Running as: $([Security.Principal.WindowsIdentity]::GetCurrent().Name)"

# --- Paths ---
$stageFile      = 'C:\Enrollment\setup-stage.txt'
$dispatcherPath = $MyInvocation.MyCommand.Path
$enrollDir      = 'C:\Enrollment'
$shopfloorDir   = Join-Path $enrollDir 'shopfloor-setup\Shopfloor'
$runOnceKey     = 'HKLM:\SOFTWARE\Microsoft\Windows\CurrentVersion\RunOnce'

# --- Stage file ---
if (-not (Test-Path -LiteralPath $stageFile)) {
    Write-Host "No stage file found at $stageFile - setup complete, nothing to do."
    try { Stop-Transcript | Out-Null } catch {}
    exit 0
}

$stage = (Get-Content -LiteralPath $stageFile -First 1).Trim()
Write-Host "Current stage: $stage"

# --- Helper: re-register this dispatcher to run at next logon ---
function Register-NextRun {
    try {
        Set-ItemProperty -Path $runOnceKey -Name 'ShopfloorSetup' `
            -Value "powershell.exe -NoProfile -ExecutionPolicy Bypass -File `"$dispatcherPath`"" `
            -Type String -Force -ErrorAction Stop
        Write-Host "RunOnce registered for next logon."
    } catch {
        Write-Warning "Failed to register RunOnce: $_"
    }
}

# --- Stage dispatch ---
switch ($stage) {

    'shopfloor-setup' {
        Write-Host ""
        Write-Host "--- Stage: shopfloor-setup ---"
        $script = Join-Path $enrollDir 'Run-ShopfloorSetup.ps1'
        if (-not (Test-Path -LiteralPath $script)) {
            Write-Warning "Run-ShopfloorSetup.ps1 not found at $script"
            break
        }

        # -FromDispatcher bypasses the stage-file gate at the top of
        # Run-ShopfloorSetup (which would otherwise see the stage file
        # and exit immediately thinking it should defer to us).
        & $script -FromDispatcher

        # Cancel whatever reboot Run-ShopfloorSetup scheduled (shutdown /r
        # /t 10) so we can advance the stage and re-register RunOnce WITHOUT
        # racing a 10-second fuse. Then reboot on our own terms.
        Write-Host "Run-ShopfloorSetup.ps1 finished. Canceling its reboot so we can advance safely."
        cmd /c "shutdown /a 2>nul" *>$null
        Set-Content -LiteralPath $stageFile -Value 'sync-intune' -Force
        Register-NextRun
        Write-Host "Stage advanced to sync-intune. Rebooting."
        shutdown /r /t 5
    }

    'sync-intune' {
        Write-Host ""
        Write-Host "--- Stage: sync-intune ---"
        $monitor = Join-Path $shopfloorDir 'lib\Monitor-IntuneProgress.ps1'
        if (-not (Test-Path -LiteralPath $monitor)) {
            Write-Warning "Monitor-IntuneProgress.ps1 not found at $monitor"
            break
        }

        # Launch the monitor in THIS console window (interactive - tech sees
        # QR code + status table). -Unattended makes it exit with a status
        # code instead of prompting/rebooting, so we keep control here.
        Write-Host "Launching Intune sync monitor (unattended mode)..."
        & $monitor -Unattended
        $monitorExit = $LASTEXITCODE
        Write-Host "Monitor exited with code: $monitorExit"

        switch ($monitorExit) {
            0 {
                # Post-reboot install complete. All milestones reached.
                # Chain directly to Configure-PC (same session, no reboot).
                Write-Host "Post-reboot install complete. Launching Configure-PC..."
                Set-Content -LiteralPath $stageFile -Value 'configure-pc' -Force

                $configScript = Join-Path $shopfloorDir 'Configure-PC.ps1'
                if (Test-Path -LiteralPath $configScript) {
                    & $configScript
                } else {
                    Write-Warning "Configure-PC.ps1 not found at $configScript"
                }

                # Chain complete.
                Write-Host "Imaging chain complete. Removing stage file."
                Remove-Item -LiteralPath $stageFile -Force -ErrorAction SilentlyContinue
            }
            2 {
                # Pre-reboot deployment done. User confirmed reboot.
                # Re-register RunOnce so we come back after the reboot to
                # monitor the post-reboot install phase.
                Write-Host "Pre-reboot phase complete. Rebooting..."
                Register-NextRun
                # Stage stays at 'sync-intune' - on next boot the monitor
                # will detect post-reboot state and monitor until done.
                shutdown /r /t 5
            }
            default {
                Write-Warning "Unexpected monitor exit code: $monitorExit"
                Write-Host "Re-registering for retry on next logon."
                Register-NextRun
            }
        }
    }

    'configure-pc' {
        Write-Host ""
        Write-Host "--- Stage: configure-pc ---"
        $configScript = Join-Path $shopfloorDir 'Configure-PC.ps1'
        if (Test-Path -LiteralPath $configScript) {
            & $configScript
        } else {
            Write-Warning "Configure-PC.ps1 not found at $configScript"
        }

        Write-Host "Imaging chain complete. Removing stage file."
        Remove-Item -LiteralPath $stageFile -Force -ErrorAction SilentlyContinue
    }

    default {
        Write-Warning "Unknown stage: '$stage'. Removing stage file."
        Remove-Item -LiteralPath $stageFile -Force -ErrorAction SilentlyContinue
    }
}

Write-Host ""
Write-Host "=== Stage-Dispatcher.ps1 finished $(Get-Date -Format 'yyyy-MM-dd HH:mm:ss') ==="
try { Stop-Transcript | Out-Null } catch {}
