<#
.SYNOPSIS
   Creates bulk students, assigns ALL available licenses, ensures the Temporary Access Pass (TAP)
   policy is enabled for all users, issues a TAP, and exports a comprehensive report to a CSV file
   in the script's directory.

.NOTES
   Requires Microsoft Graph permissions/scopes including:
   - User.ReadWrite.All
   - Directory.ReadWrite.All
   - RoleManagement.ReadWrite.Directory
   - UserAuthenticationMethod.ReadWrite.All
   - Policy.ReadWrite.AuthenticationMethod

   The signed-in admin also needs sufficient Entra role rights to update authentication method policy,
   create TAPs for users, and manage directory role membership.
#>

param(
    [Nullable[int]]$Count,
    [string]$Region
)

# --- Configuration Variables ---
$DefaultStudentCount = 24
$StudentCount = if ($PSBoundParameters.ContainsKey('Count') -and $Count -ne $null) {
    Write-Host "Using count from bootstrap: $Count" -ForegroundColor Cyan
    $Count
} else {
    $DefaultStudentCount
}

$DefaultUsageLocation = "GB"   # REQUIRED FOR LICENSING (e.g., "US", "GB", "CA")
$UsageLocation = if ($PSBoundParameters.ContainsKey('Region') -and $Region) {
    Write-Host "Using region from bootstrap: $Region" -ForegroundColor Cyan
    $Region
} else {
    $DefaultUsageLocation
}

$BaseName = "student"
$DefaultPassword = "LTRE1360!"
$TapLifetimeMinutes = 43200 # 30 days

# Instructor admin account settings
$InstructorEnabled      = $true
$InstructorUserName     = "instructorAdmin"
$InstructorDisplayName  = "Instructor Admin"
$InstructorMailNickname = "instAdm"
$InstructorPassword     = $DefaultPassword
$InstructorRoleName     = "Global Administrator"  # Entra role to assign to the instructor
$AssignInstructorRole   = $true               # Set to $true to assign the instructor role

# TAP policy targeting mode for this script version:
# This script enables TAP for ALL users in the tenant.
$TapEnableAllUsers = $true

# --- Initialize Report Array ---
$ReportData = [System.Collections.Generic.List[PSCustomObject]]::new()

