# 06-OrganizeDesktop.ps1 - Single source of truth for the Public Desktop
# layout. Builds three category folders and fills them, so end users see
# a clean desktop instead of 20+ loose icons.
#
# Creates three subfolders under C:\Users\Public\Desktop:
#   Office\          - Excel, Word, PowerPoint, Outlook, OneNote, etc.
#   Shopfloor Tools\ - UDC, eDNC, HostExplorer, MarkZebra, PC-DMIS, Hexagon,
#                      WJ Shopfloor, Defect_Tracker, NTLARS, etc.
#   Web Links\       - .url shortcuts (internal portals)
#
# Two phases:
#   PHASE 1 (Invoke-DesktopSweep): walks everything at the desktop ROOT,
#   classifies each .lnk / .url by filename / extension / .lnk target, and
#   moves it into the right category folder. Unknown items are LEFT AT
#   THE ROOT on purpose. eDNC.lnk and NTLARS.lnk are allowlisted to stay
#   at the root - end users click them too often to bury in a folder.
#
#   PHASE 2 (Add-ShopfloorToolsApps): materializes specific app shortcuts
#   into Shopfloor Tools\ from known .exe paths (UDC, eDNC, NTLARS). For
#   MSI-advertised shortcuts that don't have a resolvable target (WJ
#   Shopfloor, Defect_Tracker), finds the existing .lnk wherever it lives
#   and copies it in. This runs AFTER the sweep so it overwrites any
#   stale copies the sweeper may have moved in.
#
# Also registers an "Organize Public Desktop" scheduled task that re-runs
# the sweep (phase 1) at every logon. Phase 2 doesn't re-run - the
# shortcuts it creates are stable and 06 on first run populates them.
#
# 07-TaskbarLayout.ps1 reads the result of this script and writes
# LayoutModification.xml to pin specific shortcuts to the taskbar.

$ErrorActionPreference = 'Continue'

# ============================================================================
# Admin check - writing to C:\Users\Public\Desktop\* and registering
# system-wide scheduled tasks both require elevation. Fail fast if not
# admin so we don't spam half-successful garbage across the output.
# ============================================================================
$isAdmin = ([Security.Principal.WindowsPrincipal][Security.Principal.WindowsIdentity]::GetCurrent()).IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)
if (-not $isAdmin) {
    Write-Host ""
    Write-Host "ERROR: 06-OrganizeDesktop.ps1 must run as Administrator." -ForegroundColor Red
    Write-Host "       Re-run from an elevated PowerShell:" -ForegroundColor Red
    Write-Host "         Start -> type 'powershell' -> right-click -> Run as administrator" -ForegroundColor Red
    Write-Host ""
    exit 1
}

$publicDesktop       = 'C:\Users\Public\Desktop'
$shopfloorToolsDir   = Join-Path $publicDesktop 'Shopfloor Tools'
$scriptPath          = $MyInvocation.MyCommand.Path

# Filenames that always stay at the desktop root regardless of classification.
# End users click these many times a day and an extra folder click is real
# friction. Phase 2 also drops these here as a post-sweep safety net.
$keepAtRoot = @(
    'eDNC.lnk',
    'NTLARS.lnk'
)

