#requires -Version 5.1
<#
.SYNOPSIS
    Cross-platform bulk assignment manager for Intune policies and apps.
.DESCRIPTION
    Add or remove assignments across multiple Intune object types
    (Device Configuration, Compliance, Settings Catalog, Apps, Scripts, etc.)
    in a single operation. Uses fzf when available; falls back to numbered menus.
    Integrates with the IntuneManagement headless auth stack.
.EXAMPLE
    ./Scripts/Bulk-AssignmentManager.ps1 -TenantId "contoso.onmicrosoft.com" -AppId "..." -Secret "..."
#>
[CmdletBinding()]
param(
    [Parameter(Mandatory = $true)]
    [string]$TenantId,

    [string]$AppId,

    [string]$Secret,

    [string]$Certificate,

    [ValidateSet("AppOnly","Browser","DeviceCode")]
    [string]$AuthMode = "AppOnly",

    [string]$RedirectUri,

    [string]$SettingsFile
)

$ErrorActionPreference = "Stop"

#region Helper functions
function Test-FzfAvailable
{
    return [bool](Get-Command fzf -ErrorAction SilentlyContinue)
}

function Show-FzfMenu
{
    param(
        [Parameter(Mandatory)]
        [string[]]$Items,
        [string]$Header = "Select one",
        [switch]$Multi
    )
    $argsList = @("--header=$Header")
    if($Multi) { $argsList += "--multi" }
    $selected = $Items | fzf @argsList
    if(-not $selected) { return $null }
    if($Multi) { return @($selected -split "`r?`n" | Where-Object { $_ }) }
    return $selected
}

function Show-NumberedMenu
{
    param(
        [Parameter(Mandatory)]
        [string[]]$Items,
        [string]$Header = "Select one or more",
        [switch]$Multi
    )
    Write-Host "`n$Header" -ForegroundColor Cyan
    for($i=0; $i -lt $Items.Count; $i++)
    {
        Write-Host "  $($i+1). $($Items[$i])"
    }
    if($Multi)
    {
        $prompt = "Enter numbers separated by commas (e.g. 1,3,5) or 'all'"
    }
    else
    {
        $prompt = "Enter a number"
    }
    $choice = Read-Host $prompt
    if($choice -eq "all" -and $Multi) { return $Items }
    $indices = $choice -split "," | ForEach-Object { $_.Trim() } | Where-Object { $_ -match "^\d+$" } | ForEach-Object { [int]$_ - 1 } | Where-Object { $_ -ge 0 -and $_ -lt $Items.Count }
    if($Multi)
    {
        return $Items[$indices] | Select-Object -Unique
    }
    else
    {
        if($indices.Count -eq 0) { return $null }
        return $Items[$indices[0]]
    }
}

function Select-MenuItem
{
    param(
        [Parameter(Mandatory)]
        [string[]]$Items,
        [string]$Header = "Select one",
        [switch]$Multi
    )
    if(Test-FzfAvailable)
    {
        return Show-FzfMenu -Items $Items -Header $Header -Multi:$Multi
    }
    return Show-NumberedMenu -Items $Items -Header $Header -Multi:$Multi
}

function Read-YesNo
{
    param(
        [string]$Prompt,
        [bool]$Default = $false
    )
    $defaultChar = if($Default) { "Y" } else { "N" }
    $response = Read-Host "$Prompt [Y/n] (default: $defaultChar)"
    if([string]::IsNullOrWhiteSpace($response)) { return $Default }
    return $response -match "^\s*y"
}

function Get-DefaultSettingsPath
{
    if($IsWindows -or $env:OS -eq "Windows_NT")
    {
        if($env:LOCALAPPDATA) { return (Join-Path $env:LOCALAPPDATA "macOS_IntuneManagement\Settings.json") }
        return (Join-Path $env:USERPROFILE "AppData\Local\macOS_IntuneManagement\Settings.json")
    }
    if($IsMacOS) { return (Join-Path $HOME "Library/Application Support/macOS_IntuneManagement/Settings.json") }
    return (Join-Path $HOME ".local/share/macOS_IntuneManagement/Settings.json")
}
#endregion