# --- Helper: Ensure TAP policy is enabled and compatible with requested lifetime ---
function EnableTapPolicy {
    param(
        [Parameter(Mandatory)]
        [int]$RequestedLifetimeMinutes
    )

    $tapPolicyUri = "https://graph.microsoft.com/v1.0/policies/authenticationMethodsPolicy/authenticationMethodConfigurations/temporaryAccessPass"

    Write-Host "`nChecking TAP policy..." -ForegroundColor Cyan

    try {
        $policy = Invoke-MgGraphRequest -Method GET -Uri $tapPolicyUri
    }
    catch {
        throw "Failed to read TAP policy. Ensure your account has Policy.ReadWrite.AuthenticationMethod consent and the required Entra role. Raw error: $($_.Exception.Message)"
    }

    $patchBody = [ordered]@{
        "@odata.type" = "#microsoft.graph.temporaryAccessPassAuthenticationMethodConfiguration"
    }

    $needsPatch = $false

    # 1) Ensure TAP method is enabled
    if ($policy.state -ne "enabled") {
        $patchBody.state = "enabled"
        $needsPatch = $true
        Write-Host " - TAP policy is disabled; will enable it." -ForegroundColor Yellow
    }
    else {
        Write-Host " - TAP policy already enabled." -ForegroundColor DarkGray
    }

    # 2) Ensure includeTargets contains all_users (this overwrites targeting if patched)
    $hasAllUsersTarget = $false
    if ($null -ne $policy.includeTargets) {
        foreach ($target in $policy.includeTargets) {
            if ($target.id -eq "all_users" -and $target.targetType -eq "group") {
                $hasAllUsersTarget = $true
                break
            }
        }
    }

    if ($TapEnableAllUsers -and -not $hasAllUsersTarget) {
        $patchBody.includeTargets = @(
            @{
                "@odata.type" = "#microsoft.graph.authenticationMethodTarget"
                targetType = "group"
                id = "all_users"
                isRegistrationRequired = $false
            }
        )
        $needsPatch = $true
        Write-Host " - TAP policy is not targeted to all users; will set includeTargets = all_users." -ForegroundColor Yellow
    }
    elseif ($TapEnableAllUsers) {
        Write-Host " - TAP policy already targets all users." -ForegroundColor DarkGray
    }

    # 3) Ensure policy allows reusable TAPs if the script creates reusable TAPs
    if ($policy.isUsableOnce -eq $true) {
        $patchBody.isUsableOnce = $false
        $needsPatch = $true
        Write-Host " - TAP policy is currently one-time-use only; will allow reusable TAPs." -ForegroundColor Yellow
    }
    else {
        Write-Host " - TAP policy already allows reusable TAPs." -ForegroundColor DarkGray
    }

    # 4) Ensure requested TAP lifetime is the policy min/max/default
    $currentMin = if ($null -ne $policy.minimumLifetimeInMinutes) { [int]$policy.minimumLifetimeInMinutes } else { 10 }
    $currentMax = if ($null -ne $policy.maximumLifetimeInMinutes) { [int]$policy.maximumLifetimeInMinutes } else { 480 }
    $currentDefault = if ($null -ne $policy.defaultLifetimeInMinutes) { [int]$policy.defaultLifetimeInMinutes } else { 60 }

    if ($RequestedLifetimeMinutes -ne $currentMin) {
        $patchBody.minimumLifetimeInMinutes = $RequestedLifetimeMinutes
        $needsPatch = $true
        Write-Host " - Setting TAP policy minimum to $RequestedLifetimeMinutes." -ForegroundColor Yellow
    }

    if ($RequestedLifetimeMinutes -ne $currentMax) {
        $patchBody.maximumLifetimeInMinutes = $RequestedLifetimeMinutes
        $needsPatch = $true
        Write-Host " - Setting TAP policy maximum to $RequestedLifetimeMinutes." -ForegroundColor Yellow
    }

    if ($RequestedLifetimeMinutes -ne $currentDefault) {
        $patchBody.defaultLifetimeInMinutes = $RequestedLifetimeMinutes
        $needsPatch = $true
        Write-Host " - Setting TAP policy default to $RequestedLifetimeMinutes." -ForegroundColor Yellow
    }

    if ($needsPatch) {
        $jsonBody = $patchBody | ConvertTo-Json -Depth 10

        try {
            Invoke-MgGraphRequest `
                -Method PATCH `
                -Uri $tapPolicyUri `
                -Headers @{ "Content-Type" = "application/json" } `
                -Body $jsonBody | Out-Null

            Write-Host " - TAP policy updated successfully." -ForegroundColor Green
        }
        catch {
            throw "Failed to update TAP policy. Raw error: $($_.Exception.Message)"
        }

        # Re-read policy to confirm
        try {
            $confirmedPolicy = Invoke-MgGraphRequest -Method GET -Uri $tapPolicyUri
            Write-Host " - Confirmed TAP policy state: $($confirmedPolicy.state)" -ForegroundColor Green
            Write-Host "WAITING: Allowing new TAP policy time to sync before creating users" -ForegroundColor Yellow
            Start-Sleep -Seconds 20
        }
        catch {
            Write-Host " - WARNING: TAP policy update appears to have succeeded, but confirmation GET failed: $($_.Exception.Message)" -ForegroundColor Yellow
        }
    }
    else {
        Write-Host " - No TAP policy changes required." -ForegroundColor Green
    }
}

