<#
.SYNOPSIS
    Replaces a license for Microsoft 365 users.
.DESCRIPTION
    This script replaces a specified license (SKU) with another license for Microsoft 365 users.
.PARAMETER InputObject
    The Microsoft Graph user objects to process.
.PARAMETER RemoveSkuPartNumber
    The SKU Part Number of the license to be removed.
.PARAMETER RemoveSkuId
    The SKU Id of the license to be removed.
.PARAMETER AddSkuPartNumber
    The SKU Part Number of the license to be added.
.PARAMETER AddSkuId
    The SKU Id of the license to be added.
.PARAMETER All
    If specified, all users with the license to be removed will be processed.
.PARAMETER WhatIf
    If specified, the actions will be displayed without making any changes.
#>
function Set-License {
    [CmdletBinding()]
    param (
        [Parameter(ValueFromPipeline=$true)]
        [array]
        $InputObject,
        [Parameter(Mandatory=$true,
        ParameterSetName='SkuPartNumber')]
        [string]
        $RemoveSkuPartNumber,
        [Parameter(Mandatory=$true,
        ParameterSetName='SkuId')]
        [guid]
        $RemoveSkuId,
        [Parameter(Mandatory=$true,
        ParameterSetName='SkuPartNumber')]
        [string]
        $AddSkuPartNumber,
        [Parameter(Mandatory=$true,
        ParameterSetName='SkuId')]
        [guid]
        $AddSkuId,
        [Parameter()]
        [switch]
        $All,
        [Parameter()]
        [switch]
        $WhatIf
    )

    begin {
            if ($null -eq (Get-MgContext)) {
            throw 'You are not connected to Microsoft Graph. Please connect using Connect-MgGraph.'
        }

        if ($PSCmdlet.ParameterSetName -eq 'SkuPartNumber') {
            # Validate that the Remove SKU exists and get its Id
            $RemoveSku = (Get-MgSubscribedSku -All | Where-Object { $_.SkuPartNumber -eq $RemoveSkuPartNumber })
            if ($null -eq $RemoveSku) {
                throw "The SKU Part Number '$RemoveSkuPartNumber' was not found."
            }
            $RemoveSkuId = $RemoveSku.SkuId

            # Validate that the Add SKU exists and get its Id
            $AddSku = (Get-MgSubscribedSku -All | Where-Object { $_.SkuPartNumber -eq $AddSkuPartNumber })
            if ($null -eq $AddSku) {
                throw "The SKU Part Number '$AddSkuPartNumber' was not found."
            }
            $AddSkuId = $AddSku.SkuId
        } elseif ($PSCmdlet.ParameterSetName -eq 'SkuId') {
            # Validate that the Remove SKU exists
            $RemoveSku = (Get-MgSubscribedSku -All | Where-Object { $_.SkuId -eq $RemoveSkuId })
            if ($null -eq $RemoveSku) {
                throw "The SKU Id '$RemoveSkuId' was not found."
            }

            # Validate that the Add SKU exists
            $AddSku = (Get-MgSubscribedSku -All | Where-Object { $_.SkuId -eq $AddSkuId })
            if ($null -eq $AddSku) {
                throw "The SKU Id '$AddSkuId' was not found."
            }
        }
    }

    process {
        $Users = @()
        if ($All) {
            # Get all users with the old SKU assigned
            $Users = Get-MgUser -Filter "assignedLicenses/any(x:x/skuId eq $RemoveSkuId)" -All
        } else {
            # Get users from the pipeline with the old SKU assigned
            $Users += (Get-MgUser -UserId $InputObject.Id -Property Id,DisplayName,UserPrincipalName,AssignedLicenses) | Where-Object { $RemoveSkuId -in $_.AssignedLicenses.SkuId }
        }

        foreach ($User in $Users) {
            if ($WhatIf) {
                Write-Host "WhatIf: Replacing license for user '$($User.UserPrincipalName)' - Removing SKU Id '$RemoveSkuId', Adding SKU Id '$AddSkuId'"
            } else {
                # Replace the license
                Set-MgUserLicense -UserId $User.Id -RemoveLicenses @($RemoveSkuId) -AddLicenses @{}
                Set-MgUserLicense -UserId $User.Id -RemoveLicenses @() -AddLicenses @{ SkuId = $AddSkuId }
            }
        }
    }
}

Export-ModuleMember -Function Set-License