[CmdletBinding()]
param(
    [Parameter(Mandatory)]
    [ValidatePattern('^[0-9a-fA-F-]{36}$')]
    [string]$SubscriptionId,

    [ValidateNotNullOrEmpty()]
    [string]$Location = 'westeurope',

    [ValidateNotNullOrEmpty()]
    [string]$VmSize = 'Standard_D2s_v5',

    [ValidateRange(1, 10)]
    [int]$HostCount = 2,

    [switch]$CreateTestUser,

    [ValidatePattern('^[a-z0-9-]{3,32}$')]
    [string]$TestUserPrefix = 'avd-workshop',

    [ValidateSet('RequireInteractive','AllowConfigurationOnly')]
    [string]$ProfileEvidenceMode = 'AllowConfigurationOnly',

    [ValidateRange(5, 120)]
    [int]$InteractiveTimeoutMinutes = 20,

    [ValidateSet('Preflight','Identity','Foundation','Storage','HostPool','SessionHosts','FSLogix','Monitoring','Scaling','TestProfile','TestScaling','Validate','All')]
    [string]$Stage = 'All'
)

Set-StrictMode -Version Latest
$ErrorActionPreference = 'Stop'
. (Join-Path $PSScriptRoot 'AvdWorkshop.Common.ps1')

$script:Paths = Get-AvdWorkshopPaths -ScriptRoot $PSScriptRoot
Initialize-AvdWorkshopDirectory -Paths $script:Paths
$script:State = $null

function Save-State {
    Write-AvdWorkshopState -State $script:State -Path $script:Paths.StatePath
}

function Set-StageComplete {
    param([Parameter(Mandatory)][string]$Name)

    $timestamp = [DateTime]::UtcNow.ToString('o')
    if (@($script:State.Stages.PSObject.Properties | ForEach-Object Name) -contains $Name) {
        $script:State.Stages.$Name = $timestamp
    }
    else {
        $script:State.Stages | Add-Member -NotePropertyName $Name -NotePropertyValue $timestamp
    }
    Save-State
}

function Assert-ExpectedResourceGroup {
    if (-not $script:State) { return }
    $existing = Invoke-AvdAzJson -Arguments @('group','show','--name',$script:State.ResourceGroup) -AllowFailure -Quiet
    if (-not $existing) { return }

    if ([string]$existing.tags.workshop -ne 'azure-virtual-desktop-deep-dive' -or
        [string]$existing.tags.'workshop-session' -ne [string]$script:State.SessionId) {
        throw "Resource group $($script:State.ResourceGroup) exists but isn't owned by this workshop session. No changes were made."
    }
}

function Ensure-State {
    if (-not $script:State) {
        $script:State = Read-AvdWorkshopState -Path $script:Paths.StatePath -AllowMissing
    }
    if (-not $script:State) { Invoke-Preflight }
    Assert-ExpectedResourceGroup
}

function Get-EligibleLicenseInventory {
    $eligibleSkuParts = @(
        'SPE_E3','SPE_E5','SPE_F1','SPE_F3','SPB','ENTERPRISEPACK','ENTERPRISEPREMIUM',
        'WIN10_PRO_ENT_SUB','WIN10_VDA_E3','WIN10_VDA_E5','M365EDU_A3_FACULTY',
        'M365EDU_A5_FACULTY','M365EDU_A3_STUDENT','M365EDU_A5_STUDENT'
    )
    $response = Invoke-AvdGraphJson -Method GET -Path 'subscribedSkus'
    $inventory = foreach ($sku in @($response.value)) {
        if ($eligibleSkuParts -notcontains [string]$sku.skuPartNumber) { continue }
        $enabled = [int]$sku.prepaidUnits.enabled
        $consumed = [int]$sku.consumedUnits
        [pscustomobject]@{
            SkuId = [string]$sku.skuId
            SkuPartNumber = [string]$sku.skuPartNumber
            Enabled = $enabled
            Consumed = $consumed
            Available = [Math]::Max(0, $enabled - $consumed)
            CapabilityStatus = [string]$sku.capabilityStatus
        }
    }
    @($inventory)
}

