# --- Load Required Assemblies for UI ---
Add-Type -AssemblyName System.Windows.Forms
Add-Type -AssemblyName System.Drawing

#region --- UI Functions ---
# --- Function to Show Desktop Naming Form ---
function Show-DesktopNamingForm {
    param (
        [System.Collections.Generic.List[string]]$BuildingCodes,
        [string]$PrefillBuilding,
        [string]$PrefillRoom,
        [string]$PrefillAssetTag
    )

    $form = New-Object System.Windows.Forms.Form
    $form.Text = 'Desktop Naming'
    $form.Size = New-Object System.Drawing.Size(320, 240)
    $form.StartPosition = 'CenterScreen'
    $form.FormBorderStyle = 'FixedDialog'
    $form.TopMost = $true
    $form.MaximizeBox = $false
    $form.MinimizeBox = $false

    $labelBuilding = New-Object System.Windows.Forms.Label
    $labelBuilding.Location = New-Object System.Drawing.Point(20, 20)
    $labelBuilding.Size = New-Object System.Drawing.Size(100, 20)
    $labelBuilding.Text = 'Building Code:'
    $form.Controls.Add($labelBuilding)

    $comboBuilding = New-Object System.Windows.Forms.ComboBox
    $comboBuilding.Location = New-Object System.Drawing.Point(130, 18)
    $comboBuilding.Size = New-Object System.Drawing.Size(150, 20)
    $comboBuilding.DropDownStyle = 'DropDownList'
    $BuildingCodes.ForEach({ [void]$comboBuilding.Items.Add($_) })
    $comboBuilding.SelectedItem = $PrefillBuilding
    $form.Controls.Add($comboBuilding)

    $labelRoom = New-Object System.Windows.Forms.Label
    $labelRoom.Location = New-Object System.Drawing.Point(20, 60)
    $labelRoom.Size = New-Object System.Drawing.Size(100, 20)
    $labelRoom.Text = 'Room Number:'
    $form.Controls.Add($labelRoom)

    $textRoom = New-Object System.Windows.Forms.TextBox
    $textRoom.Location = New-Object System.Drawing.Point(130, 58)
    $textRoom.Size = New-Object System.Drawing.Size(150, 20)
    $textRoom.Text = $PrefillRoom
    $form.Controls.Add($textRoom)

    $labelAssetTag = New-Object System.Windows.Forms.Label
    $labelAssetTag.Location = New-Object System.Drawing.Point(20, 100)
    $labelAssetTag.Size = New-Object System.Drawing.Size(100, 20)
    $labelAssetTag.Text = 'Asset Tag:'
    $form.Controls.Add($labelAssetTag)

    $textAssetTag = New-Object System.Windows.Forms.TextBox
    $textAssetTag.Location = New-Object System.Drawing.Point(130, 98)
    $textAssetTag.Size = New-Object System.Drawing.Size(150, 20)
    $textAssetTag.Text = $PrefillAssetTag
    $form.Controls.Add($textAssetTag)

    $buttonOK = New-Object System.Windows.Forms.Button
    $buttonOK.Location = New-Object System.Drawing.Point(105, 150) # Centered the OK button
    $buttonOK.Size = New-Object System.Drawing.Size(90, 30)
    $buttonOK.Text = 'OK'
    $form.Controls.Add($buttonOK)

    # Event handler for OK button to perform validation
    $buttonOK.Add_Click({
        if ([string]::IsNullOrWhiteSpace($comboBuilding.SelectedItem) -or [string]::IsNullOrWhiteSpace($textRoom.Text) -or [string]::IsNullOrWhiteSpace($textAssetTag.Text)) {
            [System.Windows.Forms.MessageBox]::Show('Building Code, Room Number, and Asset Tag cannot be empty.', 'Validation Error', 'OK', 'Error')
            $form.DialogResult = [System.Windows.Forms.DialogResult]::None # Prevent form from closing
        } else {
            $form.DialogResult = [System.Windows.Forms.DialogResult]::OK
        }
    })

    # Loop until valid input is provided
    do {
        $dialogResult = $form.ShowDialog()
    } while ($dialogResult -eq [System.Windows.Forms.DialogResult]::None) # Keep showing if validation failed

    if ($dialogResult -eq [System.Windows.Forms.DialogResult]::OK) {
        return [pscustomobject]@{
            BuildingCode = $comboBuilding.SelectedItem
            RoomNumber   = $textRoom.Text
            AssetTag     = $textAssetTag.Text
        }
    }
    # This block should technically not be reached if OK requires valid input
    return $null
}