#region Initialize Runtime
$projectRoot = Split-Path -Parent $PSScriptRoot
$runtimeModule = Join-Path $projectRoot "Runtime/IntuneManagement.Runtime.psd1"
if(-not (Test-Path $runtimeModule))
{
    throw "Could not find IntuneManagement.Runtime.psd1 in $projectRoot"
}

$settingsPath = $SettingsFile
if(-not $settingsPath)
{
    $settingsPath = Get-DefaultSettingsPath
}

# Pre-load auth from settings
if($AuthMode -eq "AppOnly" -and (Test-Path $settingsPath) -and (-not $AppId -or (-not $Secret -and -not $Certificate)))
{
    try
    {
        $raw = Get-Content -Path $settingsPath -Raw -ErrorAction Stop
        $settingsObj = ConvertFrom-Json $raw -AsHashtable -ErrorAction Stop
        if($settingsObj -and $settingsObj.ContainsKey($TenantId))
        {
            $tenantNode = $settingsObj[$TenantId]
            if(-not $AppId -and $tenantNode.ContainsKey("GraphAzureAppId"))
            {
                $AppId = $tenantNode["GraphAzureAppId"]
            }
            if(-not $Secret -and $tenantNode.ContainsKey("GraphAzureAppSecret"))
            {
                $Secret = $tenantNode["GraphAzureAppSecret"]
            }
            if(-not $Certificate -and $tenantNode.ContainsKey("GraphAzureAppCert"))
            {
                $Certificate = $tenantNode["GraphAzureAppCert"]
            }
        }

        if(-not $Secret -and $IsMacOS -and $AppId)
        {
            try
            {
                $keychainSecret = security find-generic-password -a "IntuneManagement" -s "IntuneMgmt-$AppId" -w 2>$null
                if($keychainSecret) { $Secret = $keychainSecret }
            }
            catch { }
        }
    }
    catch { }
}

$invokeParams = @{
    Silent = $true
    JSonSettings = $true
    JSonFile = $settingsPath
    TenantId = $TenantId
    AppId = $AppId
    AuthMode = $AuthMode
}
if($RedirectUri) { $invokeParams.RedirectUri = $RedirectUri }
if($AuthMode -eq "AppOnly" -and $Secret) { $invokeParams.Secret = $Secret }
elseif($AuthMode -eq "AppOnly") { $invokeParams.Certificate = $Certificate }

Import-Module $runtimeModule -Force
Initialize-IntuneManagementRuntime -View "IntuneGraphAPI" @invokeParams
#endregion

#region Ensure Graph connectivity
if(-not (Get-Command Invoke-GraphRequest -ErrorAction SilentlyContinue))
{
    throw "Graph runtime did not load Invoke-GraphRequest. Aborting."
}

Write-Host "`nConnecting to Microsoft Graph..." -ForegroundColor Cyan
try
{
    $org = Invoke-GraphRequest "/organization"
    Write-Host "Connected to tenant: $($org.value[0].displayName) ($($org.value[0].id))" -ForegroundColor Green
}
catch
{
    throw "Failed to connect to Graph. Ensure auth parameters are correct. Error: $_"
}
#endregion