function Invoke-Preflight {
    Write-Host 'Stage: Preflight (read-only Azure and Microsoft Entra checks)' -ForegroundColor Cyan
    if ($PSVersionTable.PSVersion.Major -lt 7) { throw 'PowerShell 7 or later is required.' }
    if (-not (Get-Command az -ErrorAction SilentlyContinue)) { throw 'Azure CLI is required.' }

    $account = Invoke-AvdAzJson -Arguments @('account','show') -Quiet
    if (-not $account -or [string]$account.state -ne 'Enabled') { throw 'Sign in with Azure CLI to an enabled subscription.' }
    if ([string]$account.id -ne $SubscriptionId) {
        $requested = Invoke-AvdAzJson -Arguments @('account','show','--subscription',$SubscriptionId) -Quiet
        if (-not $requested -or [string]$requested.state -ne 'Enabled') { throw 'The requested subscription is unavailable or disabled.' }
        $account = $requested
    }

    $existingState = Read-AvdWorkshopState -Path $script:Paths.StatePath -AllowMissing
    if ($existingState) {
        if ([string]$existingState.SubscriptionId -ne $SubscriptionId) { throw 'Private runtime state belongs to another subscription.' }
        if ([string]$existingState.Location -ne $Location -or [string]$existingState.VmSize -ne $VmSize -or [int]$existingState.HostCount -ne $HostCount) {
            throw 'Location, VM size, and host count must match the active private workshop state.'
        }
        $script:State = $existingState
        Assert-ExpectedResourceGroup
        Write-Host 'Preflight passed; the existing private workshop state is consistent.' -ForegroundColor Green
        return
    }

    $resourceGroup = 'rg-avd-deepdive-we'
    if (Invoke-AvdAzJson -Arguments @('group','show','--name',$resourceGroup) -AllowFailure -Quiet) {
        throw "Resource group $resourceGroup already exists and no matching private state exists. Refusing to adopt it."
    }

    $providers = @('Microsoft.DesktopVirtualization','Microsoft.Compute','Microsoft.Network','Microsoft.Storage','Microsoft.Insights','Microsoft.OperationalInsights')
    $providerState = [ordered]@{}
    foreach ($provider in $providers) {
        $providerState[$provider] = [string](Invoke-AvdAzJson -Arguments @('provider','show','--namespace',$provider,'--query','registrationState') -Quiet)
    }

    $skuQuery = "[?name=='$VmSize' && resourceType=='virtualMachines'] | [0]"
    $sku = Invoke-AvdAzJson -Arguments @('vm','list-skus','--location',$Location,'--size',$VmSize,'--all','--query',$skuQuery) -Quiet
    if (-not $sku) { throw "VM size $VmSize isn't available in $Location." }
    $restrictions = @($sku.restrictions | Where-Object { $_.reasonCode -and $_.reasonCode -ne 'NotAvailableForSubscription' })
    if (@($sku.restrictions).Count -gt 0 -and $restrictions.Count -gt 0) { throw "VM size $VmSize is restricted in $Location." }

    $image = Invoke-AvdAzJson -Arguments @(
        'vm','image','show','--location',$Location,'--urn','MicrosoftWindowsDesktop:windows-11:win11-24h2-avd:latest'
    ) -Quiet
    if (-not $image -or [string]$image.hyperVGeneration -ne 'V2') { throw 'The Windows 11 24H2 AVD Gen2 image is not available.' }

    $usage = Invoke-AvdAzJson -Arguments @('vm','list-usage','--location',$Location) -Quiet
    $regional = @($usage | Where-Object { $_.name.value -eq 'cores' })[0]
    $dsv5 = @($usage | Where-Object { $_.name.value -eq 'standardDSv5Family' })[0]
    $neededCores = 2 * $HostCount
    if (($regional.limit - $regional.currentValue) -lt $neededCores -or ($dsv5.limit - $dsv5.currentValue) -lt $neededCores) {
        throw "Insufficient quota for $HostCount $VmSize hosts ($neededCores vCPUs required)."
    }

    $signedInUser = Invoke-AvdAzJson -Arguments @('ad','signed-in-user','show') -Quiet
    $assignments = Invoke-AvdAzJson -Arguments @(
        'role','assignment','list','--assignee-object-id',[string]$signedInUser.id,'--all','--include-inherited'
    ) -Quiet
    $roleNames = @($assignments | ForEach-Object roleDefinitionName)
    if (($roleNames -notcontains 'Owner') -and ($roleNames -notcontains 'User Access Administrator')) {
        throw 'Owner or User Access Administrator is required for the autoscale and user-login role assignments.'
    }

    $organization = Invoke-AvdGraphJson -Method GET -Path 'organization?$select=id,verifiedDomains'
    $defaultDomain = [string](@($organization.value[0].verifiedDomains | Where-Object isDefault)[0].name)
    if (-not $defaultDomain) { throw 'No default verified Microsoft Entra tenant domain was found.' }

    $licenseInventory = Get-EligibleLicenseInventory
    $availableLicense = @($licenseInventory | Where-Object { $_.Available -gt 0 -and $_.CapabilityStatus -eq 'Enabled' } | Select-Object -First 1)

    $avdServicePrincipal = Invoke-AvdGraphJson -Method GET -Path "servicePrincipals?`$filter=appId eq '9cdead84-a844-4324-93f2-b2e6bb768d07'"
    $avdSp = $avdServicePrincipal.value | Select-Object -First 1
    if (-not $avdSp) { throw 'The Windows Virtual Desktop service principal was not found in this tenant.' }

    $subscriptionScope = "/subscriptions/$SubscriptionId"
    $autoscaleBaseline = Invoke-AvdAzJson -Arguments @(
        'role','assignment','list','--scope',$subscriptionScope,'--assignee-object-id',[string]$avdSp.id,
        '--include-inherited','--query',"[?roleDefinitionName=='Desktop Virtualization Power On Off Contributor']"
    ) -Quiet

    $suffix = New-AvdWorkshopSuffix
    $userPrincipalName = "$TestUserPrefix-$suffix@$defaultDomain"
    $userFilter = [uri]::EscapeDataString("userPrincipalName eq '$userPrincipalName'")
    $existingUser = Invoke-AvdGraphJson -Method GET -Path "users?`$filter=$userFilter"
    if (@($existingUser.value).Count -gt 0) { throw 'The generated workshop identity unexpectedly already exists.' }

    $scope = "/subscriptions/$SubscriptionId"
    $resourceGroupId = "$scope/resourceGroups/$resourceGroup"
    $hostPoolName = 'hp-avd-pooled-we'
    $applicationGroupName = 'dag-avd-pooled-we'
    $workspaceName = 'ws-avd-deepdive-we'

    $script:State = [pscustomobject]@{
        SchemaVersion = 1
        CreatedUtc = [DateTime]::UtcNow.ToString('o')
        UpdatedUtc = [DateTime]::UtcNow.ToString('o')
        SessionId = $suffix
        SubscriptionId = $SubscriptionId
        SubscriptionName = [string]$account.name
        TenantId = [string]$account.tenantId
        Scope = $scope
        Location = $Location
        VmSize = $VmSize
        HostCount = $HostCount
        ProfileEvidenceMode = $ProfileEvidenceMode
        InteractiveTimeoutMinutes = $InteractiveTimeoutMinutes
        ResourceGroup = $resourceGroup
        ResourceGroupId = $resourceGroupId
        WorkshopTag = 'azure-virtual-desktop-deep-dive'
        VnetName = 'vnet-avd-deepdive-we'
        SessionSubnetName = 'snet-avd-sessionhosts'
        PrivateEndpointSubnetName = 'snet-avd-private-endpoints'
        NsgName = 'nsg-avd-sessionhosts-we'
        NatGatewayName = 'ng-avd-outbound-we'
        NatPublicIpName = 'pip-avd-nat-we'
        StorageAccount = "stavd$suffix"
        FileShareName = 'fslogix'
        PrivateEndpointName = 'pe-avd-fslogix-we'
        PrivateDnsZone = 'privatelink.file.core.windows.net'
        PrivateDnsLinkName = 'link-avd-vnet-we'
        HostPoolName = $hostPoolName
        HostPoolId = "$resourceGroupId/providers/Microsoft.DesktopVirtualization/hostPools/$hostPoolName"
        ApplicationGroupName = $applicationGroupName
        ApplicationGroupId = "$resourceGroupId/providers/Microsoft.DesktopVirtualization/applicationGroups/$applicationGroupName"
        WorkspaceName = $workspaceName
        WorkspaceId = "$resourceGroupId/providers/Microsoft.DesktopVirtualization/workspaces/$workspaceName"
        ScalingPlanName = 'sp-avd-power-we'
        ScalingPlanId = "$resourceGroupId/providers/Microsoft.DesktopVirtualization/scalingPlans/sp-avd-power-we"
        LogAnalyticsName = 'law-avd-deepdive-we'
        DcrName = 'dcr-avd-sessionhosts-we'
        DcrId = "$resourceGroupId/providers/Microsoft.Insights/dataCollectionRules/dcr-avd-sessionhosts-we"
        VmNames = @(1..$HostCount | ForEach-Object { 'vm-avd-{0:d2}' -f $_ })
        NicNames = @(1..$HostCount | ForEach-Object { 'nic-avd-{0:d2}' -f $_ })
        ComputerNames = @(1..$HostCount | ForEach-Object { 'avdhost-{0:d2}' -f $_ })
        AvdApiVersion = '2025-10-10'
        NetworkApiVersion = '2024-06-01'
        MonitorApiVersion = '2024-03-11'
        DiagnosticsApiVersion = '2021-05-01-preview'
        RegistrationToken = $null
        LocalAdminPassword = New-AvdWorkshopPassword
        TestUserPrefix = $TestUserPrefix
        TestUserPrincipalName = $userPrincipalName
        TestUserPassword = $null
        TestUserObjectId = $null
        TestUserCreated = $false
        TestUserRoleAssignmentIds = @()
        AssignedLicenseSkuId = $null
        AssignedLicensePartNumber = $null
        LicenseInventoryBaseline = @($licenseInventory)
        LicenseAvailableAtPreflight = ($availableLicense.Count -gt 0)
        TenantDefaultDomain = $defaultDomain
        AvdServicePrincipalObjectId = [string]$avdSp.id
        AutoscaleRoleBaselinePresent = (@($autoscaleBaseline).Count -gt 0)
        AutoscaleRoleAssignmentId = $null
        AutoscaleRoleCreated = $false
        ProviderRegistrationBaseline = $providerState
        StorageApplicationObjectId = $null
        StorageServicePrincipalObjectId = $null
        StorageApplicationCreated = $false
        StorageConsentGrantId = $null
        StorageConsentCreated = $false
        ProfilePersistenceStatus = 'Pending'
        ProfileFirstHost = $null
        ProfileSecondHost = $null
        ProfileVhdxCount = $null
        ScalingTestStatus = 'Pending'
        ScalingActivityCorrelationId = $null
        OriginalScalingSchedule = $null
        Stages = [pscustomobject]@{}
    }
    Save-State
    Write-Host 'Preflight passed. No Azure or Microsoft Entra resources were changed.' -ForegroundColor Green
    if (-not $script:State.LicenseAvailableAtPreflight) {
        Write-Warning 'No eligible unused AVD license is available. AllowConfigurationOnly will publish ProfilePersistenceStatus = NotTested-LicenseUnavailable.'
    }
}