# --- Function to Show Laptop Asset Tag Prompt ---
function Show-LaptopAssetTagPrompt {
    param(
        [string]$PrefillAssetTag = ""
    )

    $form = New-Object System.Windows.Forms.Form
    $form.Text = 'Asset Tag Required'
    $form.Size = New-Object System.Drawing.Size(320, 180)
    $form.StartPosition = 'CenterScreen'
    $form.FormBorderStyle = 'FixedDialog'
    $form.TopMost = $true
    $form.MaximizeBox = $false
    $form.MinimizeBox = $false

    $labelInfo = New-Object System.Windows.Forms.Label
    $labelInfo.Location = New-Object System.Drawing.Point(20, 20)
    $labelInfo.Size = New-Object System.Drawing.Size(280, 40)
    $labelInfo.Text = "No asset tag found please enter the device asset tag"
    $form.Controls.Add($labelInfo)

    $textAssetTag = New-Object System.Windows.Forms.TextBox
    $textAssetTag.Location = New-Object System.Drawing.Point(20, 70)
    $textAssetTag.Size = New-Object System.Drawing.Size(260, 20)
    $textAssetTag.Text = $PrefillAssetTag
    $form.Controls.Add($textAssetTag)

    $buttonOK = New-Object System.Windows.Forms.Button
    $buttonOK.Location = New-Object System.Drawing.Point(105, 110)
    $buttonOK.Size = New-Object System.Drawing.Size(90, 30)
    $buttonOK.Text = 'OK'
    $form.Controls.Add($buttonOK)

    # Event handler for OK button to perform validation
    $buttonOK.Add_Click({
        if ([string]::IsNullOrWhiteSpace($textAssetTag.Text)) {
            [System.Windows.Forms.MessageBox]::Show('Asset Tag cannot be empty.', 'Validation Error', 'OK', 'Error')
            $form.DialogResult = [System.Windows.Forms.DialogResult]::None # Prevent form from closing
        } else {
            $form.DialogResult = [System.Windows.Forms.DialogResult]::OK
        }
    })

    $form.ActiveControl = $textAssetTag

    # Loop until valid input is provided
    do {
        $dialogResult = $form.ShowDialog()
    } while ($dialogResult -eq [System.Windows.Forms.DialogResult]::None) # Keep showing if validation failed

    if ($dialogResult -eq [System.Windows.Forms.DialogResult]::OK) {
        return $textAssetTag.Text
    }
    # This block should technically not be reached if OK requires valid input
    return $null
}
#endregion

# --- Helper Function to Update AD Description ---
function Update-ADComputerDescription {
    param (
        [Parameter(Mandatory=$true)]
        [ADComputer]$ADComputerObject,
        [Parameter(Mandatory=$true)]
        [string]$NewDescription
    )
    try {
        Set-ADComputer -Identity $ADComputerObject -Description $NewDescription -ErrorAction Stop
        Write-Output "Successfully updated AD object '$($ADComputerObject.Name)' description to '$NewDescription'."
    } catch {
        Write-Warning "Failed to update AD object description for '$($ADComputerObject.Name)'. Error: $($_.Exception.Message)"
    }
}

---
## Main Script Logic
---

