# AVD App Group assignment map (Unassign -> Assign)

$subscriptionId = "xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx"
$roleName = "Desktop Virtualization User"

# 解除リスト: アプリケーショングループごとに対象を配列で指定
$unassignMap = @(
    @{
        ResourceGroup = "RG1"
        AppGroupName  = "RG1-DAG"
        Groups        = @("AVD Group1","AVD Group2")
        Users         = @("user01@xxx.onmicrosoft.com", "user02@xxx.onmicrosoft.com")
    },
    @{
        ResourceGroup = "RG2"
        AppGroupName  = "RG-2-DAG"
        Groups        = @("Entra-Group-A", "Entra-Group-C")
        Users         = @()
    }
)

# 割り当てリスト: アプリケーショングループごとに対象を配列で指定
$assignMap = @(
    @{
        ResourceGroup = "RG1"
        AppGroupName  = "RG1-DAG"
        Groups        = @("Entra-Group-D")
        Users         = @("user10@xxx.onmicrosoft.com")
    },
    @{
        ResourceGroup = "RG3"
        AppGroupName  = "RG3-DAG"
        Groups        = @("Entra-Group-A", "Entra-Group-C")
        Users         = @()
    }
)

function Write-Result {
    param([string]$Level, [string]$Message)
    Write-Output ("[{0}] {1}" -f $Level, $Message)
}

function Resolve-AppGroup {
    param([string]$ResourceGroup, [string]$AppGroupName)
    try {
        $resolved = Get-AzWvdApplicationGroup -ResourceGroupName $ResourceGroup -Name $AppGroupName -ErrorAction Stop
        if (-not $resolved) {
            Write-Result "SKIP" ("Application group not found: {0}" -f $AppGroupName)
            return $null
        }
        return $resolved
    } catch {
        Write-Result "SKIP" ("Application group not found: {0}" -f $AppGroupName)
        return $null
    }
}

function Resolve-User {
    param([string]$Upn)
    try {
        $resolved = $null
        if ($Upn -match "^[0-9a-fA-F-]{36}$") {
            $resolved = Get-AzADUser -ObjectId $Upn -ErrorAction Stop
        } else {
            $resolved = Get-AzADUser -UserPrincipalName $Upn -ErrorAction Stop
        }

        if (-not $resolved) {
            Write-Result "SKIP" ("Entra user not found: {0}" -f $Upn)
            return $null
        }
        if ($resolved -is [array]) {
            Write-Result "SKIP" ("Multiple Entra users matched: {0}" -f $Upn)
            return $null
        }
        return $resolved
    } catch {
        Write-Result "SKIP" ("Entra user not found: {0}. {1}" -f $Upn, $_.Exception.Message)
        return $null
    }
}

function Resolve-Group {
    param([string]$DisplayName)
    try {
        $resolved = $null
        if ($DisplayName -match "^[0-9a-fA-F-]{36}$") {
            $resolved = Get-AzADGroup -ObjectId $DisplayName -ErrorAction Stop
        } else {
            $resolved = Get-AzADGroup -DisplayName $DisplayName -ErrorAction Stop
        }

        if (-not $resolved) {
            Write-Result "SKIP" ("Entra group not found: {0}" -f $DisplayName)
            return $null
        }
        if ($resolved -is [array]) {
            $exact = $resolved | Where-Object { $_.DisplayName -eq $DisplayName }
            if ($exact.Count -eq 1) {
                return $exact[0]
            }
            Write-Result "SKIP" ("Multiple Entra groups matched: {0}" -f $DisplayName)
            return $null
        }
        return $resolved
    } catch {
        Write-Result "SKIP" ("Entra group not found: {0}. {1}" -f $DisplayName, $_.Exception.Message)
        return $null
    }
}

function Unassign-Principal {
    param([string]$PrincipalId, [string]$PrincipalName, [string]$Scope)
    if (-not $PrincipalId) { return }
    try {
        $existing = Get-AzRoleAssignment -ObjectId $PrincipalId -Scope $Scope -ErrorAction SilentlyContinue
        if ($existing) {
            $existing | ForEach-Object {
                Remove-AzRoleAssignment -Scope $_.Scope -ObjectId $_.ObjectId -RoleDefinitionName $_.RoleDefinitionName -ErrorAction Stop
            }
            Write-Result "OK" ("Unassigned: {0}" -f $PrincipalName)
        } else {
            Write-Result "SKIP" ("No assignment found: {0}" -f $PrincipalName)
        }
    } catch {
        Write-Result "ERROR" ("Unassign failed: {0}. {1}" -f $PrincipalName, $_.Exception.Message)
    }
}

function Assign-Principal {
    param([string]$PrincipalId, [string]$PrincipalName, [string]$Scope, [string]$RoleName)
    if (-not $PrincipalId) { return }
    try {
        New-AzRoleAssignment -ObjectId $PrincipalId -Scope $Scope -RoleDefinitionName $RoleName -ErrorAction Stop | Out-Null
        Write-Result "OK" ("Assigned: {0}" -f $PrincipalName)
    } catch {
        Write-Result "ERROR" ("Assign failed: {0}. {1}" -f $PrincipalName, $_.Exception.Message)
    }
}

# Connect
try {
    Write-Result "INFO" "Connecting to Azure with Managed Identity..."
    $context = Connect-AzAccount -Identity -ErrorAction Stop
    Write-Result "OK" ("Connected as: {0}" -f $context.Context.Account.Id)
    
    Write-Result "INFO" ("Selecting subscription: {0}" -f $subscriptionId)
    Select-AzSubscription -SubscriptionId $subscriptionId -ErrorAction Stop | Out-Null
    Write-Result "OK" "Subscription selected"
} catch {
    Write-Result "ERROR" ("Connection failed: {0}" -f $_.Exception.Message)
    throw
}

# ---- Unassign ----
Write-Result "INFO" "Starting unassignment process..."
foreach ($item in $unassignMap) {
        Write-Result "INFO" ("Processing AppGroup: {0}/{1}" -f $item.ResourceGroup, $item.AppGroupName)
    $appGroup = Resolve-AppGroup -ResourceGroup $item.ResourceGroup -AppGroupName $item.AppGroupName
    if (-not $appGroup) { continue }

    foreach ($g in $item.Groups) {
        $group = Resolve-Group -DisplayName $g
        if ($group) { Unassign-Principal -PrincipalId $group.Id -PrincipalName $g -Scope $appGroup.Id }
    }
    foreach ($u in $item.Users) {
        $user = Resolve-User -Upn $u
        if ($user) { Unassign-Principal -PrincipalId $user.Id -PrincipalName $u -Scope $appGroup.Id }
    }
}

# ---- Assign ----
Write-Result "INFO" "Starting assignment process..."
foreach ($item in $assignMap) {
        Write-Result "INFO" ("Processing AppGroup: {0}/{1}" -f $item.ResourceGroup, $item.AppGroupName)
    $appGroup = Resolve-AppGroup -ResourceGroup $item.ResourceGroup -AppGroupName $item.AppGroupName
    if (-not $appGroup) { continue }

    foreach ($g in $item.Groups) {
        $group = Resolve-Group -DisplayName $g
        if ($group) { Assign-Principal -PrincipalId $group.Id -PrincipalName $g -Scope $appGroup.Id -RoleName $roleName }
    }
    foreach ($u in $item.Users) {
        $user = Resolve-User -Upn $u
        if ($user) { Assign-Principal -PrincipalId $user.Id -PrincipalName $u -Scope $appGroup.Id -RoleName $roleName }
    }
}

Write-Result "OK" "Completed."