# ============================================================================
# Phase 1: Sweep loose shortcuts at desktop root into category folders
# ============================================================================
function Invoke-DesktopSweep {
    param([string]$DesktopPath)

    if (-not (Test-Path -LiteralPath $DesktopPath)) {
        Write-Host "Public desktop not found at $DesktopPath - skipping sweep."
        return
    }

    # Category definitions. Order matters: first match wins. Each category
    # can classify by file Extension, file Name (regex), or .lnk Target
    # (regex, resolved via WScript.Shell).
    $categories = [ordered]@{
        'Office' = @{
            Name = @(
                '^(Excel|Word|PowerPoint|Outlook|OneNote|Access|Publisher|OneDrive)(\s|$)'
            )
            Target = @(
                '\\(EXCEL|WINWORD|POWERPNT|OUTLOOK|ONENOTE|MSACCESS|MSPUB)\.EXE$',
                '\\OneDrive\.exe$'
            )
        }
        'Shopfloor Tools' = @{
            Name = @(
                '^UDC',
                'eDNC', '\bDNC\b', 'DncMain', 'GE DNC', 'NTLARS',
                'Host\s*Explorer', 'ShopFloor', 'TN3270', 'TN5250', 'HE\s*3270', 'HE\s*5250',
                'OpenText',
                'Defect[_\s-]?Tracker',
                'MarkZebra', 'Zebra',
                'PC-?DMIS',
                'Hexagon', 'CLM\s*Tools',
                'Blancco',
                'Keyence'
            )
            Target = @(
                '\\UDC\.exe$',
                'DncMain\.exe$', 'NTLARS\.exe$', 'DNCRemote',
                'HostExplorer\.exe$', 'hostex32\.exe$', 'humicon15\.exe$',
                'PC-?DMIS',
                'Hexagon',
                'Defect_Tracker\.application$'
            )
        }
        'Web Links' = @{
            Extension = @('.url')
        }
    }

    # Ensure category folders exist
    foreach ($cat in $categories.Keys) {
        $dir = Join-Path $DesktopPath $cat
        if (-not (Test-Path -LiteralPath $dir)) {
            try {
                New-Item -ItemType Directory -Path $dir -Force -ErrorAction Stop | Out-Null
            } catch {
                Write-Warning "Failed to create category folder '$cat' : $_"
            }
        }
    }

    # WScript.Shell for resolving .lnk targets
    $shell = $null
    try { $shell = New-Object -ComObject WScript.Shell } catch {}

    # Only sweep files at the desktop root, never recurse into the category
    # folders themselves (that would loop things back).
    $items = @(Get-ChildItem -LiteralPath $DesktopPath -File -ErrorAction SilentlyContinue)

    $moved   = 0
    $skipped = 0

    foreach ($item in $items) {
        # Allowlist check: never sweep items the user wants to keep visible
        if ($keepAtRoot -contains $item.Name) {
            Write-Host "  keep:  $($item.Name) (allowlisted)"
            continue
        }

        $category = $null

        # 1. Extension match (cheapest)
        foreach ($cat in $categories.Keys) {
            $exts = $categories[$cat].Extension
            if ($exts -and ($exts -contains $item.Extension)) {
                $category = $cat
                break
            }
        }

        # 2. File name regex match
        if (-not $category) {
            foreach ($cat in $categories.Keys) {
                $pats = $categories[$cat].Name
                if (-not $pats) { continue }
                foreach ($p in $pats) {
                    if ($item.BaseName -match $p) {
                        $category = $cat
                        break
                    }
                }
                if ($category) { break }
            }
        }

        # 3. .lnk target resolution
        if (-not $category -and $item.Extension -ieq '.lnk' -and $shell) {
            try {
                $sc = $shell.CreateShortcut($item.FullName)
                $target = $sc.TargetPath
                if ($target) {
                    foreach ($cat in $categories.Keys) {
                        $pats = $categories[$cat].Target
                        if (-not $pats) { continue }
                        foreach ($p in $pats) {
                            if ($target -match $p) {
                                $category = $cat
                                break
                            }
                        }
                        if ($category) { break }
                    }
                }
            } catch {}
        }

        if ($category) {
            $destDir  = Join-Path $DesktopPath $category
            $destPath = Join-Path $destDir $item.Name
            try {
                Move-Item -LiteralPath $item.FullName -Destination $destPath -Force -ErrorAction Stop
                Write-Host "  moved: $($item.Name) -> $category\"
                $moved++
            } catch {
                Write-Warning "Failed to move $($item.Name) : $_"
            }
        } else {
            $skipped++
        }
    }

    if ($shell) {
        try { [System.Runtime.InteropServices.Marshal]::ReleaseComObject($shell) | Out-Null } catch {}
    }

    Write-Host "Sweep complete: $moved moved, $skipped unclassified (left at root)."
}