function Ensure-DirectoryRoleAssignment {
    param(
        [Parameter(Mandatory)][string]$UserId,
        [Parameter(Mandatory)][string]$RoleDisplayName
    )

    $roleTemplate = Get-MgDirectoryRoleTemplate -All | Where-Object { $_.DisplayName -eq $RoleDisplayName } | Select-Object -First 1
    if (-not $roleTemplate) {
        throw "Role template '$RoleDisplayName' not found. Confirm the role name is valid."
    }

    $directoryRole = Get-MgDirectoryRoleByRoleTemplateId -RoleTemplateId $roleTemplate.Id -ErrorAction SilentlyContinue
    if (-not $directoryRole) {
        Write-Host " -> Activating directory role '$RoleDisplayName'..." -ForegroundColor Yellow
        $directoryRole = New-MgDirectoryRole -BodyParameter @{ roleTemplateId = $roleTemplate.Id } -ErrorAction Stop
        Start-Sleep -Seconds 3
    }

    if (-not $directoryRole -or -not $directoryRole.Id) {
        throw "Unable to resolve or activate directory role '$RoleDisplayName'."
    }

    $existingMember = Get-MgDirectoryRoleMember -DirectoryRoleId $directoryRole.Id -All | Where-Object { $_.Id -eq $UserId }
    if ($existingMember) {
        return "Already member of $RoleDisplayName"
    }

    $odataId = "https://graph.microsoft.com/v1.0/directoryObjects/$UserId"
    New-MgDirectoryRoleMemberByRef -DirectoryRoleId $directoryRole.Id -OdataId $odataId -ErrorAction Stop | Out-Null

    return "Assigned $RoleDisplayName"
}