#region Object type registry (assignable types)
$assignableTypes = @(
    [PSCustomObject]@{ Title = "Applications"; API = "/deviceAppManagement/mobileApps"; AssignmentsType = "mobileAppAssignments"; AssignmentODataType = "#microsoft.graph.mobileAppAssignment"; HasIntent = $true; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "Device Configuration"; API = "/deviceManagement/deviceConfigurations"; AssignmentsType = "assignments"; AssignmentODataType = "#microsoft.graph.deviceConfigurationAssignment"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "Settings Catalog"; API = "/deviceManagement/configurationPolicies"; AssignmentsType = "assignments"; AssignmentODataType = "#microsoft.graph.deviceManagementConfigurationPolicyAssignment"; HasIntent = $false; NameProp = "name" },
    [PSCustomObject]@{ Title = "Compliance Policies"; API = "/deviceManagement/deviceCompliancePolicies"; AssignmentsType = "assignments"; AssignmentODataType = "#microsoft.graph.deviceCompliancePolicyAssignment"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "Administrative Templates"; API = "/deviceManagement/groupPolicyConfigurations"; AssignmentsType = "assignments"; AssignmentODataType = "#microsoft.graph.groupPolicyConfigurationAssignment"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "Endpoint Security"; API = "/deviceManagement/intents"; AssignmentsType = "assignments"; AssignmentODataType = "#microsoft.graph.deviceManagementIntentAssignment"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "App Protection"; API = "/deviceAppManagement/managedAppPolicies"; AssignmentsType = "assignments"; AssignmentODataType = "#microsoft.graph.targetedManagedAppPolicyAssignment"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "App Configuration (Device)"; API = "/deviceAppManagement/mobileAppConfigurations"; AssignmentsType = "assignments"; AssignmentODataType = "#microsoft.graph.managedDeviceMobileAppConfigurationAssignment"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "Platform Scripts"; API = "/deviceManagement/deviceManagementScripts"; AssignmentsType = "deviceManagementScriptAssignments"; AssignmentODataType = "#microsoft.graph.deviceManagementScriptAssignment"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "macOS Scripts"; API = "/deviceManagement/deviceShellScripts"; AssignmentsType = "deviceManagementScriptAssignments"; AssignmentODataType = "#microsoft.graph.deviceManagementScriptAssignment"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "Device Health Scripts"; API = "/deviceManagement/deviceHealthScripts"; AssignmentsType = "deviceHealthScriptAssignments"; AssignmentODataType = "#microsoft.graph.deviceHealthScriptAssignment"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "macOS Custom Attributes"; API = "/deviceManagement/deviceCustomAttributeShellScripts"; AssignmentsType = "deviceManagementScriptAssignments"; AssignmentODataType = "#microsoft.graph.deviceManagementScriptAssignment"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "Enrollment Restrictions"; API = "/deviceManagement/deviceEnrollmentConfigurations"; AssignmentsType = "enrollmentConfigurationAssignments"; AssignmentODataType = "#microsoft.graph.enrollmentConfigurationAssignment"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "Enrollment Status Page"; API = "/deviceManagement/deviceEnrollmentConfigurations"; AssignmentsType = "enrollmentConfigurationAssignments"; AssignmentODataType = "#microsoft.graph.enrollmentConfigurationAssignment"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "Autopilot"; API = "/deviceManagement/windowsAutopilotDeploymentProfiles"; AssignmentsType = "assignments"; AssignmentODataType = "#microsoft.graph.windowsAutopilotDeploymentProfileAssignment"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "Terms and Conditions"; API = "/deviceManagement/termsAndConditions"; AssignmentsType = "assignments"; AssignmentODataType = "#microsoft.graph.termsAndConditionsAssignment"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "Policy Sets"; API = "/deviceAppManagement/policySets"; AssignmentsType = "assignments"; AssignmentODataType = "#microsoft.graph.policySetAssignment"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "Update Policies"; API = "/deviceManagement/windowsUpdateForBusinessConfigurations"; AssignmentsType = "assignments"; AssignmentODataType = "#microsoft.graph.windowsUpdateForBusinessConfigurationAssignment"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "Feature Updates"; API = "/deviceManagement/windowsFeatureUpdateProfiles"; AssignmentsType = "assignments"; AssignmentODataType = "#microsoft.graph.windowsFeatureUpdateProfileAssignment"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "Quality Updates"; API = "/deviceManagement/windowsQualityUpdateProfiles"; AssignmentsType = "assignments"; AssignmentODataType = "#microsoft.graph.windowsQualityUpdateProfileAssignment"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "Device Management Intents"; API = "/deviceManagement/intents"; AssignmentsType = "assignments"; AssignmentODataType = "#microsoft.graph.deviceManagementIntentAssignment"; HasIntent = $false; NameProp = "displayName" }
)
#endregion