function Invoke-Identity {
    Ensure-State
    Write-Host 'Stage: Identity' -ForegroundColor Cyan
    if (-not $CreateTestUser) {
        throw 'Identity creation is intentionally gated. Re-run Stage Identity with -CreateTestUser only after action-time confirmation.'
    }

    $s = $script:State
    $filter = [uri]::EscapeDataString("userPrincipalName eq '$($s.TestUserPrincipalName)'")
    $existing = Invoke-AvdGraphJson -Method GET -Path "users?`$filter=$filter"
    if (@($existing.value).Count -gt 0 -and -not [bool]$s.TestUserCreated) {
        throw 'The planned user already exists but is not owned by this workshop state. Refusing to adopt it.'
    }

    if (-not [bool]$s.TestUserCreated) {
        $password = New-AvdWorkshopPassword
        $body = @{
            accountEnabled = $true
            displayName = 'AVD Workshop Temporary User'
            mailNickname = "$($s.TestUserPrefix)-$($s.SessionId)"
            userPrincipalName = [string]$s.TestUserPrincipalName
            userType = 'Member'
            usageLocation = 'DE'
            passwordProfile = @{
                forceChangePasswordNextSignIn = $false
                password = $password
            }
        }
        $created = Invoke-AvdGraphJson -Method POST -Path 'users' -Body $body
        $script:State.TestUserPassword = $password
        $script:State.TestUserObjectId = [string]$created.id
        $script:State.TestUserCreated = $true
        Save-State
    }

    $inventory = Get-EligibleLicenseInventory
    $license = @($inventory | Where-Object { $_.Available -gt 0 -and $_.CapabilityStatus -eq 'Enabled' } | Select-Object -First 1)
    if ($license.Count -gt 0 -and -not $s.AssignedLicenseSkuId) {
        Invoke-AvdGraphJson -Method POST -Path "users/$($s.TestUserObjectId)/assignLicense" -Body @{
            addLicenses = @(@{ skuId = [string]$license[0].SkuId })
            removeLicenses = @()
        } | Out-Null
        $script:State.AssignedLicenseSkuId = [string]$license[0].SkuId
        $script:State.AssignedLicensePartNumber = [string]$license[0].SkuPartNumber
    }

    foreach ($assignment in @(
        @{ Role = 'Desktop Virtualization User'; Scope = [string]$s.ApplicationGroupId },
        @{ Role = 'Virtual Machine User Login'; Scope = [string]$s.ResourceGroupId }
    )) {
        $current = Invoke-AvdAzJson -Arguments @(
            'role','assignment','list','--assignee-object-id',[string]$s.TestUserObjectId,
            '--scope',$assignment.Scope,'--query',"[?roleDefinitionName=='$($assignment.Role)'] | [0]"
        ) -Quiet
        if (-not $current) {
            $current = Invoke-AvdAzJson -Arguments @(
                'role','assignment','create','--assignee-object-id',[string]$s.TestUserObjectId,
                '--assignee-principal-type','User','--role',$assignment.Role,'--scope',$assignment.Scope
            ) -Quiet
        }
        if ([string]$current.id -and @($script:State.TestUserRoleAssignmentIds) -notcontains [string]$current.id) {
            $script:State.TestUserRoleAssignmentIds = @($script:State.TestUserRoleAssignmentIds) + [string]$current.id
        }
    }

    if (-not $script:State.AssignedLicenseSkuId) {
        $script:State.ProfilePersistenceStatus = 'NotTested-LicenseUnavailable'
        Write-Warning 'The user was created without a license; interactive profile validation remains configuration-only.'
    }
    Set-StageComplete -Name 'Identity'
}

function Invoke-Foundation {
    Ensure-State
    Write-Host 'Stage: Foundation' -ForegroundColor Cyan
    $s = $script:State
    Invoke-AvdAzCli -Arguments @('account','set','--subscription',$SubscriptionId) -Quiet | Out-Null

    foreach ($provider in @('Microsoft.DesktopVirtualization','Microsoft.Compute','Microsoft.Network','Microsoft.Storage','Microsoft.Insights','Microsoft.OperationalInsights')) {
        Invoke-AvdAzCli -Arguments @('provider','register','--namespace',$provider,'--wait') -Quiet | Out-Null
    }

    $tags = @('workshop=azure-virtual-desktop-deep-dive',"workshop-session=$($s.SessionId)",'environment=lab','owner=beyondcloudwithchriz','workload=avd')
    if (-not (Invoke-AvdAzJson -Arguments @('group','show','--name',$s.ResourceGroup) -AllowFailure -Quiet)) {
        Invoke-AvdAzCli -Arguments (@('group','create','--name',$s.ResourceGroup,'--location',$s.Location,'--tags') + $tags) -Quiet | Out-Null
    }
    Assert-ExpectedResourceGroup

    if (-not (Invoke-AvdAzJson -Arguments @('network','nsg','show','--resource-group',$s.ResourceGroup,'--name',$s.NsgName) -AllowFailure -Quiet)) {
        Invoke-AvdAzCli -Arguments (@('network','nsg','create','--resource-group',$s.ResourceGroup,'--name',$s.NsgName,'--location',$s.Location,'--tags') + $tags) -Quiet | Out-Null
    }
    if (-not (Invoke-AvdAzJson -Arguments @('network','vnet','show','--resource-group',$s.ResourceGroup,'--name',$s.VnetName) -AllowFailure -Quiet)) {
        Invoke-AvdAzCli -Arguments (@('network','vnet','create','--resource-group',$s.ResourceGroup,'--name',$s.VnetName,'--location',$s.Location,'--address-prefixes','10.210.0.0/16','--tags') + $tags) -Quiet | Out-Null
    }
    if (-not (Invoke-AvdAzJson -Arguments @('network','vnet','subnet','show','--resource-group',$s.ResourceGroup,'--vnet-name',$s.VnetName,'--name',$s.SessionSubnetName) -AllowFailure -Quiet)) {
        Invoke-AvdAzCli -Arguments @('network','vnet','subnet','create','--resource-group',$s.ResourceGroup,'--vnet-name',$s.VnetName,'--name',$s.SessionSubnetName,'--address-prefixes','10.210.1.0/24','--network-security-group',$s.NsgName) -Quiet | Out-Null
    }
    if (-not (Invoke-AvdAzJson -Arguments @('network','vnet','subnet','show','--resource-group',$s.ResourceGroup,'--vnet-name',$s.VnetName,'--name',$s.PrivateEndpointSubnetName) -AllowFailure -Quiet)) {
        Invoke-AvdAzCli -Arguments @('network','vnet','subnet','create','--resource-group',$s.ResourceGroup,'--vnet-name',$s.VnetName,'--name',$s.PrivateEndpointSubnetName,'--address-prefixes','10.210.2.0/24','--disable-private-endpoint-network-policies','true') -Quiet | Out-Null
    }

    if (-not (Invoke-AvdAzJson -Arguments @('network','public-ip','show','--resource-group',$s.ResourceGroup,'--name',$s.NatPublicIpName) -AllowFailure -Quiet)) {
        Invoke-AvdAzCli -Arguments (@('network','public-ip','create','--resource-group',$s.ResourceGroup,'--name',$s.NatPublicIpName,'--location',$s.Location,'--sku','Standard','--allocation-method','Static','--version','IPv4','--tags') + $tags) -Quiet | Out-Null
    }
    if (-not (Invoke-AvdAzJson -Arguments @('network','nat','gateway','show','--resource-group',$s.ResourceGroup,'--name',$s.NatGatewayName) -AllowFailure -Quiet)) {
        Invoke-AvdAzCli -Arguments (@('network','nat','gateway','create','--resource-group',$s.ResourceGroup,'--name',$s.NatGatewayName,'--location',$s.Location,'--public-ip-addresses',$s.NatPublicIpName,'--idle-timeout','10','--tags') + $tags) -Quiet | Out-Null
    }
    Invoke-AvdAzCli -Arguments @('network','vnet','subnet','update','--resource-group',$s.ResourceGroup,'--vnet-name',$s.VnetName,'--name',$s.SessionSubnetName,'--nat-gateway',$s.NatGatewayName) -Quiet | Out-Null
    Set-StageComplete -Name 'Foundation'
}