function Process-User {
    param(
        [Parameter(Mandatory)][string]$UserPrincipalName,
        [Parameter(Mandatory)][string]$DisplayName,
        [string]$MailNickname = $null,
        [string]$Password = $DefaultPassword,
        [switch]$AssignRole,
        [string]$RoleDisplayName = $null
    )

    $AssignedToUser = @()
    $LicenseStatus = "Not evaluated"
    $Warnings = @()
    $TapPasscode = "N/A"
    $Status = "Success"
    $ErrorMessage = ""

    if (-not $MailNickname) {
        $MailNickname = $UserPrincipalName.Split("@")[0]
    }

    Write-Host "`nProcessing $UserPrincipalName..." -ForegroundColor Cyan

    try {
        $ExistingUser = Get-MgUser -Filter "UserPrincipalName eq '$UserPrincipalName'" -ErrorAction SilentlyContinue

        if ($ExistingUser) {
            Write-Host " -> User already exists. Skipping creation." -ForegroundColor DarkGray
            $TargetUserId = $ExistingUser.Id
        }
        else {
            $PasswordProfile = @{
                Password = $Password
                ForceChangePasswordNextSignIn = $false
            }

            $NewUser = New-MgUser -DisplayName $DisplayName `
                                  -UserPrincipalName $UserPrincipalName `
                                  -MailNickname $MailNickname `
                                  -AccountEnabled `
                                  -PasswordProfile $PasswordProfile `
                                  -UsageLocation $UsageLocation

            Write-Host " -> User created successfully." -ForegroundColor Green
            $TargetUserId = $NewUser.Id
        }

        if ($AvailableSkus.Count -gt 0) {
            $LicensesToAdd = @()
            $LicensePoolsToAdd = @()
            $SkippedLicensePools = @()
            $FailedLicensePools = @()
            $ExistingLicenses = Get-MgUserLicenseDetail -UserId $TargetUserId

            foreach ($Sku in $AvailableSkus) {
                if ($ExistingLicenses -and $ExistingLicenses.SkuId -contains $Sku.SkuId) {
                    Write-Host " -> License already assigned: $($Sku.SkuPartNumber)" -ForegroundColor DarkGray
                    $AssignedToUser += $Sku.SkuPartNumber
                    continue
                }

                if ($Sku.RemainingUnits -le 0) {
                    Write-Host " -> No available licenses remaining for: $($Sku.SkuPartNumber)" -ForegroundColor Yellow
                    $SkippedLicensePools += $Sku.SkuPartNumber
                    continue
                }

                $LicensesToAdd += @{ SkuId = $Sku.SkuId }
                $LicensePoolsToAdd += $Sku
                Write-Host " -> Queuing license for assignment: $($Sku.SkuPartNumber) ($($Sku.RemainingUnits) remaining before assignment)" -ForegroundColor Cyan
            }

            if ($LicensesToAdd.Count -gt 0) {
                try {
                    Set-MgUserLicense -UserId $TargetUserId -AddLicenses $LicensesToAdd -RemoveLicenses @() -ErrorAction Stop | Out-Null

                    foreach ($Sku in $LicensePoolsToAdd) {
                        $Sku.RemainingUnits = [Math]::Max(0, ([int]$Sku.RemainingUnits - 1))
                        $AssignedToUser += $Sku.SkuPartNumber
                    }

                    Write-Host " -> Successfully assigned $($LicensesToAdd.Count) new license(s)." -ForegroundColor Green
                }
                catch {
                    $failedLicenseNames = @($LicensePoolsToAdd | Select-Object -ExpandProperty SkuPartNumber)
                    $licenseError = $_.Exception.Message
                    $FailedLicensePools += $failedLicenseNames
                    $Warnings += "License assignment failed for $($failedLicenseNames -join ', '): $licenseError"

                    if ($licenseError -match '(?i)(no.*available|available licenses|license.*capacity|subscription.*capacity)') {
                        $SkippedLicensePools += $failedLicenseNames
                        $FailedLicensePools = @($FailedLicensePools | Where-Object { $failedLicenseNames -notcontains $_ })
                        foreach ($Sku in $LicensePoolsToAdd) {
                            $Sku.RemainingUnits = 0
                        }
                    }

                    Write-Host " -> WARNING: Could not assign license(s): $($failedLicenseNames -join ', ')" -ForegroundColor Yellow
                    Write-Host "    $licenseError" -ForegroundColor Yellow
                }
            }

            $SkippedLicensePools = @($SkippedLicensePools | Select-Object -Unique)
            $FailedLicensePools = @($FailedLicensePools | Select-Object -Unique)
            $LicenseProblems = @()
            if ($SkippedLicensePools.Count -gt 0) {
                $LicenseProblems += "no licenses available: $($SkippedLicensePools -join ', ')"
            }
            if ($FailedLicensePools.Count -gt 0) {
                $LicenseProblems += "assignment failed: $($FailedLicensePools -join ', ')"
            }

            if ($AssignedToUser.Count -eq 0 -and $LicenseProblems.Count -gt 0) {
                $LicenseStatus = "Unlicensed"
                $Warnings += "No license applied: $($LicenseProblems -join '; ')"
                Write-Host " -> User will remain unlicensed; $($LicenseProblems -join '; ')." -ForegroundColor Yellow
            }
            elseif ($AssignedToUser.Count -gt 0 -and $LicenseProblems.Count -gt 0) {
                $LicenseStatus = "Partially licensed; $($LicenseProblems -join '; ')"
                $Warnings += $LicenseStatus
                Write-Host " -> User partially licensed; $($LicenseProblems -join '; ')." -ForegroundColor Yellow
            }
            elseif ($AssignedToUser.Count -gt 0) {
                $LicenseStatus = "Licensed"
            }
            else {
                $LicenseStatus = "Unlicensed"
                $Warnings += "No license applied: no licenses were assigned"
                Write-Host " -> User will remain unlicensed; no licenses were assigned." -ForegroundColor Yellow
            }
        }
        else {
            $LicenseStatus = "Unlicensed"
            $Warnings += "No license applied: no available licenses were found in the tenant"
            Write-Host " -> User will remain unlicensed; no available licenses were found in the tenant." -ForegroundColor Yellow
        }

        $AuthReady = $false

        for ($j = 1; $j -le 10; $j++) {
            try {
                Get-MgUserAuthenticationMethod `
                    -UserId $TargetUserId `
                    -ErrorAction Stop | Out-Null

                $AuthReady = $true
                break
            }
            catch {
                Write-Host " -> Waiting for authentication methods service..." -ForegroundColor Yellow
                Start-Sleep 5
            }
        }

        if (-not $AuthReady) {
            throw "Authentication methods endpoint never became available."
        }

        $Body = @{
            lifetimeInMinutes = $TapLifetimeMinutes
            isUsableOnce      = $false
        }

        try {
            if ($AssignRole -and $RoleDisplayName) {
                try {
                    $RoleStatus = Ensure-DirectoryRoleAssignment -UserId $TargetUserId -RoleDisplayName $RoleDisplayName
                    Write-Host " -> $RoleStatus" -ForegroundColor Green
                }
                catch {
                    Write-Warning "Failed to assign role '$RoleDisplayName': $($_.Exception.Message)"
                }
            }

            $Tap = New-MgUserAuthenticationTemporaryAccessPassMethod `
                -UserId $TargetUserId `
                -BodyParameter $Body `
                -ErrorAction Stop

            $TapPasscode = $Tap.TemporaryAccessPass
            Write-Host " -> TAP Created! Passcode: $TapPasscode" -ForegroundColor Green
        }
        catch {
            $Warnings += "Failed to create TAP: $($_.Exception.Message)"
            Write-Warning "Failed to create TAP: $($_.Exception.Message)"
        }
    }
    catch {
        $Status = "Failed"
        $ErrorMessage = $_.Exception.Message
        Write-Host " -> ERROR processing $UserPrincipalName : $ErrorMessage" -ForegroundColor Red
    }

    $WarningMessage = ($Warnings | Select-Object -Unique) -join " | "
    $ReportErrorMessage = if ($ErrorMessage) { $ErrorMessage } else { $WarningMessage }

    return [PSCustomObject]@{
        "UserPrincipalName"   = $UserPrincipalName
        "DisplayName"         = $DisplayName
        "TemporaryAccessPass" = $TapPasscode
        "LicensesAssigned"    = if ($AssignedToUser.Count -gt 0) { ($AssignedToUser | Select-Object -Unique) -join ", " } else { "None" }
        "LicenseStatus"       = $LicenseStatus
        "RoleAssigned"        = if ($AssignRole -and $RoleDisplayName) { $RoleDisplayName } else { "None" }
        "Status"              = if ($Status -eq "Success" -and $Warnings.Count -gt 0) { "SuccessWithWarnings" } else { $Status }
        "WarningMessage"      = $WarningMessage
        "ErrorMessage"        = $ReportErrorMessage
    }
}