# ============================================================================
# Phase 2: Materialize specific app shortcuts into Shopfloor Tools\
#
# Runs AFTER the sweep. Creates fresh .lnk files for apps we can target by
# .exe path (UDC, eDNC, NTLARS) and copies existing .lnk files for apps
# that are MSI-advertised and have no resolvable target (WJ Shopfloor,
# Defect_Tracker). Overwrites whatever the sweeper may have moved in -
# this is intentional because fresh shortcuts are more reliable than
# whatever install crud might be floating around.
#
# Each app is CONDITIONAL - if its source isn't present on this PC, the
# entry is silently skipped. Safe to run on all shopfloor PC types.
# ============================================================================

function New-ShopfloorLnk {
    param(
        [string]$Path,
        [string]$Target,
        [string]$Arguments = '',
        [string]$WorkingDirectory = ''
    )
    $wsh = New-Object -ComObject WScript.Shell
    try {
        $sc = $wsh.CreateShortcut($Path)
        $sc.TargetPath = $Target
        if ($Arguments)        { $sc.Arguments = $Arguments }
        if ($WorkingDirectory) { $sc.WorkingDirectory = $WorkingDirectory }
        elseif ($Target) {
            $parent = Split-Path -Parent $Target
            if ($parent) { $sc.WorkingDirectory = $parent }
        }
        $sc.IconLocation = "$Target,0"
        $sc.Save()
        return $true
    } catch {
        Write-Warning "Failed to create $Path : $_"
        return $false
    } finally {
        try { [System.Runtime.InteropServices.Marshal]::ReleaseComObject($wsh) | Out-Null } catch {}
    }
}

# Look up an existing .lnk file across the paths where installers / ppkgs
# typically drop them. Used for apps where we don't build the .lnk from
# an .exe (because the target is an MSI-advertised Darwin descriptor).
function Find-ExistingLnk {
    param([string]$Filename)

    $candidates = @(
        (Join-Path $publicDesktop $Filename),
        (Join-Path $shopfloorToolsDir $Filename),
        "C:\ProgramData\Microsoft\Windows\Start Menu\Programs\$Filename",
        "C:\ProgramData\Microsoft\Windows\Start Menu\Programs\Shopfloor Tools\$Filename"
    )
    foreach ($c in $candidates) {
        if (Test-Path -LiteralPath $c) { return $c }
    }
    return $null
}

function Add-ShopfloorToolsApps {
    if (-not (Test-Path -LiteralPath $shopfloorToolsDir)) {
        try {
            New-Item -ItemType Directory -Path $shopfloorToolsDir -Force -ErrorAction Stop | Out-Null
        } catch {
            Write-Warning "Failed to create $shopfloorToolsDir : $_"
            return
        }
    }

    # App registry - each entry describes how to materialize one shortcut
    # into Shopfloor Tools\.
    #
    #   Kind = 'exe'       -> build a fresh .lnk from ExePath
    #   Kind = 'existing'  -> copy an existing .lnk via Find-ExistingLnk
    $apps = @(
        @{ Name = 'UDC';            Kind = 'exe';      ExePath = 'C:\Program Files\UDC\UDC.exe' }
        @{ Name = 'eDNC';           Kind = 'exe';      ExePath = 'C:\Program Files (x86)\Dnc\bin\DncMain.exe' }
        @{ Name = 'NTLARS';         Kind = 'exe';      ExePath = 'C:\Program Files (x86)\Dnc\Common\NTLARS.exe' }
        @{ Name = 'WJ Shopfloor';   Kind = 'existing'; SourceName = 'WJ Shopfloor.lnk' }
        @{ Name = 'Defect_Tracker'; Kind = 'existing'; SourceName = 'Defect_Tracker.lnk' }
    )

    foreach ($app in $apps) {
        $dest = Join-Path $shopfloorToolsDir "$($app.Name).lnk"

        switch ($app.Kind) {
            'exe' {
                if (Test-Path -LiteralPath $app.ExePath) {
                    if (New-ShopfloorLnk -Path $dest -Target $app.ExePath) {
                        Write-Host "  created: $($app.Name) -> $dest"
                    }
                } else {
                    Write-Host "  skip:    $($app.Name) - target not installed ($($app.ExePath))" -ForegroundColor DarkGray
                }
            }

            'existing' {
                $src = Find-ExistingLnk $app.SourceName
                if ($src) {
                    try {
                        Copy-Item -LiteralPath $src -Destination $dest -Force -ErrorAction Stop
                        Write-Host "  copied:  $($app.Name) from $src"
                    } catch {
                        Write-Warning "Failed to copy $src -> $dest : $_"
                    }
                } else {
                    Write-Host "  skip:    $($app.Name) - no source .lnk found" -ForegroundColor DarkGray
                }
            }
        }
    }
}