#region Action selection
Clear-Host
Write-Host "========================================" -ForegroundColor Cyan
Write-Host "  Intune Bulk Assignment Manager" -ForegroundColor Cyan
Write-Host "========================================" -ForegroundColor Cyan

$action = Select-MenuItem -Items @("Add assignments","Remove assignments") -Header "Select action"
if(-not $action) { Write-Host "Cancelled." -ForegroundColor Yellow; exit 0 }
#endregion

#region Select object type
$typeTitles = $assignableTypes | ForEach-Object { $_.Title }
$selectedTypeTitle = Select-MenuItem -Items $typeTitles -Header "Select object type"
if(-not $selectedTypeTitle) { Write-Host "Cancelled." -ForegroundColor Yellow; exit 0 }
$objectType = $assignableTypes | Where-Object { $_.Title -eq $selectedTypeTitle } | Select-Object -First 1
#endregion

#region Load objects
Write-Host "`nLoading $($objectType.Title) objects..." -ForegroundColor Cyan
$api = "$($objectType.API)?`$select=id,$($objectType.NameProp)&`$orderby=$($objectType.NameProp)"
$objectsResponse = Invoke-GraphRequest $api -AllPages
$objects = $objectsResponse.value | Where-Object { $_ } | Sort-Object $objectType.NameProp
Write-Host "Found $($objects.Count) objects." -ForegroundColor Green

$filter = Read-Host "`nFilter by name (optional, press Enter to skip)"
if(-not [string]::IsNullOrWhiteSpace($filter))
{
    $objects = $objects | Where-Object { $_."$($objectType.NameProp)" -like "*$filter*" }
    Write-Host "Filtered to $($objects.Count) objects." -ForegroundColor Green
}

if($objects.Count -eq 0)
{
    Write-Host "No objects found. Exiting." -ForegroundColor Yellow
    exit 0
}

$objectDisplays = $objects | ForEach-Object { "$($_."$($objectType.NameProp)") [$($_.id)]" }
$selectedDisplays = Select-MenuItem -Items $objectDisplays -Header "Select objects (multi-select)" -Multi
if(-not $selectedDisplays)
{
    Write-Host "No objects selected. Exiting." -ForegroundColor Yellow
    exit 0
}

$selectedObjects = @()
foreach($disp in $selectedDisplays)
{
    $id = $disp -replace '.*\[(.*?)\]$', '$1'
    $obj = $objects | Where-Object { $_.id -eq $id } | Select-Object -First 1
    if($obj) { $selectedObjects += $obj }
}
Write-Host "Selected $($selectedObjects.Count) objects." -ForegroundColor Green
#endregion

#region Load groups & filters
Write-Host "`nLoading Azure AD groups..." -ForegroundColor Cyan
$groupsResponse = Invoke-GraphRequest "/groups?`$select=id,displayName&`$orderby=displayName" -AllPages
$groups = $groupsResponse.value | Where-Object { $_.displayName } | Sort-Object displayName
Write-Host "Found $($groups.Count) groups." -ForegroundColor Green

Write-Host "`nLoading assignment filters..." -ForegroundColor Cyan
$filtersResponse = Invoke-GraphRequest "/deviceManagement/assignmentFilters?`$select=id,displayName&`$orderby=displayName"
$assignmentFilters = $filtersResponse.value | Where-Object { $_.displayName } | Sort-Object displayName
Write-Host "Found $($assignmentFilters.Count) filters." -ForegroundColor Green
#endregion