function Invoke-Storage {
    Ensure-State
    Write-Host 'Stage: Storage and Microsoft Entra Kerberos' -ForegroundColor Cyan
    $s = $script:State
    $tags = @('workshop=azure-virtual-desktop-deep-dive',"workshop-session=$($s.SessionId)",'environment=lab','owner=beyondcloudwithchriz','workload=fslogix')

    if (-not (Invoke-AvdAzJson -Arguments @('storage','account','show','--resource-group',$s.ResourceGroup,'--name',$s.StorageAccount) -AllowFailure -Quiet)) {
        Invoke-AvdAzCli -Arguments (@(
            'storage','account','create','--resource-group',$s.ResourceGroup,'--name',$s.StorageAccount,
            '--location',$s.Location,'--sku','Standard_LRS','--kind','StorageV2','--https-only','true',
            '--min-tls-version','TLS1_2','--allow-blob-public-access','false','--public-network-access','Disabled',
            '--default-action','Deny','--allow-shared-key-access','true','--tags'
        ) + $tags) -Quiet | Out-Null
    }
    Invoke-AvdAzCli -Arguments @(
        'storage','account','update','--resource-group',$s.ResourceGroup,'--name',$s.StorageAccount,
        '--enable-files-aadkerb','true','--default-share-permission','StorageFileDataSmbShareContributor'
    ) -Quiet | Out-Null

    $shareId = "$($s.ResourceGroupId)/providers/Microsoft.Storage/storageAccounts/$($s.StorageAccount)/fileServices/default/shares/$($s.FileShareName)"
    if (-not (Invoke-AvdAzJson -Arguments @('resource','show','--ids',$shareId) -AllowFailure -Quiet)) {
        Invoke-AvdAzCli -Arguments @('storage','share-rm','create','--resource-group',$s.ResourceGroup,'--storage-account',$s.StorageAccount,'--name',$s.FileShareName,'--quota','100','--access-tier','TransactionOptimized') -Quiet | Out-Null
    }

    if (-not (Invoke-AvdAzJson -Arguments @('network','private-dns','zone','show','--resource-group',$s.ResourceGroup,'--name',$s.PrivateDnsZone) -AllowFailure -Quiet)) {
        Invoke-AvdAzCli -Arguments @('network','private-dns','zone','create','--resource-group',$s.ResourceGroup,'--name',$s.PrivateDnsZone) -Quiet | Out-Null
    }
    if (-not (Invoke-AvdAzJson -Arguments @('network','private-dns','link','vnet','show','--resource-group',$s.ResourceGroup,'--zone-name',$s.PrivateDnsZone,'--name',$s.PrivateDnsLinkName) -AllowFailure -Quiet)) {
        Invoke-AvdAzCli -Arguments @('network','private-dns','link','vnet','create','--resource-group',$s.ResourceGroup,'--zone-name',$s.PrivateDnsZone,'--name',$s.PrivateDnsLinkName,'--virtual-network',$s.VnetName,'--registration-enabled','false') -Quiet | Out-Null
    }
    $storageId = "$($s.ResourceGroupId)/providers/Microsoft.Storage/storageAccounts/$($s.StorageAccount)"
    if (-not (Invoke-AvdAzJson -Arguments @('network','private-endpoint','show','--resource-group',$s.ResourceGroup,'--name',$s.PrivateEndpointName) -AllowFailure -Quiet)) {
        Invoke-AvdAzCli -Arguments @(
            'network','private-endpoint','create','--resource-group',$s.ResourceGroup,'--name',$s.PrivateEndpointName,
            '--location',$s.Location,'--vnet-name',$s.VnetName,'--subnet',$s.PrivateEndpointSubnetName,
            '--private-connection-resource-id',$storageId,'--group-id','file','--connection-name','fslogix-file-connection'
        ) -Quiet | Out-Null
    }
    if (-not (Invoke-AvdAzJson -Arguments @('network','private-endpoint','dns-zone-group','show','--resource-group',$s.ResourceGroup,'--endpoint-name',$s.PrivateEndpointName,'--name','fslogix-dns') -AllowFailure -Quiet)) {
        Invoke-AvdAzCli -Arguments @(
            'network','private-endpoint','dns-zone-group','create','--resource-group',$s.ResourceGroup,'--endpoint-name',$s.PrivateEndpointName,
            '--name','fslogix-dns','--private-dns-zone',$s.PrivateDnsZone,'--zone-name','file'
        ) -Quiet | Out-Null
    }
    $dnsDeadline = [DateTime]::UtcNow.AddMinutes(3)
    do {
        $privateRecord = Invoke-AvdAzJson -Arguments @(
            'network','private-dns','record-set','a','show','--resource-group',$s.ResourceGroup,
            '--zone-name',$s.PrivateDnsZone,'--name',$s.StorageAccount
        ) -AllowFailure -Quiet
        if (-not $privateRecord) { Start-Sleep -Seconds 10 }
    } while (-not $privateRecord -and [DateTime]::UtcNow -lt $dnsDeadline)
    if (-not $privateRecord -or @($privateRecord.aRecords).Count -ne 1) {
        throw 'The private endpoint DNS zone group did not produce exactly one Azure Files A record.'
    }

    $displayName = "[Storage Account] $($s.StorageAccount).file.core.windows.net"
    $filter = [uri]::EscapeDataString("displayName eq '$displayName'")
    $storageApp = $null
    $deadline = [DateTime]::UtcNow.AddMinutes(5)
    do {
        $apps = Invoke-AvdGraphJson -Method GET -Path "applications?`$filter=$filter"
        $storageApp = $apps.value | Select-Object -First 1
        if (-not $storageApp) { Start-Sleep -Seconds 15 }
    } while (-not $storageApp -and [DateTime]::UtcNow -lt $deadline)
    if (-not $storageApp) { throw 'The autogenerated Microsoft Entra Kerberos storage application did not appear within five minutes.' }

    if (-not $s.StorageApplicationObjectId) {
        $script:State.StorageApplicationObjectId = [string]$storageApp.id
        $script:State.StorageApplicationCreated = $true
    }
    $spFilter = [uri]::EscapeDataString("appId eq '$($storageApp.appId)'")
    $storageSp = Invoke-AvdGraphJson -Method GET -Path "servicePrincipals?`$filter=$spFilter"
    $storageSpObject = $storageSp.value | Select-Object -First 1
    if (-not $storageSpObject) { throw 'The storage service principal was not found.' }
    $script:State.StorageServicePrincipalObjectId = [string]$storageSpObject.id

    $tags = @($storageApp.tags)
    if ($tags -notcontains 'kdc_enable_cloud_group_sids') { $tags += 'kdc_enable_cloud_group_sids' }
    $uris = @($storageApp.identifierUris)
    $privateFqdn = "$($s.StorageAccount).privatelink.file.core.windows.net"
    $publicFqdn = "$($s.StorageAccount).file.core.windows.net"
    foreach ($existingUri in @($storageApp.identifierUris)) {
        if ([string]$existingUri -like "*$publicFqdn") {
            $candidate = [string]$existingUri -replace [regex]::Escape($publicFqdn), $privateFqdn
            if ($uris -notcontains $candidate) { $uris += $candidate }
        }
    }
    Invoke-AvdGraphJson -Method PATCH -Path "applications/$($storageApp.id)" -Body @{ tags = $tags; identifierUris = $uris } | Out-Null

    $graphSpResponse = Invoke-AvdGraphJson -Method GET -Path "servicePrincipals?`$filter=appId eq '00000003-0000-0000-c000-000000000000'"
    $graphSp = $graphSpResponse.value | Select-Object -First 1
    if (-not $graphSp) { throw 'The Microsoft Graph service principal was not found.' }
    $grantFilter = [uri]::EscapeDataString("clientId eq '$($storageSpObject.id)' and resourceId eq '$($graphSp.id)'")
    $grantResponse = Invoke-AvdGraphJson -Method GET -Path "oauth2PermissionGrants?`$filter=$grantFilter"
    $grant = $grantResponse.value | Select-Object -First 1
    if (-not $grant) {
        $grant = Invoke-AvdGraphJson -Method POST -Path 'oauth2PermissionGrants' -Body @{
            clientId = [string]$storageSpObject.id
            consentType = 'AllPrincipals'
            resourceId = [string]$graphSp.id
            scope = 'openid profile User.Read'
        }
        $script:State.StorageConsentCreated = $true
    }
    $script:State.StorageConsentGrantId = [string]$grant.id
    Set-StageComplete -Name 'Storage'
}