# --- MAIN SCRIPT EXECUTION ---
# --- 1. Connect to Microsoft Graph ---
$RequiredGraphModules = @(
    'Microsoft.Graph.Authentication',
    'Microsoft.Graph.Identity.DirectoryManagement',
    'Microsoft.Graph.Users',
    'Microsoft.Graph.Users.Actions',
    'Microsoft.Graph.Identity.SignIns'
)

if (-not (Get-Command Connect-MgGraph -ErrorAction SilentlyContinue)) {
    try {
        foreach ($moduleName in $RequiredGraphModules) {
            if (-not (Get-Module -ListAvailable -Name $moduleName)) {
                throw "Required Graph module not installed: $moduleName"
            }
            Import-Module -Name $moduleName -ErrorAction Stop
        }
    }
    catch {
        Write-Host "ERROR: Required Microsoft Graph PowerShell modules are not available." -ForegroundColor Red
        Write-Host "Install them with the following commands:" -ForegroundColor Yellow
        foreach ($moduleName in $RequiredGraphModules) {
            Write-Host "    Install-Module $moduleName -Scope CurrentUser" -ForegroundColor Yellow
        }
        throw
    }
}

Write-Host "Connecting to Microsoft Graph..." -ForegroundColor Cyan
Disconnect-MgGraph -ErrorAction SilentlyContinue
Connect-MgGraph -Scopes @(
    "User.ReadWrite.All",
    "Directory.ReadWrite.All",
    "RoleManagement.ReadWrite.Directory",
    "UserAuthenticationMethod.ReadWrite.All",
    "Policy.ReadWrite.AuthenticationMethod"
)

$graphContext = Get-MgContext
Write-Host "Connected as: $($graphContext.Account)" -ForegroundColor Green
Write-Host "Granted Graph scopes: $($graphContext.Scopes -join ', ')" -ForegroundColor Green
if (-not ($graphContext.Scopes -contains 'RoleManagement.ReadWrite.Directory')) {
    throw "RoleManagement.ReadWrite.Directory was not granted during sign-in. Re-run the script and consent when prompted."
}