#region Add assignments flow
if($action -eq "Add assignments")
{
    $groupDisplays = $groups | ForEach-Object { "$($_.displayName) [$($_.id)]" }
    $selectedGroupDisplays = Select-MenuItem -Items $groupDisplays -Header "Select target groups (multi-select)" -Multi
    $selectedGroups = @()
    if($selectedGroupDisplays)
    {
        foreach($disp in $selectedGroupDisplays)
        {
            $id = $disp -replace '.*\[(.*?)\]$', '$1'
            $grp = $groups | Where-Object { $_.id -eq $id } | Select-Object -First 1
            if($grp) { $selectedGroups += $grp }
        }
    }

    $allUsers = Read-YesNo -Prompt "Target All Users?" -Default $false
    $allDevices = Read-YesNo -Prompt "Target All Devices?" -Default $false

    if(($selectedGroups.Count -eq 0) -and -not $allUsers -and -not $allDevices)
    {
        Write-Host "No targets selected. Exiting." -ForegroundColor Yellow
        exit 0
    }

    $intent = $null
    if($objectType.HasIntent)
    {
        $intent = Select-MenuItem -Items @("required","available","uninstall") -Header "Select assignment intent"
        if(-not $intent) { $intent = "required" }
        if($intent -eq "available")
        {
            Write-Host "Note: All Devices cannot be targeted with Available intent." -ForegroundColor DarkGray
            $allDevices = $false
        }
    }

    $includeExclude = "include"
    if($selectedGroups.Count -gt 0)
    {
        $includeExclude = Select-MenuItem -Items @("include","exclude") -Header "Group target mode"
        if(-not $includeExclude) { $includeExclude = "include" }
    }

    $filterDisplay = "(none)"
    if($assignmentFilters.Count -gt 0)
    {
        $filterDisplays = @("(none)") + ($assignmentFilters | ForEach-Object { "$($_.displayName) [$($_.id)]" })
        $filterSelection = Select-MenuItem -Items $filterDisplays -Header "Select assignment filter (optional)"
        if($filterSelection -and $filterSelection -ne "(none)")
        {
            $filterId = $filterSelection -replace '.*\[(.*?)\]$', '$1'
            $filterObj = $assignmentFilters | Where-Object { $_.id -eq $filterId } | Select-Object -First 1
            if($filterObj) { $filterDisplay = $filterObj.displayName }
        }
    }

    # Review
    Clear-Host
    Write-Host "Review add-assignment operation:" -ForegroundColor Green
    Write-Host "  Object Type : $($objectType.Title)"
    Write-Host "  Objects     : $($selectedObjects.Count)"
    Write-Host "  Groups      : $($selectedGroups.Count)"
    Write-Host "  All Users   : $allUsers"
    Write-Host "  All Devices : $allDevices"
    if($intent) { Write-Host "  Intent      : $intent" }
    Write-Host "  Mode        : $includeExclude"
    Write-Host "  Filter      : $filterDisplay"
    $confirm = Read-Host "`nProceed? [Y/n]"
    if(-not ([string]::IsNullOrWhiteSpace($confirm) -or $confirm -match "^\s*y"))
    {
        Write-Host "Cancelled." -ForegroundColor Yellow
        exit 0
    }

    # Execute
    $success = 0
    $skipped = 0
    $failed = 0

    foreach($obj in $selectedObjects)
    {
        Write-Host "`nProcessing: $($obj."$($objectType.NameProp)")" -ForegroundColor Cyan

        try
        {
            $existing = Invoke-GraphRequest "$($objectType.API)/$($obj.id)/assignments"
            $existingTargets = $existing.value
        }
        catch
        {
            Write-Host "  ERROR: Could not load existing assignments" -ForegroundColor Red
            $failed++
            continue
        }

        function Test-AssignmentExists
        {
            param($targetType, $groupId)
            foreach($ea in $existingTargets)
            {
                $t = $ea.target
                if($targetType -eq "group" -and $t."@odata.type" -eq "#microsoft.graph.groupAssignmentTarget" -and $t.groupId -eq $groupId) { return $true }
                if($targetType -eq "allUsers" -and $t."@odata.type" -eq "#microsoft.graph.allLicensedUsersAssignmentTarget") { return $true }
                if($targetType -eq "allDevices" -and $t."@odata.type" -eq "#microsoft.graph.allDevicesAssignmentTarget") { return $true }
                if($targetType -eq "excludeGroup" -and $t."@odata.type" -eq "#microsoft.graph.exclusionGroupAssignmentTarget" -and $t.groupId -eq $groupId) { return $true }
            }
            return $false
        }

        $payloads = @()

        foreach($grp in $selectedGroups)
        {
            $targetTypeName = if($includeExclude -eq "exclude") { "excludeGroup" } else { "group" }
            $odataType = if($includeExclude -eq "exclude") { "#microsoft.graph.exclusionGroupAssignmentTarget" } else { "#microsoft.graph.groupAssignmentTarget" }

            if(Test-AssignmentExists -targetType $targetTypeName -groupId $grp.id)
            {
                Write-Host "  SKIP: $($grp.displayName) ($includeExclude) already assigned" -ForegroundColor DarkYellow
                $skipped++
                continue
            }

            $targetPayload = @{
                "@odata.type" = $odataType
                groupId = $grp.id
            }
            if($filterObj)
            {
                $targetPayload["deviceAndAppManagementAssignmentFilterId"] = $filterObj.id
                $targetPayload["deviceAndAppManagementAssignmentFilterType"] = "include"
            }

            $assignmentPayload = @{
                "@odata.type" = $objectType.AssignmentODataType
                target = $targetPayload
            }
            if($objectType.HasIntent -and $intent)
            {
                $assignmentPayload.intent = $intent
            }

            $payloads += $assignmentPayload
            Write-Host "  QUEUE: Group -> $($grp.displayName) ($includeExclude)" -ForegroundColor Gray
        }

        if($allUsers)
        {
            if(Test-AssignmentExists -targetType "allUsers")
            {
                Write-Host "  SKIP: All Users already assigned" -ForegroundColor DarkYellow
                $skipped++
            }
            else
            {
                $targetPayload = @{
                    "@odata.type" = "#microsoft.graph.allLicensedUsersAssignmentTarget"
                }
                if($filterObj)
                {
                    $targetPayload["deviceAndAppManagementAssignmentFilterId"] = $filterObj.id
                    $targetPayload["deviceAndAppManagementAssignmentFilterType"] = "include"
                }
                $assignmentPayload = @{
                    "@odata.type" = $objectType.AssignmentODataType
                    target = $targetPayload
                }
                if($objectType.HasIntent -and $intent)
                {
                    $assignmentPayload.intent = $intent
                }
                $payloads += $assignmentPayload
                Write-Host "  QUEUE: All Users" -ForegroundColor Gray
            }
        }

        if($allDevices)
        {
            if(Test-AssignmentExists -targetType "allDevices")
            {
                Write-Host "  SKIP: All Devices already assigned" -ForegroundColor DarkYellow
                $skipped++
            }
            else
            {
                $targetPayload = @{
                    "@odata.type" = "#microsoft.graph.allDevicesAssignmentTarget"
                }
                if($filterObj)
                {
                    $targetPayload["deviceAndAppManagementAssignmentFilterId"] = $filterObj.id
                    $targetPayload["deviceAndAppManagementAssignmentFilterType"] = "include"
                }
                $assignmentPayload = @{
                    "@odata.type" = $objectType.AssignmentODataType
                    target = $targetPayload
                }
                if($objectType.HasIntent -and $intent)
                {
                    $assignmentPayload.intent = $intent
                }
                $payloads += $assignmentPayload
                Write-Host "  QUEUE: All Devices" -ForegroundColor Gray
            }
        }

        if($payloads.Count -eq 0)
        {
            continue
        }

        # Merge existing + new assignments and POST to /assign (the standard Intune bulk endpoint)
        try
        {
            $allAssignments = @()

            # Clean existing assignments (remove id/source, preserve structure)
            foreach($ea in $existingTargets)
            {
                $clean = $ea | ConvertTo-Json -Depth 50 | ConvertFrom-Json
                if($clean.PSObject.Properties["id"]) { $clean.PSObject.Properties.Remove("id") }
                if($clean.PSObject.Properties["source"]) { $clean.PSObject.Properties.Remove("source") }
                if(-not $clean."@odata.type")
                {
                    $clean | Add-Member -NotePropertyName "@odata.type" -NotePropertyValue $objectType.AssignmentODataType -Force
                }
                $allAssignments += $clean
            }

            foreach($p in $payloads)
            {
                $allAssignments += $p
            }

            $assignPayload = @{
                $objectType.AssignmentsType = $allAssignments
            } | ConvertTo-Json -Depth 50 -Compress

            $null = Invoke-GraphRequest "$($objectType.API)/$($obj.id)/assign" -HttpMethod POST -Content $assignPayload
            Write-Host "  OK: Assigned $($payloads.Count) new target(s)" -ForegroundColor Green
            $success += $payloads.Count
        }
        catch
        {
            Write-Host "  ERROR: Failed to assign. $($_.Exception.Message)" -ForegroundColor Red
            $failed += $payloads.Count
        }
    }

    Write-Host "`n========================================" -ForegroundColor Cyan
    Write-Host "  Add Assignments Complete" -ForegroundColor Cyan
    Write-Host "========================================" -ForegroundColor Cyan
    Write-Host "  Success : $success"
    Write-Host "  Skipped : $skipped"
    Write-Host "  Failed  : $failed"
}
#endregion