function Invoke-HostPool {
    Ensure-State
    Write-Host 'Stage: Host pool, application group, and workspace' -ForegroundColor Cyan
    $s = $script:State
    $tags = @{ workshop = 'azure-virtual-desktop-deep-dive'; 'workshop-session' = [string]$s.SessionId; environment = 'lab'; owner = 'beyondcloudwithchriz' }

    $hostPoolUri = "https://management.azure.com$($s.HostPoolId)?api-version=$($s.AvdApiVersion)"
    $existing = Invoke-AvdAzRestJson -Method GET -Uri $hostPoolUri -AllowFailure -Quiet
    if (-not $existing) {
        $expiry = [DateTime]::UtcNow.AddHours(24).ToString('o')
        $hostPool = Invoke-AvdAzRestJson -Method PUT -Uri $hostPoolUri -Body @{
            location = [string]$s.Location
            tags = $tags
            properties = @{
                friendlyName = 'Azure Workshopz AVD pooled host pool'
                description = 'Pooled, Entra-joined AVD host pool for the deep-dive workshop.'
                hostPoolType = 'Pooled'
                loadBalancerType = 'DepthFirst'
                preferredAppGroupType = 'Desktop'
                maxSessionLimit = 10
                startVMOnConnect = $true
                validationEnvironment = $false
                customRdpProperty = 'targetisaadjoined:i:1;enablerdsaadauth:i:1;redirectclipboard:i:1;drivestoredirect:s:;'
                registrationInfo = @{ expirationTime = $expiry; registrationTokenOperation = 'Update' }
                agentUpdate = @{
                    type = 'Scheduled'
                    useSessionHostLocalTime = $false
                    maintenanceWindowTimeZone = 'W. Europe Standard Time'
                    maintenanceWindows = @(@{ dayOfWeek = 'Saturday'; hour = 2 }, @{ dayOfWeek = 'Sunday'; hour = 2 })
                }
            }
        } -Quiet
        $script:State.RegistrationToken = [string]$hostPool.properties.registrationInfo.token
    }
    if (-not $script:State.RegistrationToken) {
        $token = Invoke-AvdAzRestJson -Method POST -Uri "https://management.azure.com$($s.HostPoolId)/retrieveRegistrationToken?api-version=$($s.AvdApiVersion)" -Body @{} -Quiet
        $script:State.RegistrationToken = [string]$token.token
    }

    $applicationGroupUri = "https://management.azure.com$($s.ApplicationGroupId)?api-version=$($s.AvdApiVersion)"
    if (-not (Invoke-AvdAzRestJson -Method GET -Uri $applicationGroupUri -AllowFailure -Quiet)) {
        Invoke-AvdAzRestJson -Method PUT -Uri $applicationGroupUri -Body @{
            location = [string]$s.Location
            tags = $tags
            properties = @{
                hostPoolArmPath = [string]$s.HostPoolId
                applicationGroupType = 'Desktop'
                friendlyName = 'Azure Workshopz pooled desktop'
                description = 'Desktop application group for the Azure Virtual Desktop deep-dive workshop.'
            }
        } -Quiet | Out-Null
    }

    $workspaceUri = "https://management.azure.com$($s.WorkspaceId)?api-version=$($s.AvdApiVersion)"
    if (-not (Invoke-AvdAzRestJson -Method GET -Uri $workspaceUri -AllowFailure -Quiet)) {
        Invoke-AvdAzRestJson -Method PUT -Uri $workspaceUri -Body @{
            location = [string]$s.Location
            tags = $tags
            properties = @{
                friendlyName = 'Azure Workshopz AVD workspace'
                description = 'Published workspace for the AVD deep-dive workshop.'
                applicationGroupReferences = @([string]$s.ApplicationGroupId)
                publicNetworkAccess = 'Enabled'
            }
        } -Quiet | Out-Null
    }
    Set-StageComplete -Name 'HostPool'
}

function Invoke-SessionHosts {
    Ensure-State
    Write-Host 'Stage: Entra-joined Trusted Launch session hosts' -ForegroundColor Cyan
    $s = $script:State
    if (-not $s.RegistrationToken) { throw 'Host-pool registration token is missing. Run Stage HostPool.' }
    $tags = @('workshop=azure-virtual-desktop-deep-dive',"workshop-session=$($s.SessionId)",'environment=lab','owner=beyondcloudwithchriz','workload=avd-session-host')

    $dscSettingsPath = Join-Path $script:Paths.RuntimeRoot 'dsc-settings.json'
    $dscProtectedPath = Join-Path $script:Paths.RuntimeRoot 'dsc-protected-settings.json'
    $dscSettings = @{
        modulesUrl = 'https://raw.githubusercontent.com/Azure/RDS-Templates/master/ARM-wvd-templates/DSC/Configuration.zip'
        configurationFunction = 'Configuration.ps1\AddSessionHost'
        properties = @{
            hostPoolName = [string]$s.HostPoolName
            aadJoin = $true
            sessionHostConfigurationLastUpdateTime = [DateTime]::UtcNow.ToString('o')
        }
    }
    $dscProtected = @{ properties = @{ registrationInfoToken = [string]$s.RegistrationToken } }
    [System.IO.File]::WriteAllText($dscSettingsPath, ($dscSettings | ConvertTo-Json -Depth 20), [System.Text.UTF8Encoding]::new($false))
    [System.IO.File]::WriteAllText($dscProtectedPath, ($dscProtected | ConvertTo-Json -Depth 20), [System.Text.UTF8Encoding]::new($false))

    for ($index = 0; $index -lt [int]$s.HostCount; $index++) {
        $vmName = [string]$s.VmNames[$index]
        $nicName = [string]$s.NicNames[$index]
        $computerName = [string]$s.ComputerNames[$index]
        if (-not (Invoke-AvdAzJson -Arguments @('network','nic','show','--resource-group',$s.ResourceGroup,'--name',$nicName) -AllowFailure -Quiet)) {
            Invoke-AvdAzCli -Arguments (@(
                'network','nic','create','--resource-group',$s.ResourceGroup,'--name',$nicName,'--location',$s.Location,
                '--vnet-name',$s.VnetName,'--subnet',$s.SessionSubnetName,'--network-security-group',$s.NsgName,'--tags'
            ) + $tags) -Quiet | Out-Null
        }
        if (-not (Invoke-AvdAzJson -Arguments @('vm','show','--resource-group',$s.ResourceGroup,'--name',$vmName) -AllowFailure -Quiet)) {
            Invoke-AvdAzCli -Arguments (@(
                'vm','create','--resource-group',$s.ResourceGroup,'--name',$vmName,'--location',$s.Location,
                '--nics',$nicName,'--computer-name',$computerName,
                '--image','MicrosoftWindowsDesktop:windows-11:win11-24h2-avd:latest','--size',$s.VmSize,
                '--admin-username','avdlocaladmin','--admin-password',[string]$s.LocalAdminPassword,
                '--license-type','Windows_Client','--security-type','TrustedLaunch','--enable-secure-boot','true','--enable-vtpm','true',
                '--assign-identity','--storage-sku','StandardSSD_LRS','--tags'
            ) + $tags) -Quiet | Out-Null
        }

        $aadExtension = Invoke-AvdAzJson -Arguments @('vm','extension','show','--resource-group',$s.ResourceGroup,'--vm-name',$vmName,'--name','AADLoginForWindows') -AllowFailure -Quiet
        if (-not $aadExtension -or [string]$aadExtension.provisioningState -ne 'Succeeded') {
            Invoke-AvdAzCli -Arguments @(
                'vm','extension','set','--resource-group',$s.ResourceGroup,'--vm-name',$vmName,
                '--publisher','Microsoft.Azure.ActiveDirectory','--name','AADLoginForWindows','--version','1.0'
            ) -Quiet | Out-Null
        }

        $dscExtension = Invoke-AvdAzJson -Arguments @('vm','extension','show','--resource-group',$s.ResourceGroup,'--vm-name',$vmName,'--name','Microsoft.PowerShell.DSC') -AllowFailure -Quiet
        if (-not $dscExtension -or [string]$dscExtension.provisioningState -ne 'Succeeded') {
            Invoke-AvdAzCli -Arguments @(
                'vm','extension','set','--resource-group',$s.ResourceGroup,'--vm-name',$vmName,
                '--publisher','Microsoft.Powershell','--name','DSC','--extension-instance-name','Microsoft.PowerShell.DSC',
                '--version','2.83','--settings',"@$dscSettingsPath",'--protected-settings',"@$dscProtectedPath"
            ) -Quiet | Out-Null
        }
    }

    Wait-AvdSessionHosts -State $s -TimeoutSeconds 2400 | Out-Null
    Set-StageComplete -Name 'SessionHosts'
}