try {
    # 1. Initialize SCCM Task Sequence Environment
    $tsenv = New-Object -COMObject Microsoft.SMS.TSEnvironment
    Write-Output "Successfully connected to Task Sequence environment."

    # 2. Gather Hardware Information
    $serialNumber = (Get-WmiObject -Class Win32_BIOS).SerialNumber.Trim()
    $chassisTypes = (Get-WmiObject -Class Win32_SystemEnclosure).ChassisTypes
    $smbiosGuid = (Get-WmiObject -Class Win32_ComputerSystemProduct).UUID
    $laptopChassisTypes = @(8, 9, 10, 11, 12, 14, 18, 21, 31, 32)
    $isLaptop = $chassisTypes | Where-Object { $laptopChassisTypes -contains $_ }

    Write-Output "Serial Number: $serialNumber"
    Write-Output "Machine Type: $(if($isLaptop){'Laptop'} else {'Desktop'})"

    # 3. Find existing AD Computer Object
    $adComputer = $null
    $existingAssetTag = $null
    try {
        # This requires the AD module to be installed on the environment where the script runs
        Import-Module ActiveDirectory -ErrorAction SilentlyContinue
        # Ensure 'description' is always retrieved
        $adComputer = Get-ADComputer -Filter "servicePrincipalName -like '*$smbiosGuid*'" -Properties description, Name -ErrorAction SilentlyContinue
        if ($adComputer) {
            Write-Output "Found existing AD computer object: $($adComputer.Name)"
            $existingAssetTag = $adComputer.description
            if ([string]::IsNullOrWhiteSpace($existingAssetTag)) {
                Write-Output "Existing AD object has no Asset Tag in its description."
            } else {
                Write-Output "Existing Asset Tag in AD: $existingAssetTag"
            }
        }
        else { Write-Output "No existing AD object found for this machine." }
    }
    catch {
        Write-Warning "Could not query Active Directory. Error: $($_.Exception.Message)"
    }

    ---
    ## Laptop Logic
    ---
    if ($isLaptop) {
        $computerName = "LT-$($serialNumber)"
        Write-Output "Determined Laptop Computer Name: $computerName"

        $assetTagToStore = $existingAssetTag
        # Only prompt if existing asset tag is null or empty
        if ([string]::IsNullOrWhiteSpace($existingAssetTag)) {
            Write-Output "AD object exists but description (Asset Tag) is missing or empty. Prompting user."
            $promptedAssetTag = Show-LaptopAssetTagPrompt -PrefillAssetTag $existingAssetTag

            # Since the prompt now enforces input, $promptedAssetTag will always contain a value if the function returns.
            $assetTagToStore = $promptedAssetTag
        }

        # Update AD description if an AD object was found and an asset tag is determined
        # This condition also ensures we don't write back the same value if it was already there
        # and not changed by the user.
        if ($adComputer -and ($assetTagToStore -ne $existingAssetTag)) {
            Update-ADComputerDescription -ADComputerObject $adComputer -NewDescription $assetTagToStore
        }

        $tsenv.Value('OSDComputerName') = $computerName
        Write-Output "Set OSDComputerName to '$computerName' in Task Sequence environment."
    }
    ---
    ## Desktop Logic
    ---
    else {
        $buildingCodeFile = "\\scssccm2\Sources\Script_Sources\BuildingCodeList.txt"
        if (-not (Test-Path $buildingCodeFile)) {
            Write-Error "Cannot find Building Code list at '$buildingCodeFile'. Cannot continue."
            exit 1 # Exit task sequence if building code file is missing
        }
        $buildingCodes = New-Object System.Collections.Generic.List[string]
        Get-Content $buildingCodeFile | Where-Object { -not [string]::IsNullOrWhiteSpace($_) } | ForEach-Object { $buildingCodes.Add($_.Trim()) }

        $prefill = @{ Building = ''; Room = ''; AssetTag = $existingAssetTag } # Pre-fill AssetTag from AD
        if ($adComputer) {
            # Try to pre-fill building and room from existing AD computer name
            $nameParts = $adComputer.Name.Split('-')
            if ($nameParts.Count -ge 2) {
                if ($buildingCodes.Contains($nameParts[0])) { $prefill.Building = $nameParts[0] }
                $prefill.Room = $nameParts[1]
            }
        }

        $userInput = Show-DesktopNamingForm -BuildingCodes $buildingCodes -PrefillBuilding $prefill.Building -PrefillRoom $prefill.Room -PrefillAssetTag $prefill.AssetTag

        # Since desktop naming form now enforces input, $userInput will never be $null here
        $last5OfSerial = if ($serialNumber.Length -gt 5) { $serialNumber.Substring($serialNumber.Length - 5) } else { $serialNumber }
        $computerName = "$($userInput.BuildingCode)-$($userInput.RoomNumber)-$($last5OfSerial)"

        $tsenv.Value('OSDComputerName') = $computerName
        Write-Output "Set OSDComputerName to '$computerName' in Task Sequence environment."

        # Update AD description with the entered Asset Tag
        if ($adComputer -and ($userInput.AssetTag -ne $existingAssetTag)) {
            Update-ADComputerDescription -ADComputerObject $adComputer -NewDescription $userInput.AssetTag
        }
    }
}
catch {
    Write-Error "A critical error occurred in the naming script: $($_.Exception.Message)"
    # Exit with a non-zero code to fail the Task Sequence step
    exit 1
}