#requires -Version 5.1
<#
.SYNOPSIS
    Backup and restore Intune policy/app assignments.
.DESCRIPTION
    Backs up assignments for selected object types to a JSON file,
    or restores assignments from a previously created backup.
    Works cross-platform (macOS/Linux/Windows) using the headless auth stack.
.EXAMPLE
    # Backup
    ./Scripts/Backup-Restore-Assignments.ps1 -TenantId "..." -Mode Backup -OutputPath ./backups/assignments-backup.json

    # Restore
    ./Scripts/Backup-Restore-Assignments.ps1 -TenantId "..." -Mode Restore -InputPath ./backups/assignments-backup.json
#>
[CmdletBinding()]
param(
    [Parameter(Mandatory = $true)]
    [string]$TenantId,

    [Parameter(Mandatory = $true)]
    [ValidateSet("Backup","Restore")]
    [string]$Mode,

    [string]$OutputPath,

    [string]$InputPath,

    [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 --bind=space:toggle
    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 Validate paths
if($Mode -eq "Backup" -and -not $OutputPath)
{
    throw "Backup mode requires -OutputPath."
}
if($Mode -eq "Restore" -and -not $InputPath)
{
    throw "Restore mode requires -InputPath."
}
if($Mode -eq "Restore" -and -not (Test-Path $InputPath))
{
    throw "Input file not found: $InputPath"
}
#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
$assignableTypes = @(
    [PSCustomObject]@{ Title = "Applications"; API = "/deviceAppManagement/mobileApps"; AssignmentsType = "mobileAppAssignments"; HasIntent = $true; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "Device Configuration"; API = "/deviceManagement/deviceConfigurations"; AssignmentsType = "assignments"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "Settings Catalog"; API = "/deviceManagement/configurationPolicies"; AssignmentsType = "assignments"; HasIntent = $false; NameProp = "name" },
    [PSCustomObject]@{ Title = "Compliance Policies"; API = "/deviceManagement/deviceCompliancePolicies"; AssignmentsType = "assignments"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "Administrative Templates"; API = "/deviceManagement/groupPolicyConfigurations"; AssignmentsType = "assignments"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "Endpoint Security"; API = "/deviceManagement/intents"; AssignmentsType = "assignments"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "App Protection"; API = "/deviceAppManagement/managedAppPolicies"; AssignmentsType = "assignments"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "App Configuration (Device)"; API = "/deviceAppManagement/mobileAppConfigurations"; AssignmentsType = "assignments"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "Platform Scripts"; API = "/deviceManagement/deviceManagementScripts"; AssignmentsType = "deviceManagementScriptAssignments"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "macOS Scripts"; API = "/deviceManagement/deviceShellScripts"; AssignmentsType = "deviceManagementScriptAssignments"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "Device Health Scripts"; API = "/deviceManagement/deviceHealthScripts"; AssignmentsType = "deviceHealthScriptAssignments"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "macOS Custom Attributes"; API = "/deviceManagement/deviceCustomAttributeShellScripts"; AssignmentsType = "deviceManagementScriptAssignments"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "Enrollment Restrictions"; API = "/deviceManagement/deviceEnrollmentConfigurations"; AssignmentsType = "enrollmentConfigurationAssignments"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "Enrollment Status Page"; API = "/deviceManagement/deviceEnrollmentConfigurations"; AssignmentsType = "enrollmentConfigurationAssignments"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "Autopilot"; API = "/deviceManagement/windowsAutopilotDeploymentProfiles"; AssignmentsType = "assignments"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "Terms and Conditions"; API = "/deviceManagement/termsAndConditions"; AssignmentsType = "assignments"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "Policy Sets"; API = "/deviceAppManagement/policySets"; AssignmentsType = "assignments"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "Update Policies"; API = "/deviceManagement/windowsUpdateForBusinessConfigurations"; AssignmentsType = "assignments"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "Feature Updates"; API = "/deviceManagement/windowsFeatureUpdateProfiles"; AssignmentsType = "assignments"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "Quality Updates"; API = "/deviceManagement/windowsQualityUpdateProfiles"; AssignmentsType = "assignments"; HasIntent = $false; NameProp = "displayName" },
    [PSCustomObject]@{ Title = "Device Management Intents"; API = "/deviceManagement/intents"; AssignmentsType = "assignments"; HasIntent = $false; NameProp = "displayName" }
)
#endregion

#region BACKUP
if($Mode -eq "Backup")
{
    $typeTitles = $assignableTypes | ForEach-Object { $_.Title }
    $selectedTypeTitles = Select-MenuItem -Items $typeTitles -Header "Select object types to back up (multi-select)" -Multi
    if(-not $selectedTypeTitles)
    {
        Write-Host "No types selected. Exiting." -ForegroundColor Yellow
        exit 0
    }

    # Preload groups for name resolution in backup
    Write-Host "`nLoading groups for backup resolution..." -ForegroundColor Cyan
    $backupGroupsResponse = Invoke-GraphRequest "/groups?`$select=id,displayName&`$orderby=displayName" -AllPages
    $backupGroups = @{}
    foreach($g in $backupGroupsResponse.value)
    {
        $backupGroups[$g.id] = $g.displayName
    }

    $backupData = @{
        TenantId = $org.value[0].id
        TenantName = $org.value[0].displayName
        Created = (Get-Date -Format "o")
        Groups = $backupGroups
        Objects = @()
    }

    foreach($typeTitle in $selectedTypeTitles)
    {
        $objectType = $assignableTypes | Where-Object { $_.Title -eq $typeTitle } | Select-Object -First 1
        Write-Host "`nBacking up $($objectType.Title) assignments..." -ForegroundColor Cyan

        try
        {
            $objectsResponse = Invoke-GraphRequest "$($objectType.API)?`$select=id,$($objectType.NameProp)&`$orderby=$($objectType.NameProp)"
            $objects = $objectsResponse.value | Where-Object { $_ }
            Write-Host "  Found $($objects.Count) objects" -ForegroundColor Green

            foreach($obj in $objects)
            {
                try
                {
                    $assignmentsResponse = Invoke-GraphRequest "$($objectType.API)/$($obj.id)/assignments"
                    $assignments = $assignmentsResponse.value
                    if($assignments.Count -gt 0)
                    {
                        # Enrich assignments with group display names for cross-tenant restore
                        $enrichedAssignments = $assignments | ConvertTo-Json -Depth 50 | ConvertFrom-Json
                        foreach($ass in $enrichedAssignments)
                        {
                            if($ass.target.groupId -and $backupGroups.ContainsKey($ass.target.groupId))
                            {
                                $ass.target | Add-Member -NotePropertyName "_backupGroupName" -NotePropertyValue $backupGroups[$ass.target.groupId] -Force
                            }
                        }
                        $backupData.Objects += [PSCustomObject]@{
                            ObjectType = $objectType.Title
                            ObjectId = $obj.id
                            ObjectName = if($objectType.NameProp -eq "name") { $obj.name } else { $obj.displayName }
                            NameProp = $objectType.NameProp
                            API = $objectType.API
                            AssignmentsType = $objectType.AssignmentsType
                            Assignments = $enrichedAssignments
                        }
                    }
                }
                catch
                {
                    Write-Host "  WARNING: Could not backup assignments for $($obj."$($objectType.NameProp)")" -ForegroundColor DarkYellow
                }
            }
        }
        catch
        {
            Write-Host "  WARNING: Could not load objects for $($objectType.Title)" -ForegroundColor DarkYellow
        }
    }

    $backupJson = $backupData | ConvertTo-Json -Depth 50
    $OutputPath = (Resolve-Path (Split-Path -Parent $OutputPath) -ErrorAction SilentlyContinue).Path + "/" + (Split-Path -Leaf $OutputPath)
    $backupJson | Out-File -LiteralPath $OutputPath -Encoding utf8 -Force

    $totalAssignments = 0
    foreach($obj in $backupData.Objects) { $totalAssignments += $obj.Assignments.Count }

    Write-Host "`n========================================" -ForegroundColor Cyan
    Write-Host "  Backup Complete" -ForegroundColor Cyan
    Write-Host "========================================" -ForegroundColor Cyan
    Write-Host "  File          : $OutputPath"
    Write-Host "  Objects       : $($backupData.Objects.Count)"
    Write-Host "  Assignments   : $totalAssignments"
}
#endregion

#region RESTORE
elseif($Mode -eq "Restore")
{
    $backup = Get-Content $InputPath -Raw | ConvertFrom-Json

    Write-Host "`nBackup info:" -ForegroundColor Cyan
    Write-Host "  Tenant : $($backup.TenantName) ($($backup.TenantId))"
    Write-Host "  Created: $($backup.Created)"
    Write-Host "  Objects: $($backup.Objects.Count)"

    $currentTenantId = $org.value[0].id
    if($backup.TenantId -ne $currentTenantId)
    {
        Write-Host "`nWARNING: Backup is from a different tenant!" -ForegroundColor Yellow
        if(-not (Read-YesNo -Prompt "Continue anyway?" -Default $false))
        {
            Write-Host "Cancelled." -ForegroundColor Yellow
            exit 0
        }
    }

    # Resolve group names to IDs in current tenant if needed
    Write-Host "`nLoading current tenant groups for name resolution..." -ForegroundColor Cyan
    $currentGroupsResponse = Invoke-GraphRequest "/groups?`$select=id,displayName&`$orderby=displayName" -AllPages
    $currentGroups = $currentGroupsResponse.value

    $success = 0
    $skipped = 0
    $failed = 0

    foreach($entry in $backup.Objects)
    {
        Write-Host "`nRestoring: $($entry.ObjectName) ($($entry.ObjectType))" -ForegroundColor Cyan

        # Try to find the object in current tenant by displayName
        $nameProp = ?? $entry.NameProp "displayName"
        $searchUrl = "$($entry.API)?`$filter=$nameProp eq '$([uri]::EscapeDataString($entry.ObjectName))'&`$select=id,$nameProp"
        try
        {
            $searchResult = Invoke-GraphRequest $searchUrl
            $targetObj = $searchResult.value | Select-Object -First 1
        }
        catch
        {
            $targetObj = $null
        }

        if(-not $targetObj)
        {
            Write-Host "  SKIP: Object '$($entry.ObjectName)' not found in current tenant" -ForegroundColor DarkYellow
            $failed++
            continue
        }

        # Load existing assignments to avoid duplicates
        try
        {
            $existing = Invoke-GraphRequest "$($entry.API)/$($targetObj.id)/assignments"
            $existingTargets = $existing.value
        }
        catch
        {
            Write-Host "  ERROR: Could not load existing assignments" -ForegroundColor Red
            $failed++
            continue
        }

        function Test-BackupAssignmentExists
        {
            param($assignment, $existingList)
            $t = $assignment.target
            foreach($ea in $existingList)
            {
                $et = $ea.target
                if($t."@odata.type" -ne $et."@odata.type") { continue }
                if($t.groupId -and $t.groupId -ne $et.groupId) { continue }
                # Also match intent for apps
                if($entry.AssignmentsType -eq "mobileAppAssignments" -and ($assignment.intent -ne $ea.intent)) { continue }
                return $true
            }
            return $false
        }

        foreach($assignment in $entry.Assignments)
        {
            # Clone assignment to avoid modifying backup data
            $restoredAssignment = $assignment | ConvertTo-Json -Depth 50 | ConvertFrom-Json

            # Remove Id
            if($restoredAssignment.PSObject.Properties["id"])
            {
                $restoredAssignment.PSObject.Properties.Remove("id")
            }

            # Map group IDs if cross-tenant
            if($backup.TenantId -ne $currentTenantId -and $restoredAssignment.target.groupId)
            {
                $originalGroupName = $restoredAssignment.target."_backupGroupName"
                if($originalGroupName)
                {
                    $matchedGroup = $currentGroups | Where-Object { $_.displayName -eq $originalGroupName } | Select-Object -First 1
                    if($matchedGroup)
                    {
                        Write-Host "  MAPPED: Group '$originalGroupName' -> $($matchedGroup.id)" -ForegroundColor Gray
                        $restoredAssignment.target.groupId = $matchedGroup.id
                    }
                    else
                    {
                        Write-Host "  SKIP: Could not find group '$originalGroupName' in current tenant" -ForegroundColor DarkYellow
                        $skipped++
                        continue
                    }
                }
                else
                {
                    Write-Host "  SKIP: Cross-tenant restore cannot resolve group without name mapping" -ForegroundColor DarkYellow
                    $skipped++
                    continue
                }
            }

            # Clean up internal property before sending
            if($restoredAssignment.target.PSObject.Properties["_backupGroupName"])
            {
                $restoredAssignment.target.PSObject.Properties.Remove("_backupGroupName")
            }

            if(Test-BackupAssignmentExists -assignment $restoredAssignment -existingList $existingTargets)
            {
                Write-Host "  SKIP: Assignment already exists" -ForegroundColor DarkYellow
                $skipped++
                continue
            }

            # Prepare payload
            $payload = @{
                $entry.AssignmentsType = @($restoredAssignment)
            }

            try
            {
                $body = $payload | ConvertTo-Json -Depth 50 -Compress
                $null = Invoke-GraphRequest "$($entry.API)/$($targetObj.id)/assign" -HttpMethod POST -Content $body
                Write-Host "  OK: Restored assignment" -ForegroundColor Green
                $success++
            }
            catch
            {
                Write-Host "  ERROR: Failed to restore assignment. $($_.Exception.Message)" -ForegroundColor Red
                $failed++
            }
        }
    }

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