#region Remove assignments flow
elseif($action -eq "Remove assignments")
{
    # Gather all existing assignments across selected objects
    Write-Host "`nLoading existing assignments..." -ForegroundColor Cyan
    $allAssignments = @()
    foreach($obj in $selectedObjects)
    {
        try
        {
            $existing = Invoke-GraphRequest "$($objectType.API)/$($obj.id)/assignments"
            foreach($ass in $existing.value)
            {
                $targetDesc = "Unknown"
                $targetType = $ass.target."@odata.type"
                if($targetType -eq "#microsoft.graph.groupAssignmentTarget")
                {
                    $grp = $groups | Where-Object { $_.id -eq $ass.target.groupId } | Select-Object -First 1
                    $targetDesc = "Include: $(if($grp){$grp.displayName}else{$ass.target.groupId})"
                }
                elseif($targetType -eq "#microsoft.graph.exclusionGroupAssignmentTarget")
                {
                    $grp = $groups | Where-Object { $_.id -eq $ass.target.groupId } | Select-Object -First 1
                    $targetDesc = "Exclude: $(if($grp){$grp.displayName}else{$ass.target.groupId})"
                }
                elseif($targetType -eq "#microsoft.graph.allLicensedUsersAssignmentTarget")
                {
                    $targetDesc = "All Users"
                }
                elseif($targetType -eq "#microsoft.graph.allDevicesAssignmentTarget")
                {
                    $targetDesc = "All Devices"
                }

                $allAssignments += [PSCustomObject]@{
                    ObjectId = $obj.id
                    ObjectName = $obj."$($objectType.NameProp)"
                    AssignmentId = $ass.id
                    TargetDesc = $targetDesc
                    TargetType = $targetType
                    GroupId = $ass.target.groupId
                }
            }
        }
        catch
        {
            Write-Host "  WARNING: Could not load assignments for $($obj."$($objectType.NameProp)")" -ForegroundColor DarkYellow
        }
    }

    if($allAssignments.Count -eq 0)
    {
        Write-Host "No assignments found to remove. Exiting." -ForegroundColor Yellow
        exit 0
    }

    # Deduplicate by target description for selection
    $uniqueTargets = $allAssignments | Select-Object -Property TargetDesc, TargetType, GroupId -Unique
    $targetDisplays = $uniqueTargets | ForEach-Object { $_.TargetDesc }
    $selectedTargetDisplays = Select-MenuItem -Items $targetDisplays -Header "Select assignments to remove (multi-select)" -Multi
    if(-not $selectedTargetDisplays)
    {
        Write-Host "No targets selected. Exiting." -ForegroundColor Yellow
        exit 0
    }

    # Review
    Clear-Host
    Write-Host "Review remove-assignment operation:" -ForegroundColor Green
    Write-Host "  Object Type    : $($objectType.Title)"
    Write-Host "  Objects        : $($selectedObjects.Count)"
    Write-Host "  Targets to remove:" -ForegroundColor Yellow
    foreach($td in $selectedTargetDisplays)
    {
        $count = ($allAssignments | Where-Object { $_.TargetDesc -eq $td } | Measure-Object).Count
        Write-Host "                 - $td ($count occurrence$(if($count -ne 1){'s'}))"
    }
    $confirm = Read-Host "`nProceed? [Y/n]"
    if(-not ([string]::IsNullOrWhiteSpace($confirm) -or $confirm -match "^\s*y"))
    {
        Write-Host "Cancelled." -ForegroundColor Yellow
        exit 0
    }

    # Helper: compute TargetDesc for an assignment
    function Get-AssignmentTargetDesc
    {
        param($Ass)
        $tt = $Ass.target."@odata.type"
        switch($tt)
        {
            "#microsoft.graph.groupAssignmentTarget"
            {
                $grp = $groups | Where-Object { $_.id -eq $Ass.target.groupId } | Select-Object -First 1
                return "Include: $(if($grp){$grp.displayName}else{$Ass.target.groupId})"
            }
            "#microsoft.graph.exclusionGroupAssignmentTarget"
            {
                $grp = $groups | Where-Object { $_.id -eq $Ass.target.groupId } | Select-Object -First 1
                return "Exclude: $(if($grp){$grp.displayName}else{$Ass.target.groupId})"
            }
            "#microsoft.graph.allLicensedUsersAssignmentTarget" { return "All Users" }
            "#microsoft.graph.allDevicesAssignmentTarget"       { return "All Devices" }
            default { return "Unknown" }
        }
    }

    # Execute
    $success = 0
    $failed = 0
    foreach($obj in $selectedObjects)
    {
        $objAssignments = $allAssignments | Where-Object { $_.ObjectId -eq $obj.id -and $_.TargetDesc -in $selectedTargetDisplays }
        if($objAssignments.Count -eq 0) { continue }

        Write-Host "`nProcessing: $($obj."$($objectType.NameProp)")" -ForegroundColor Cyan
        try
        {
            $existing = Invoke-GraphRequest "$($objectType.API)/$($obj.id)/assignments"
            $remaining = @()
            foreach($ea in $existing.value)
            {
                $desc = Get-AssignmentTargetDesc -Ass $ea
                if($desc -in $selectedTargetDisplays)
                {
                    continue
                }
                # Sanitize for re-post
                $clean = $ea | ConvertTo-Json -Depth 50 | ConvertFrom-Json
                if($clean.PSObject.Properties["id"])     { $clean.PSObject.Properties.Remove("id") }
                if($clean.PSObject.Properties["source"]) { $clean.PSObject.Properties.Remove("source") }
                if(-not $clean."@odata.type")
                {
                    $clean | Add-Member -NotePropertyName "@odata.type" -NotePropertyValue $objectType.AssignmentODataType -Force
                }
                $remaining += $clean
            }

            $assignPayload = @{
                $objectType.AssignmentsType = $remaining
            } | ConvertTo-Json -Depth 50 -Compress

            $null = Invoke-GraphRequest "$($objectType.API)/$($obj.id)/assign" -HttpMethod POST -Content $assignPayload
            foreach($ass in $objAssignments)
            {
                Write-Host "  OK: Removed $($ass.TargetDesc)" -ForegroundColor Green
                $success++
            }
        }
        catch
        {
            foreach($ass in $objAssignments)
            {
                Write-Host "  ERROR: Failed to remove $($ass.TargetDesc). $($_.Exception.Message)" -ForegroundColor Red
                $failed++
            }
        }
    }

    Write-Host "`n========================================" -ForegroundColor Cyan
    Write-Host "  Remove Assignments Complete" -ForegroundColor Cyan
    Write-Host "========================================" -ForegroundColor Cyan
    Write-Host "  Success : $success"
    Write-Host "  Failed  : $failed"
}
#endregion