function Invoke-FSLogix {
    Ensure-State
    Write-Host 'Stage: FSLogix configuration and private Azure Files verification' -ForegroundColor Cyan
    $s = $script:State
    $configurationScript = Join-Path $PSScriptRoot 'session-host\Configure-AvdSessionHost.ps1'
    if (-not (Test-Path -LiteralPath $configurationScript)) { throw "Missing $configurationScript" }
    $storageKey = [string](Invoke-AvdAzJson -Arguments @('storage','account','keys','list','--resource-group',$s.ResourceGroup,'--account-name',$s.StorageAccount,'--query','[0].value') -Quiet)

    foreach ($vmName in @($s.VmNames)) {
        $result = Invoke-AvdAzJson -Arguments @(
            'vm','run-command','invoke','--resource-group',$s.ResourceGroup,'--name',[string]$vmName,
            '--command-id','RunPowerShellScript','--scripts',"@$configurationScript",
            '--parameters',"StorageAccountName=$($s.StorageAccount)","ShareName=$($s.FileShareName)","StorageKey=$storageKey"
        ) -Quiet
        $privateResultPath = Join-Path $script:Paths.RuntimeRoot ("{0}-fslogix-config.json" -f $vmName)
        [System.IO.File]::WriteAllText($privateResultPath, ($result | ConvertTo-Json -Depth 30), [System.Text.UTF8Encoding]::new($false))
        $stdout = [string](@($result.value | Where-Object { [string]$_.code -like '*StdOut*' })[0].message)
        $stderr = [string](@($result.value | Where-Object { [string]$_.code -like '*StdErr*' })[0].message)
        if (-not [string]::IsNullOrWhiteSpace($stderr)) { throw "FSLogix configuration failed on $vmName. See the private Run Command result." }
        if ([string]::IsNullOrWhiteSpace($stdout)) { throw "FSLogix configuration returned no proof on $vmName." }
        $proof = $stdout.Trim() | ConvertFrom-Json
        if ([string]$proof.ConfigurationStatus -ne 'Applied' -or -not [bool]$proof.PrivateDnsResolved -or -not [bool]$proof.SmbPort445Reachable -or -not [bool]$proof.CreatorOwnerAclApplied) {
            throw "FSLogix configuration proof was incomplete on $vmName."
        }
    }
    foreach ($vmName in @($s.VmNames)) {
        Invoke-AvdAzCli -Arguments @('vm','restart','--resource-group',$s.ResourceGroup,'--name',[string]$vmName) -Quiet | Out-Null
    }
    Wait-AvdSessionHosts -State $s -TimeoutSeconds 1800 | Out-Null
    Set-StageComplete -Name 'FSLogix'
}

function Invoke-Monitoring {
    Ensure-State
    Write-Host 'Stage: Azure Monitor, DCR, diagnostics, and AVD Insights data sources' -ForegroundColor Cyan
    $s = $script:State
    $tags = @('workshop=azure-virtual-desktop-deep-dive',"workshop-session=$($s.SessionId)",'environment=lab','owner=beyondcloudwithchriz')
    if (-not (Invoke-AvdAzJson -Arguments @('monitor','log-analytics','workspace','show','--resource-group',$s.ResourceGroup,'--workspace-name',$s.LogAnalyticsName) -AllowFailure -Quiet)) {
        Invoke-AvdAzCli -Arguments (@('monitor','log-analytics','workspace','create','--resource-group',$s.ResourceGroup,'--workspace-name',$s.LogAnalyticsName,'--location',$s.Location,'--retention-time','30','--tags') + $tags) -Quiet | Out-Null
    }
    $law = Invoke-AvdAzJson -Arguments @('monitor','log-analytics','workspace','show','--resource-group',$s.ResourceGroup,'--workspace-name',$s.LogAnalyticsName) -Quiet

    $dcrUri = "https://management.azure.com$($s.DcrId)?api-version=$($s.MonitorApiVersion)"
    $dcr = Invoke-AvdAzRestJson -Method GET -Uri $dcrUri -AllowFailure -Quiet
    if (-not $dcr) {
        Invoke-AvdAzRestJson -Method PUT -Uri $dcrUri -Body @{
            location = [string]$s.Location
            kind = 'Windows'
            tags = @{ workshop = 'azure-virtual-desktop-deep-dive'; 'workshop-session' = [string]$s.SessionId }
            properties = @{
                dataSources = @{
                    performanceCounters = @(@{
                        name = 'avd-performance'
                        streams = @('Microsoft-Perf')
                        samplingFrequencyInSeconds = 60
                        counterSpecifiers = @(
                            '\Processor Information(_Total)\% Processor Time',
                            '\Memory\Available Mbytes',
                            '\LogicalDisk(_Total)\% Free Space',
                            '\Network Interface(*)\Bytes Total/sec'
                        )
                    })
                    windowsEventLogs = @(@{
                        name = 'avd-events'
                        streams = @('Microsoft-Event')
                        xPathQueries = @(
                            'Microsoft-FSLogix-Apps/Operational!*[System[(Level=1 or Level=2 or Level=3 or Level=4)]]',
                            'Microsoft-Windows-TerminalServices-LocalSessionManager/Operational!*[System[(Level=1 or Level=2 or Level=3 or Level=4)]]',
                            'Application!*[System[Provider[@Name=''AVDWorkshop'']]]',
                            'System!*[System[(Level=1 or Level=2 or Level=3)]]'
                        )
                    })
                }
                destinations = @{ logAnalytics = @(@{ name = 'avd-law'; workspaceResourceId = [string]$law.id }) }
                dataFlows = @(
                    @{ streams = @('Microsoft-Perf'); destinations = @('avd-law'); transformKql = 'source' },
                    @{ streams = @('Microsoft-Event'); destinations = @('avd-law'); transformKql = 'source' }
                )
            }
        } -Quiet | Out-Null
    }

    for ($index = 0; $index -lt [int]$s.HostCount; $index++) {
        $vmName = [string]$s.VmNames[$index]
        $vmId = "$($s.ResourceGroupId)/providers/Microsoft.Compute/virtualMachines/$vmName"
        $ama = Invoke-AvdAzJson -Arguments @('vm','extension','show','--resource-group',$s.ResourceGroup,'--vm-name',$vmName,'--name','AzureMonitorWindowsAgent') -AllowFailure -Quiet
        if (-not $ama -or [string]$ama.provisioningState -ne 'Succeeded') {
            Invoke-AvdAzCli -Arguments @(
                'vm','extension','set','--resource-group',$s.ResourceGroup,'--vm-name',$vmName,
                '--publisher','Microsoft.Azure.Monitor','--name','AzureMonitorWindowsAgent','--enable-auto-upgrade','true'
            ) -Quiet | Out-Null
        }
        $associationId = "$vmId/providers/Microsoft.Insights/dataCollectionRuleAssociations/avd-sessionhosts"
        $associationUri = "https://management.azure.com$associationId`?api-version=$($s.MonitorApiVersion)"
        if (-not (Invoke-AvdAzRestJson -Method GET -Uri $associationUri -AllowFailure -Quiet)) {
            Invoke-AvdAzRestJson -Method PUT -Uri $associationUri -Body @{ properties = @{ dataCollectionRuleId = [string]$s.DcrId } } -Quiet | Out-Null
        }
    }

    foreach ($resource in @(
        @{ Id = [string]$s.HostPoolId; Name = 'diag-host-pool' },
        @{ Id = [string]$s.ApplicationGroupId; Name = 'diag-application-group' },
        @{ Id = [string]$s.WorkspaceId; Name = 'diag-workspace' }
    )) {
        $diagUri = "https://management.azure.com$($resource.Id)/providers/Microsoft.Insights/diagnosticSettings/$($resource.Name)?api-version=$($s.DiagnosticsApiVersion)"
        Invoke-AvdAzRestJson -Method PUT -Uri $diagUri -Body @{
            properties = @{
                workspaceId = [string]$law.id
                logs = @(@{ categoryGroup = 'allLogs'; enabled = $true })
                metrics = @()
            }
        } -Quiet | Out-Null
    }
    Set-StageComplete -Name 'Monitoring'
}

function New-ScalingSchedule {
    param([switch]$TestWindow)

    if (-not $TestWindow) {
        return @{
            name = 'production-week'
            daysOfWeek = @('Monday','Tuesday','Wednesday','Thursday','Friday','Saturday','Sunday')
            rampUpStartTime = @{ hour = 6; minute = 0 }
            rampUpLoadBalancingAlgorithm = 'DepthFirst'
            rampUpMinimumHostsPct = 50
            rampUpCapacityThresholdPct = 75
            peakStartTime = @{ hour = 8; minute = 0 }
            peakLoadBalancingAlgorithm = 'DepthFirst'
            rampDownStartTime = @{ hour = 18; minute = 0 }
            rampDownLoadBalancingAlgorithm = 'DepthFirst'
            rampDownMinimumHostsPct = 50
            rampDownCapacityThresholdPct = 50
            rampDownForceLogoffUsers = $false
            rampDownWaitTimeMinutes = 15
            rampDownNotificationMessage = 'This workshop host is entering the autoscale ramp-down phase.'
            offPeakStartTime = @{ hour = 22; minute = 0 }
            offPeakLoadBalancingAlgorithm = 'DepthFirst'
        }
    }

    $localNow = [TimeZoneInfo]::ConvertTimeBySystemTimeZoneId([DateTime]::UtcNow, 'W. Europe Standard Time')
    $rampDown = $localNow.AddMinutes(3)
    $rampUp = $rampDown.AddHours(-2)
    $peak = $rampDown.AddHours(-1)
    $offPeak = $rampDown.AddMinutes(45)
    @{
        name = 'autoscale-proof-window'
        daysOfWeek = @('Monday','Tuesday','Wednesday','Thursday','Friday','Saturday','Sunday')
        rampUpStartTime = @{ hour = $rampUp.Hour; minute = $rampUp.Minute }
        rampUpLoadBalancingAlgorithm = 'DepthFirst'
        rampUpMinimumHostsPct = 100
        rampUpCapacityThresholdPct = 75
        peakStartTime = @{ hour = $peak.Hour; minute = $peak.Minute }
        peakLoadBalancingAlgorithm = 'DepthFirst'
        rampDownStartTime = @{ hour = $rampDown.Hour; minute = $rampDown.Minute }
        rampDownLoadBalancingAlgorithm = 'DepthFirst'
        rampDownMinimumHostsPct = 50
        rampDownCapacityThresholdPct = 50
        rampDownForceLogoffUsers = $false
        rampDownWaitTimeMinutes = 0
        rampDownNotificationMessage = 'Workshop autoscale proof window.'
        offPeakStartTime = @{ hour = $offPeak.Hour; minute = $offPeak.Minute }
        offPeakLoadBalancingAlgorithm = 'DepthFirst'
    }
}