# ============================================================================
# Phase 3: Pin eDNC and NTLARS shortcuts at the desktop root
#
# These are the two shortcuts end users touch most frequently. Allowlisted
# in the sweep, and also force-dropped here (copied from the Shopfloor
# Tools\ versions we just created) so a sweep-and-forget user still ends
# up with them loose at the root.
# ============================================================================
function Update-DesktopRootPins {
    foreach ($name in @('eDNC.lnk', 'NTLARS.lnk')) {
        $src = Join-Path $shopfloorToolsDir $name
        $dst = Join-Path $publicDesktop $name
        if (Test-Path -LiteralPath $src) {
            try {
                Copy-Item -LiteralPath $src -Destination $dst -Force -ErrorAction Stop
                Write-Host "  root:    $name"
            } catch {
                Write-Warning "Failed to drop $dst : $_"
            }
        }
    }
}

# ============================================================================
# Scheduled task registration (phase 1 re-run at every logon)
# ============================================================================
function Register-SweepScheduledTask {
    param([string]$ScriptPath)

    $taskName = 'Organize Public Desktop'

    if (-not (Test-Path -LiteralPath $ScriptPath)) {
        Write-Warning "Cannot register scheduled task - script not found at $ScriptPath"
        return
    }

    try {
        $action = New-ScheduledTaskAction `
            -Execute 'powershell.exe' `
            -Argument "-NoProfile -ExecutionPolicy Bypass -WindowStyle Hidden -File `"$ScriptPath`""

        $trigger = New-ScheduledTaskTrigger -AtLogOn

        # Run as whichever user logs in, at Limited (standard) rights. Public
        # desktop is writable by BUILTIN\Users so no elevation needed. Using
        # the well-known Users group SID so this works on non-English Windows.
        $principal = New-ScheduledTaskPrincipal `
            -GroupId 'S-1-5-32-545' `
            -RunLevel Limited

        $settings = New-ScheduledTaskSettingsSet `
            -AllowStartIfOnBatteries `
            -DontStopIfGoingOnBatteries `
            -StartWhenAvailable `
            -ExecutionTimeLimit (New-TimeSpan -Minutes 5)

        # -ErrorAction Stop is required because Register-ScheduledTask emits
        # non-terminating errors on failure (Access Denied etc), which try/catch
        # would otherwise miss and we'd falsely claim success.
        Register-ScheduledTask `
            -TaskName   $taskName `
            -Action     $action `
            -Trigger    $trigger `
            -Principal  $principal `
            -Settings   $settings `
            -Force `
            -ErrorAction Stop | Out-Null

        Write-Host "Scheduled task '$taskName' registered (runs at every logon)."
    } catch {
        Write-Warning "Failed to register scheduled task : $_"
    }
}

# ============================================================================
# Main
# ============================================================================
Write-Host ""
Write-Host "=== Phase 1: sweep loose shortcuts into category folders ==="
Invoke-DesktopSweep -DesktopPath $publicDesktop

Write-Host ""
Write-Host "=== Phase 2: populate Shopfloor Tools with app shortcuts ==="
Add-ShopfloorToolsApps

Write-Host ""
Write-Host "=== Phase 3: drop eDNC / NTLARS at desktop root ==="
Update-DesktopRootPins

Write-Host ""
Write-Host "=== Phase 4: register logon sweeper scheduled task ==="
Register-SweepScheduledTask -ScriptPath $scriptPath

exit 0