# --- 2. Ensure tenant TAP policy is enabled before creating user TAPs ---
try {
    EnableTapPolicy -RequestedLifetimeMinutes $TapLifetimeMinutes
}
catch {
    Write-Host "FATAL: Could not ensure TAP policy is enabled. $($_.Exception.Message)" -ForegroundColor Red
    break
}

# --- 3. Dynamically Gather Tenant Info ---
$Domain = (Get-MgDomain | Where-Object IsDefault).Id
Write-Host "`nTargeting Domain: $Domain" -ForegroundColor Green

# Find ALL available license SKUs with free units and track remaining capacity locally
$AvailableSkus = @(Get-MgSubscribedSku | Where-Object {
    $_.CapabilityStatus -eq 'Enabled' -and ($_.PrepaidUnits.Enabled - $_.ConsumedUnits) -gt 0
} | ForEach-Object {
    [PSCustomObject]@{
        SkuId = $_.SkuId
        SkuPartNumber = $_.SkuPartNumber
        EnabledUnits = [int]$_.PrepaidUnits.Enabled
        ConsumedUnits = [int]$_.ConsumedUnits
        RemainingUnits = [int]($_.PrepaidUnits.Enabled - $_.ConsumedUnits)
    }
})

if (-not $AvailableSkus) {
    Write-Host "WARNING: No available licenses found in this tenant. Users will be created without licenses." -ForegroundColor Yellow
}
else {
    Write-Host "Found $($AvailableSkus.Count) available license pool(s):" -ForegroundColor Green
    foreach ($Sku in $AvailableSkus) {
        Write-Host " - $($Sku.SkuPartNumber) (Available: $($Sku.RemainingUnits))" -ForegroundColor Green
    }
}

# --- 4. User Creation and Configuration Loop ---
if ($InstructorEnabled) {
    $InstructorPrincipalName = "$InstructorUserName@$Domain"
    $ReportData.Add((Process-User -UserPrincipalName $InstructorPrincipalName `
                                -DisplayName $InstructorDisplayName `
                                -MailNickname $InstructorMailNickname `
                                -Password $InstructorPassword `
                                -AssignRole:$AssignInstructorRole `
                                -RoleDisplayName $InstructorRoleName))
}

for ($i = 1; $i -le $StudentCount; $i++) {
    $UserPrincipalName = "$BaseName$i@$Domain"
    $DisplayName = "Student $i"

    $ReportData.Add((Process-User -UserPrincipalName $UserPrincipalName `
                                -DisplayName $DisplayName `
                                -MailNickname "$BaseName$i" `
                                -Password $DefaultPassword))
}

$UnlicensedAccounts = @($ReportData | Where-Object { $_.LicenseStatus -like 'Unlicensed*' })
$PartiallyLicensedAccounts = @($ReportData | Where-Object { $_.LicenseStatus -like 'Partially licensed*' })

Write-Host "`nLicense assignment summary:" -ForegroundColor Cyan
Write-Host " - Licensed accounts: $(($ReportData.Count - $UnlicensedAccounts.Count - $PartiallyLicensedAccounts.Count))" -ForegroundColor Green
Write-Host " - Partially licensed accounts: $($PartiallyLicensedAccounts.Count)" -ForegroundColor Yellow
Write-Host " - Unlicensed accounts: $($UnlicensedAccounts.Count)" -ForegroundColor Yellow

# --- 5. Export Data to CSV alongside Script ---
if ($ReportData.Count -gt 0) {
    # Establish saving directory: use $PSScriptRoot if run as a file, otherwise fall back to current directory
    $TargetDirectory = if ($PSScriptRoot) { $PSScriptRoot } else { $PWD }

    # Generate a unique timestamped filename so multiple tenant runs don't overwrite each other
    $Timestamp = Get-Date -Format "yyyyMMdd_HHmmss"
    $CsvPath = Join-Path $TargetDirectory "Tenant_Students_Export_$Timestamp.csv"

    Write-Host "`nExporting results to CSV..." -ForegroundColor Cyan
    $ReportData | Export-Csv -Path $CsvPath -NoTypeInformation -Encoding utf8
    Write-Host "Report successfully saved to: $CsvPath" -ForegroundColor Green
}

Write-Host "`nScript execution complete." -ForegroundColor Cyan