function Set-ScalingPlanSchedule {
    param([Parameter(Mandatory)]$Schedule)
    $s = $script:State
    $uri = "https://management.azure.com$($s.ScalingPlanId)?api-version=$($s.AvdApiVersion)"
    Invoke-AvdAzRestJson -Method PUT -Uri $uri -Body @{
        location = [string]$s.Location
        tags = @{ workshop = 'azure-virtual-desktop-deep-dive'; 'workshop-session' = [string]$s.SessionId; environment = 'lab' }
        properties = @{
            friendlyName = 'Azure Workshopz AVD power management'
            description = 'Keeps 50 percent of the two-host pool available and deallocates unused capacity.'
            timeZone = 'W. Europe Standard Time'
            hostPoolType = 'Pooled'
            exclusionTag = 'excludeFromScaling'
            schedules = @($Schedule)
            hostPoolReferences = @(@{ hostPoolArmPath = [string]$s.HostPoolId; scalingPlanEnabled = $true })
        }
    } -Quiet | Out-Null
}

function Invoke-Scaling {
    Ensure-State
    Write-Host 'Stage: Power-management scaling plan' -ForegroundColor Cyan
    $s = $script:State
    if (-not [bool]$s.AutoscaleRoleBaselinePresent -and -not [bool]$s.AutoscaleRoleCreated) {
        $assignment = Invoke-AvdAzJson -Arguments @(
            'role','assignment','create','--assignee-object-id',[string]$s.AvdServicePrincipalObjectId,
            '--assignee-principal-type','ServicePrincipal','--role','Desktop Virtualization Power On Off Contributor',
            '--scope',[string]$s.Scope
        ) -Quiet
        $script:State.AutoscaleRoleAssignmentId = [string]$assignment.id
        $script:State.AutoscaleRoleCreated = $true
        Save-State
    }

    $productionSchedule = New-ScalingSchedule
    Set-ScalingPlanSchedule -Schedule $productionSchedule
    $script:State.OriginalScalingSchedule = $productionSchedule

    $law = Invoke-AvdAzJson -Arguments @('monitor','log-analytics','workspace','show','--resource-group',$s.ResourceGroup,'--workspace-name',$s.LogAnalyticsName) -AllowFailure -Quiet
    if ($law) {
        $diagUri = "https://management.azure.com$($s.ScalingPlanId)/providers/Microsoft.Insights/diagnosticSettings/diag-scaling-plan?api-version=$($s.DiagnosticsApiVersion)"
        Invoke-AvdAzRestJson -Method PUT -Uri $diagUri -Body @{
            properties = @{ workspaceId = [string]$law.id; logs = @(@{ categoryGroup = 'allLogs'; enabled = $true }); metrics = @() }
        } -Quiet | Out-Null
    }
    Set-StageComplete -Name 'Scaling'
}

function Invoke-TestProfile {
    Ensure-State
    Write-Host 'Stage: Deliberate FSLogix profile test' -ForegroundColor Cyan
    $s = $script:State
    if (-not $s.AssignedLicenseSkuId) {
        if ($ProfileEvidenceMode -eq 'RequireInteractive') { throw 'Interactive evidence was required, but no eligible AVD license is available.' }
        $script:State.ProfilePersistenceStatus = 'NotTested-LicenseUnavailable'
        $script:State.ProfileFirstHost = $null
        $script:State.ProfileSecondHost = $null
        $script:State.ProfileVhdxCount = $null
        Set-StageComplete -Name 'TestProfile'
        Write-Warning 'FSLogix is configured, but profile persistence was not live-verified because no eligible AVD license is available.'
        return
    }

    $instructions = [ordered]@{
        Status = 'InteractiveLoginRequired'
        WebClient = 'https://client.wvd.microsoft.com/arm/webclient/index.html'
        PlannedFirstHost = [string]$s.ComputerNames[0]
        PlannedSecondHost = [string]$s.ComputerNames[1]
        TimeoutMinutes = [int]$s.InteractiveTimeoutMinutes
        Note = 'Use the private test-user credentials from runtime state. Never copy them into public evidence.'
    }
    $instructionsPath = Join-Path $script:Paths.RuntimeRoot 'interactive-profile-test.json'
    [System.IO.File]::WriteAllText($instructionsPath, ($instructions | ConvertTo-Json), [System.Text.UTF8Encoding]::new($false))
    throw "Interactive profile proof is ready. Complete the web-client host-switch sequence, then run Get-AzureAvdEvidence.ps1 -CapturePrivateProfileProof. Instructions: $instructionsPath"
}

function Invoke-TestScaling {
    Ensure-State
    Write-Host 'Stage: Deliberate autoscale test' -ForegroundColor Cyan
    $s = $script:State

    $sessionsUri = "https://management.azure.com$($s.HostPoolId)/userSessions?api-version=$($s.AvdApiVersion)"
    $sessions = Invoke-AvdAzRestJson -Method GET -Uri $sessionsUri -AllowFailure -Quiet
    if (@($sessions.value).Count -gt 0) { throw 'Autoscale proof requires an empty host pool. Log off every AVD session and rerun.' }

    $existingStates = @($s.VmNames | ForEach-Object { Get-AvdVmPowerState -ResourceGroup $s.ResourceGroup -Name ([string]$_) })
    if (@($existingStates | Where-Object { $_ -eq 'VM deallocated' }).Count -eq 1 -and @($existingStates | Where-Object { $_ -eq 'VM running' }).Count -eq 1) {
        $eventDeadline = [DateTime]::UtcNow.AddMinutes(10)
        $existingEvent = $null
        do {
            $recentEvents = Invoke-AvdAzJson -Arguments @(
                'monitor','activity-log','list','--resource-group',$s.ResourceGroup,'--offset','2h','--max-events','500'
            ) -Quiet
            $existingEvent = $recentEvents | Where-Object {
                [string]$_.operationName.value -match '(?i)virtualMachines/(deallocate|powerOff)/action' -and
                [string]$_.status.value -match '(?i)Succeeded'
            } | Sort-Object eventTimestamp -Descending | Select-Object -First 1
            if (-not $existingEvent) { Start-Sleep -Seconds 30 }
        } while (-not $existingEvent -and [DateTime]::UtcNow -lt $eventDeadline)
        if ($existingEvent) {
            Set-ScalingPlanSchedule -Schedule (New-ScalingSchedule)
            $script:State.ScalingTestStatus = 'Passed-OneRunning-OneDeallocated'
            $script:State.ScalingActivityCorrelationId = [string]$existingEvent.correlationId
            Set-StageComplete -Name 'TestScaling'
            return
        }
    }

    Wait-AvdSessionHosts -State $s -TimeoutSeconds 1200 | Out-Null
    foreach ($vmName in @($s.VmNames)) {
        Invoke-AvdAzCli -Arguments @('vm','start','--resource-group',$s.ResourceGroup,'--name',[string]$vmName,'--no-wait') -Quiet | Out-Null
    }
    foreach ($vmName in @($s.VmNames)) {
        Wait-AvdVmPowerState -ResourceGroup $s.ResourceGroup -Name ([string]$vmName) -Expected @('VM running') -TimeoutSeconds 1200 | Out-Null
    }

    $testStarted = [DateTime]::UtcNow
    $testSchedule = New-ScalingSchedule -TestWindow
    try {
        Set-ScalingPlanSchedule -Schedule $testSchedule
        $deadline = [DateTime]::UtcNow.AddMinutes(45)
        $states = @()
        do {
            $states = @($s.VmNames | ForEach-Object { Get-AvdVmPowerState -ResourceGroup $s.ResourceGroup -Name ([string]$_) })
            if (@($states | Where-Object { $_ -eq 'VM deallocated' }).Count -eq 1 -and @($states | Where-Object { $_ -eq 'VM running' }).Count -eq 1) { break }
            Start-Sleep -Seconds 30
        } while ([DateTime]::UtcNow -lt $deadline)

        if (@($states | Where-Object { $_ -eq 'VM deallocated' }).Count -ne 1 -or @($states | Where-Object { $_ -eq 'VM running' }).Count -ne 1) {
            throw "Autoscale didn't produce the required one-running/one-deallocated boundary. States: $($states -join ', ')"
        }

        $eventDeadline = [DateTime]::UtcNow.AddMinutes(10)
        $events = $null
        do {
            $activity = Invoke-AvdAzJson -Arguments @(
                'monitor','activity-log','list','--resource-group',$s.ResourceGroup,
                '--start-time',$testStarted.AddMinutes(-5).ToString('o'),'--max-events','500'
            ) -Quiet
            $events = $activity | Where-Object {
                [string]$_.operationName.value -match '(?i)virtualMachines/(deallocate|powerOff)/action' -and
                [string]$_.status.value -match '(?i)Succeeded'
            } | Sort-Object eventTimestamp -Descending | Select-Object -First 1
            if (-not $events) { Start-Sleep -Seconds 30 }
        } while (-not $events -and [DateTime]::UtcNow -lt $eventDeadline)
        if (-not $events) { throw 'The one-running/one-deallocated boundary was real, but its Activity Log completion event did not arrive within ten minutes.' }
        $script:State.ScalingTestStatus = 'Passed-OneRunning-OneDeallocated'
        $script:State.ScalingActivityCorrelationId = [string]$events.correlationId
    }
    finally {
        Set-ScalingPlanSchedule -Schedule (New-ScalingSchedule)
    }
    Set-StageComplete -Name 'TestScaling'
}

function Invoke-Validate {
    Ensure-State
    Write-Host 'Stage: Validate' -ForegroundColor Cyan
    $s = $script:State
    Assert-ExpectedResourceGroup

    $hostsResponse = Invoke-AvdAzRestJson -Method GET -Uri "https://management.azure.com$($s.HostPoolId)/sessionHosts?api-version=$($s.AvdApiVersion)" -Quiet
    $hosts = @($hostsResponse.value)
    if ($hosts.Count -ne [int]$s.HostCount) { throw 'The registered session-host count is incorrect.' }
    $allowedHostStatuses = @('Available','Unavailable')
    if ([string]$s.ScalingTestStatus -eq 'Passed-OneRunning-OneDeallocated') {
        $allowedHostStatuses += 'Shutdown'
    }
    if (@($hosts | Where-Object { $_.properties.status -notin $allowedHostStatuses }).Count -gt 0) { throw 'One or more session hosts has an unexpected AVD status.' }

    $powerStates = @()
    foreach ($index in 0..([int]$s.HostCount - 1)) {
        $vm = Invoke-AvdAzJson -Arguments @('vm','show','--resource-group',$s.ResourceGroup,'--name',[string]$s.VmNames[$index]) -Quiet
        $instanceView = Invoke-AvdAzJson -Arguments @('vm','get-instance-view','--resource-group',$s.ResourceGroup,'--name',[string]$s.VmNames[$index]) -Quiet
        $powerStates += [string](@($instanceView.instanceView.statuses | Where-Object { [string]$_.code -like 'PowerState/*' } | Select-Object -First 1).displayStatus)
        $nic = Invoke-AvdAzJson -Arguments @('network','nic','show','--resource-group',$s.ResourceGroup,'--name',[string]$s.NicNames[$index]) -Quiet
        $ipConfiguration = $nic.ipConfigurations[0]
        if ($ipConfiguration.PSObject.Properties.Name -contains 'publicIPAddress' -and $ipConfiguration.publicIPAddress) {
            throw "VM $($s.VmNames[$index]) unexpectedly has a public IP configuration."
        }
        if ([string]$vm.securityProfile.securityType -ne 'TrustedLaunch' -or -not [bool]$vm.securityProfile.uefiSettings.secureBootEnabled -or -not [bool]$vm.securityProfile.uefiSettings.vTpmEnabled) {
            throw "VM $($s.VmNames[$index]) isn't fully configured for Trusted Launch."
        }
    }
    if ([string]$s.ScalingTestStatus -eq 'Passed-OneRunning-OneDeallocated') {
        if (@($powerStates | Where-Object { $_ -eq 'VM running' }).Count -ne 1 -or @($powerStates | Where-Object { $_ -eq 'VM deallocated' }).Count -ne 1) {
            throw 'The tested autoscale boundary is no longer one running and one deallocated VM.'
        }
    }

    $storage = Invoke-AvdAzJson -Arguments @('storage','account','show','--resource-group',$s.ResourceGroup,'--name',$s.StorageAccount) -Quiet
    if ([string]$storage.azureFilesIdentityBasedAuthentication.directoryServiceOptions -ne 'AADKERB') {
        throw 'Microsoft Entra Kerberos is not enabled on Azure Files.'
    }
    $pe = Invoke-AvdAzJson -Arguments @('network','private-endpoint','show','--resource-group',$s.ResourceGroup,'--name',$s.PrivateEndpointName) -Quiet
    if ([string]$pe.privateLinkServiceConnections[0].privateLinkServiceConnectionState.status -ne 'Approved') { throw 'The Azure Files private endpoint is not approved.' }

    $scaling = Invoke-AvdAzRestJson -Method GET -Uri "https://management.azure.com$($s.ScalingPlanId)?api-version=$($s.AvdApiVersion)" -Quiet
    if ([int]$scaling.properties.schedules[0].rampDownMinimumHostsPct -ne 50) { throw 'Production scaling schedule was not restored to 50 percent minimum.' }

    $report = [ordered]@{
        ValidatedUtc = [DateTime]::UtcNow.ToString('o')
        RegisteredSessionHostCount = $hosts.Count
        SessionHostsWithPublicIpCount = 0
        TrustedLaunchHostCount = [int]$s.HostCount
        AzureFilesEntraKerberos = $true
        PrivateEndpointApproved = $true
        FslogixConfigurationAppliedHostCount = @($s.VmNames).Count
        ProfilePersistenceStatus = [string]$s.ProfilePersistenceStatus
        ScalingTestStatus = [string]$s.ScalingTestStatus
        ProductionMinimumHostsPercent = 50
    }
    $privateReport = Join-Path $script:Paths.RuntimeRoot 'validation-private.json'
    [System.IO.File]::WriteAllText($privateReport, ($report | ConvertTo-Json -Depth 10), [System.Text.UTF8Encoding]::new($false))
    Set-StageComplete -Name 'Validate'
    Write-Host 'Validation passed.' -ForegroundColor Green
}

if ($Stage -eq 'All') {
    Invoke-Preflight
    if ($CreateTestUser) { Invoke-Identity }
    else { Write-Warning 'Identity stage skipped in All because -CreateTestUser was not supplied.' }
    Invoke-Foundation
    Invoke-Storage
    Invoke-HostPool
    Invoke-SessionHosts
    Invoke-FSLogix
    Invoke-Monitoring
    Invoke-Scaling
    Invoke-TestProfile
    Invoke-TestScaling
    Invoke-Validate
}
else {
    switch ($Stage) {
        'Preflight'    { Invoke-Preflight }
        'Identity'     { Invoke-Identity }
        'Foundation'   { Invoke-Foundation }
        'Storage'      { Invoke-Storage }
        'HostPool'     { Invoke-HostPool }
        'SessionHosts' { Invoke-SessionHosts }
        'FSLogix'      { Invoke-FSLogix }
        'Monitoring'   { Invoke-Monitoring }
        'Scaling'      { Invoke-Scaling }
        'TestProfile'  { Invoke-TestProfile }
        'TestScaling'  { Invoke-TestScaling }
        'Validate'     { Invoke-Validate }
    }
}
