#requires -Version 7.0
<#
.SYNOPSIS
    Deploys a CIS M365 tenant-level baseline from a YAML manifest.

.DESCRIPTION
    Reads a baseline YAML file that mirrors the OpenIntuneBaseline schema and
    adds tenantConfig sections for Entra ID, Conditional Access, Defender,
    Exchange, SharePoint, and Teams.

    CONDITIONAL ACCESS SAFETY:
    - All CA policies are created in report-only mode by default (global switch).
    - The break-glass group is automatically excluded from every CA policy.
    - You must explicitly pass -Apply when Mode is Deploy.

.EXAMPLE
    # Assess without making any changes
    ./Scripts/Deploy-CISM365Baseline.ps1 -BaselinePath ./Baselines/mytenant-cisv7.yaml

.EXAMPLE
    # Deploy after review
    ./Scripts/Deploy-CISM365Baseline.ps1 -BaselinePath ./Baselines/mytenant-cisv7.yaml -Mode Deploy -Apply -Verbose

.EXAMPLE
    # Deploy only Conditional Access and Entra ID settings
    ./Scripts/Deploy-CISM365Baseline.ps1 -BaselinePath ./Baselines/mytenant-cisv7.yaml -Mode Deploy -Apply -Workloads EntraID,ConditionalAccess
#>
[CmdletBinding(SupportsShouldProcess)]
param(
    [Parameter(Mandatory = $true)]
    [string]$BaselinePath,

    [Parameter()]
    [string]$TenantId,

    [Parameter()]
    [ValidateSet('Assess','Deploy')]
    [string]$Mode = 'Assess',

    [Parameter()]
    [ValidateSet('EntraID','ConditionalAccess','Defender','Exchange','SharePoint','Teams','AdminCenter','Purview','PowerBI')]
    [string[]]$Workloads = @('EntraID','ConditionalAccess','Defender','Exchange','SharePoint','Teams','AdminCenter','Purview','PowerBI'),

    [Parameter()]
    [switch]$Apply,

    [Parameter()]
    [switch]$WhatIf,

    [Parameter()]
    [ValidateSet('AppOnly','Browser','DeviceCode')]
    [string]$AuthMode = 'Browser',

    [Parameter()]
    [string]$AppId,

    [Parameter()]
    [string]$Secret,

    [Parameter()]
    [string]$Certificate
)

$ErrorActionPreference = 'Stop'

#region Helper Functions
function Test-YamlModule {
    return [bool](Get-Module -ListAvailable -Name powershell-yaml)
}

function Install-YamlModule {
    Write-Host "powershell-yaml module is required but not installed." -ForegroundColor Yellow
    if (-not $WhatIf) {
        $confirm = Read-Host "Install powershell-yaml from PSGallery now? [Y/n]"
        if ($confirm -match "^\s*n") {
            throw "powershell-yaml is required. Install it with: Install-Module powershell-yaml -Scope CurrentUser -Force"
        }
        Install-Module powershell-yaml -Scope CurrentUser -Force
        Import-Module powershell-yaml -Force
    }
}

function Resolve-RelativePath {
    param([string]$Path, [string]$BasePath)
    if ([System.IO.Path]::IsPathRooted($Path)) { return $Path }
    $baseDir = Split-Path -Parent $BasePath
    return Join-Path $baseDir $Path
}

function Write-SectionHeader {
    param([string]$Title)
    Write-Host "`n========================================" -ForegroundColor Cyan
    Write-Host "  $Title" -ForegroundColor Cyan
    Write-Host "========================================" -ForegroundColor Cyan
}

function Add-Result {
    param(
        [string]$Workload,
        [string]$Control,
        [string]$Status,      # Pass, Fail, Fixed, Skipped, Error, Manual
        [string]$Message,
        [string]$Remediation = ''
    )
    $script:Results.Add([PSCustomObject]@{
        Workload    = $Workload
        Control     = $Control
        Status      = $Status
        Message     = $Message
        Remediation = $Remediation
    })
    switch ($Status) {
        'Fixed'   { $script:ChangesMade++ }
        'Skipped' { $script:ChangesSkipped++ }
        'Error'   { $script:Errors++ }
    }
}

function Invoke-ApplyMutation {
    param([string]$Name, [hashtable]$Mutation)
    if (-not $Mutation) { return $Name }

    $search = $Mutation["search"]
    $replace = $Mutation["replace"]
    $prefix = $Mutation["prefix"]

    if ($search -and $replace) {
        $Name = $Name -replace $search, $replace
    }
    elseif ($prefix) {
        if (-not $Name.StartsWith($prefix)) {
            $Name = "$prefix$Name"
        }
    }
    return $Name
}

# Common Entra admin role template IDs
$script:RoleTemplateMap = @{
    "Global Administrator"            = "62e90394-69f5-4237-9190-012177145e10"
    "Privileged Role Administrator"   = "e8611ab8-c189-46e8-94e1-60213ab1f814"
    "Security Administrator"          = "194ae4cb-b126-40b2-bd5b-6091b380977d"
    "Exchange Administrator"          = "29232cdf-9323-42fd-ade2-1d097af3e4de"
    "SharePoint Administrator"        = "f28a1f50-f6e7-4571-818b-6a12f2af6b6c"
    "Conditional Access Administrator"= "b1be1c3e-b65d-4f19-8427-f6fa0d97feb9"
    "Application Administrator"       = "9b895d92-2cd3-44c7-9d02-a6ac2d5ea5d3"
    "Cloud Application Administrator" = "158c047a-c907-4556-b7ef-446551a6b5f7"
    "User Administrator"              = "fe930be7-5e62-47db-91af-98c3a49a38b1"
    "Helpdesk Administrator"          = "729827e3-9c14-49f7-bb1b-9608f156bbb8"
    "Billing Administrator"           = "b0f54661-2d74-4c50-afa3-1ec803f12efe"
    "Authentication Administrator"    = "c4e39bd9-1100-46d3-8c65-fb160da0071f"
    "Password Administrator"          = "966707d0-3269-4727-9be2-8c3a10f19b9d"
    "Global Reader"                   = "f2ef992c-3afb-46b9-b7cf-a127eeeb959e"
}

$script:Results = [System.Collections.Generic.List[object]]::new()
$script:ChangesMade = 0
$script:ChangesSkipped = 0
$script:Errors = 0
$script:GroupCache = @{}
$script:NamedLocationCache = @{}
$script:EffectiveWhatIf = $WhatIf.IsPresent -or ($Mode -eq 'Deploy' -and -not $Apply.IsPresent)
#endregion

#region Auth
Write-SectionHeader "Authentication"

# Microsoft Graph
$GraphScopes = @(
    'Directory.Read.All','Directory.ReadWrite.All','Policy.Read.All','Policy.ReadWrite.ConditionalAccess',
    'Organization.Read.All','Organization.ReadWrite.All','RoleManagement.ReadWrite.Directory',
    'IdentityRiskyUser.Read.All','IdentityRiskEvent.Read.All','Group.ReadWrite.All'
)
Write-Host "Connecting to Microsoft Graph (mode: $AuthMode)..." -NoNewline

$connectParams = @{}
if ($TenantId) { $connectParams.TenantId = $TenantId }

switch ($AuthMode) {
    'AppOnly' {
        if (-not $AppId) { throw "AppId is required for AppOnly auth mode." }
        if ($Secret) {
            $secureSecret = ConvertTo-SecureString -String $Secret -AsPlainText -Force
            $credential = New-Object System.Management.Automation.PSCredential($AppId, $secureSecret)
            $connectParams.ClientSecretCredential = $credential
        }
        elseif ($Certificate) {
            $cert = Get-ChildItem Cert:\CurrentUser\My | Where-Object { $_.Thumbprint -eq $Certificate -or $_.Subject -eq $Certificate } | Select-Object -First 1
            if (-not $cert) { throw "Certificate not found: $Certificate" }
            $connectParams.ClientCertificateCredential = $cert
        }
        else {
            throw "Secret or Certificate is required for AppOnly auth mode."
        }
        Connect-MgGraph @connectParams -NoWelcome
    }
    'DeviceCode' {
        Connect-MgGraph -Scopes ($GraphScopes -join ',') @connectParams -UseDeviceCode -NoWelcome
    }
    default { # Browser / Interactive
        Connect-MgGraph -Scopes ($GraphScopes -join ',') @connectParams -NoWelcome
    }
}

$context = Get-MgContext
Write-Host " OK ($($context.Account))" -ForegroundColor Green

# Exchange Online
if ($Workloads -contains 'Defender' -or $Workloads -contains 'Exchange') {
    Write-Host "Connecting to Exchange Online..." -NoNewline
    Connect-ExchangeOnline -ShowBanner:$false
    Write-Host " OK" -ForegroundColor Green
}

# SharePoint
if ($Workloads -contains 'SharePoint') {
    Write-Host "Connecting to SharePoint Online..." -NoNewline
    # URL will be resolved from YAML later
    Write-Host " deferred" -ForegroundColor Yellow
}

# Teams
if ($Workloads -contains 'Teams') {
    Write-Host "Connecting to Microsoft Teams..." -NoNewline
    Connect-MicrosoftTeams
    Write-Host " OK" -ForegroundColor Green
}
#endregion

#region Load YAML
if (-not (Test-YamlModule)) { Install-YamlModule }
Import-Module powershell-yaml -Force

$baselinePathResolved = Resolve-Path $BaselinePath | Select-Object -ExpandProperty Path
if (-not (Test-Path $baselinePathResolved)) {
    throw "Baseline file not found: $BaselinePath"
}
$baselineDir = Split-Path $baselinePathResolved -Parent

Write-Host "`nLoading baseline: $baselinePathResolved" -ForegroundColor Cyan
$yamlText = Get-Content $baselinePathResolved -Raw
$yamlRoot = ConvertFrom-Yaml -Yaml $yamlText

if (-not $yamlRoot -or -not $yamlRoot.ContainsKey("baseline")) {
    throw "Invalid baseline YAML: missing 'baseline' root node."
}
$baseline = $yamlRoot["baseline"]

$globalMutation = $null
if ($baseline.ContainsKey("tenantMutation")) {
    $globalMutation = $baseline["tenantMutation"]
}

$tenantConfig = $baseline.ContainsKey("tenantConfig") ? $baseline["tenantConfig"] : @{}

Write-Host "Baseline name : $($baseline["name"])" -ForegroundColor Cyan
Write-Host "Mode          : $Mode" -ForegroundColor Cyan
Write-Host "Workloads     : $($Workloads -join ', ')" -ForegroundColor Cyan
if ($script:EffectiveWhatIf) { Write-Host "*** DRY-RUN / WHATIF MODE ***" -ForegroundColor Magenta }
#endregion

#region Resolve / create groups
function Get-OrCreateGroup {
    param([string]$DisplayName, [string]$MailNickname, [bool]$SecurityEnabled = $true)
    if ($script:GroupCache.ContainsKey($DisplayName)) {
        return $script:GroupCache[$DisplayName]
    }

    $existing = Get-MgGroup -Filter "displayName eq '$($DisplayName -replace "'","''")'" -ErrorAction SilentlyContinue | Select-Object -First 1
    if ($existing) {
        Write-Host "  Group exists: $DisplayName ($($existing.Id))" -ForegroundColor Green
        $script:GroupCache[$DisplayName] = $existing.Id
        return $existing.Id
    }

    if ($script:EffectiveWhatIf) {
        Write-Host "  [WHATIF] Would create group: $DisplayName" -ForegroundColor Magenta
        $script:GroupCache[$DisplayName] = "WHATIF-$DisplayName"
        return $script:GroupCache[$DisplayName]
    }

    Write-Host "  Creating group: $DisplayName" -ForegroundColor Yellow
    $newGrp = New-MgGroup -DisplayName $DisplayName -MailEnabled:$false -MailNickname $MailNickname -SecurityEnabled:$SecurityEnabled
    $script:GroupCache[$DisplayName] = $newGrp.Id
    Write-Host "    Created: $($newGrp.Id)" -ForegroundColor Green
    return $newGrp.Id
}

function Get-OrCreateNamedLocation {
    param(
        [string]$DisplayName,
        [string]$Type,
        [array]$CountriesAndRegions,
        [bool]$IncludeUnknownCountriesAndRegions = $false,
        [bool]$IsTrusted = $false,
        [array]$IpRanges
    )
    if ($script:NamedLocationCache.ContainsKey($DisplayName)) {
        return $script:NamedLocationCache[$DisplayName]
    }

    $existing = Get-MgIdentityConditionalAccessNamedLocation -Filter "displayName eq '$($DisplayName -replace "'","''")'" -ErrorAction SilentlyContinue | Select-Object -First 1
    if ($existing) {
        Write-Host "  Named location exists: $DisplayName ($($existing.Id))" -ForegroundColor Green
        $script:NamedLocationCache[$DisplayName] = $existing.Id
        return $existing.Id
    }

    if ($script:EffectiveWhatIf) {
        Write-Host "  [WHATIF] Would create named location: $DisplayName" -ForegroundColor Magenta
        $script:NamedLocationCache[$DisplayName] = "WHATIF-$DisplayName"
        return $script:NamedLocationCache[$DisplayName]
    }

    Write-Host "  Creating named location: $DisplayName" -ForegroundColor Yellow
    $body = @{ displayName = $DisplayName }
    if ($Type -eq 'country') {
        $body['@odata.type'] = '#microsoft.graph.countryNamedLocation'
        $body['countriesAndRegions'] = $CountriesAndRegions
        $body['includeUnknownCountriesAndRegions'] = $IncludeUnknownCountriesAndRegions
    }
    elseif ($Type -eq 'ip') {
        $body['@odata.type'] = '#microsoft.graph.ipNamedLocation'
        $body['isTrusted'] = $IsTrusted
        $body['ipRanges'] = @()
        foreach ($range in $IpRanges) {
            $body['ipRanges'] += @{
                '@odata.type' = '#microsoft.graph.iPv4CidrRange'
                cidrAddress = $range
            }
        }
    }

    $newLoc = New-MgIdentityConditionalAccessNamedLocation -BodyParameter $body
    $script:NamedLocationCache[$DisplayName] = $newLoc.Id
    Write-Host "    Created: $($newLoc.Id)" -ForegroundColor Green
    return $newLoc.Id
}

if ($baseline.ContainsKey("groups") -and $baseline["groups"]) {
    Write-SectionHeader "Resolving Groups"
    foreach ($grpDef in $baseline["groups"]) {
        $displayName = $grpDef["displayName"]
        $mailNick = $grpDef["mailNickname"]
        $secEnabled = if ($grpDef.ContainsKey("securityEnabled")) { [bool]$grpDef["securityEnabled"] } else { $true }
        $null = Get-OrCreateGroup -DisplayName $displayName -MailNickname $mailNick -SecurityEnabled $secEnabled
    }
}
#endregion

#region CA Policy Builder
function ConvertTo-CAPolicyPayload {
    param([hashtable]$PolicyDef, [string]$BreakGlassGroupId, [hashtable]$Mutation)

    $name = Invoke-ApplyMutation -Name $PolicyDef["name"] -Mutation $Mutation
    $state = $PolicyDef["state"]
    $description = $PolicyDef["description"]

    # Build conditions
    $conditions = @{}

    # Applications
    if ($PolicyDef["conditions"].ContainsKey("applications")) {
        $appConditions = @{}
        $appDef = $PolicyDef["conditions"]["applications"]
        if ($appDef.ContainsKey("includeApplications")) { $appConditions["includeApplications"] = $appDef["includeApplications"] }
        if ($appDef.ContainsKey("excludeApplications")) { $appConditions["excludeApplications"] = $appDef["excludeApplications"] }
        if ($appDef.ContainsKey("includeUserActions")) { $appConditions["includeUserActions"] = $appDef["includeUserActions"] }
        $conditions["applications"] = $appConditions
    }

    # Users
    $userConditions = @{}
    $userDef = $PolicyDef["conditions"]["users"]
    if ($userDef) {
        if ($userDef.ContainsKey("includeUsers")) { $userConditions["includeUsers"] = $userDef["includeUsers"] }
        if ($userDef.ContainsKey("excludeUsers")) { $userConditions["excludeUsers"] = $userDef["excludeUsers"] }
        if ($userDef.ContainsKey("includeGroups")) {
            $userConditions["includeGroups"] = @()
            foreach ($gn in $userDef["includeGroups"]) {
                $gid = Get-OrCreateGroup -DisplayName $gn -MailNickname ($gn -replace "\s","")
                if ($gid -notmatch "^WHATIF-") { $userConditions["includeGroups"] += $gid }
            }
        }
        if ($userDef.ContainsKey("excludeGroups")) {
            $userConditions["excludeGroups"] = @()
            foreach ($gn in $userDef["excludeGroups"]) {
                $gid = Get-OrCreateGroup -DisplayName $gn -MailNickname ($gn -replace "\s","")
                if ($gid -notmatch "^WHATIF-") { $userConditions["excludeGroups"] += $gid }
            }
        }
        # Auto-exclude break-glass group
        if ($BreakGlassGroupId -and $BreakGlassGroupId -notmatch "^WHATIF-") {
            if (-not $userConditions.ContainsKey("excludeGroups")) { $userConditions["excludeGroups"] = @() }
            if ($userConditions["excludeGroups"] -notcontains $BreakGlassGroupId) {
                $userConditions["excludeGroups"] += $BreakGlassGroupId
            }
        }
        # Resolve roles
        if ($userDef.ContainsKey("includeRoles")) {
            $userConditions["includeRoles"] = @()
            foreach ($roleName in $userDef["includeRoles"]) {
                if ($script:RoleTemplateMap.ContainsKey($roleName)) {
                    $userConditions["includeRoles"] += $script:RoleTemplateMap[$roleName]
                } else {
                    Write-Warning "Unknown role name '$roleName' in CA policy '$name'. Skipping."
                }
            }
        }
        if ($userDef.ContainsKey("excludeRoles")) {
            $userConditions["excludeRoles"] = @()
            foreach ($roleName in $userDef["excludeRoles"]) {
                if ($script:RoleTemplateMap.ContainsKey($roleName)) {
                    $userConditions["excludeRoles"] += $script:RoleTemplateMap[$roleName]
                }
            }
        }
        # Guests / external users
        if ($userDef.ContainsKey("includeGuestsOrExternalUsers")) {
            $guestDef = $userDef["includeGuestsOrExternalUsers"]
            $guestObj = @{}
            if ($guestDef.ContainsKey("guestTypes")) { $guestObj["guestTypes"] = $guestDef["guestTypes"] }
            if ($guestDef.ContainsKey("externalTenants")) {
                $extDef = $guestDef["externalTenants"]
                $guestObj["externalTenants"] = @{}
                if ($extDef.ContainsKey("membershipKind")) { $guestObj["externalTenants"]["membershipKind"] = $extDef["membershipKind"] }
            }
            $userConditions["includeGuestsOrExternalUsers"] = $guestObj
        }
        if ($userDef.ContainsKey("excludeGuestsOrExternalUsers")) {
            $guestDef = $userDef["excludeGuestsOrExternalUsers"]
            $guestObj = @{}
            if ($guestDef.ContainsKey("guestTypes")) { $guestObj["guestTypes"] = $guestDef["guestTypes"] }
            if ($guestDef.ContainsKey("externalTenants")) {
                $extDef = $guestDef["externalTenants"]
                $guestObj["externalTenants"] = @{}
                if ($extDef.ContainsKey("membershipKind")) { $guestObj["externalTenants"]["membershipKind"] = $extDef["membershipKind"] }
            }
            $userConditions["excludeGuestsOrExternalUsers"] = $guestObj
        }
    }
    $conditions["users"] = $userConditions

    # Client app types
    if ($PolicyDef["conditions"].ContainsKey("clientAppTypes")) {
        $conditions["clientAppTypes"] = $PolicyDef["conditions"]["clientAppTypes"]
    }

    # Sign-in risk
    if ($PolicyDef["conditions"].ContainsKey("signInRiskLevels")) {
        $conditions["signInRiskLevels"] = $PolicyDef["conditions"]["signInRiskLevels"]
    }

    # Locations
    if ($PolicyDef["conditions"].ContainsKey("locations")) {
        $locConditions = @{}
        $locDef = $PolicyDef["conditions"]["locations"]
        if ($locDef.ContainsKey("includeLocations")) {
            $locConditions["includeLocations"] = @()
            foreach ($loc in $locDef["includeLocations"]) {
                if ($loc -eq 'All' -or $loc -eq 'AllTrusted' -or $loc -eq 'MfaTrusted') {
                    $locConditions["includeLocations"] += $loc
                } elseif ($script:NamedLocationCache.ContainsKey($loc)) {
                    $locConditions["includeLocations"] += $script:NamedLocationCache[$loc]
                } else {
                    Write-Warning "Named location '$loc' not found in cache for policy '$name'. Passing as-is."
                    $locConditions["includeLocations"] += $loc
                }
            }
        }
        if ($locDef.ContainsKey("excludeLocations")) {
            $locConditions["excludeLocations"] = @()
            foreach ($loc in $locDef["excludeLocations"]) {
                if ($loc -eq 'All' -or $loc -eq 'AllTrusted' -or $loc -eq 'MfaTrusted') {
                    $locConditions["excludeLocations"] += $loc
                } elseif ($script:NamedLocationCache.ContainsKey($loc)) {
                    $locConditions["excludeLocations"] += $script:NamedLocationCache[$loc]
                } else {
                    Write-Warning "Named location '$loc' not found in cache for policy '$name'. Passing as-is."
                    $locConditions["excludeLocations"] += $loc
                }
            }
        }
        $conditions["locations"] = $locConditions
    }

    # Authentication flows (device code)
    if ($PolicyDef["conditions"].ContainsKey("authenticationFlows")) {
        $flowConditions = @{}
        $flowDef = $PolicyDef["conditions"]["authenticationFlows"]
        if ($flowDef.ContainsKey("deviceCodeFlow")) {
            $flowConditions["deviceCodeFlow"] = @{
                isEnabled = [bool]$flowDef["deviceCodeFlow"]["isEnabled"]
            }
        }
        $conditions["authenticationFlows"] = $flowConditions
    }

    # Build grant controls
    $grantControls = @{}
    $grantDef = $PolicyDef["grantControls"]
    if ($grantDef) {
        $grantControls["operator"] = $grantDef["operator"]
        $grantControls["builtInControls"] = $grantDef["builtInControls"]
        if ($grantDef.ContainsKey("authenticationStrength")) {
            $grantControls["authenticationStrength"] = @{
                id = $grantDef["authenticationStrength"]["id"]
            }
        }
    }

    # Build session controls
    $sessionControls = $null
    if ($PolicyDef.ContainsKey("sessionControls")) {
        $sessionControls = @{}
        $sessDef = $PolicyDef["sessionControls"]
        if ($sessDef.ContainsKey("signInFrequency")) {
            $sessionControls["signInFrequency"] = @{
                value = $sessDef["signInFrequency"]["value"]
                type = $sessDef["signInFrequency"]["type"]
                isEnabled = [bool]$sessDef["signInFrequency"]["isEnabled"]
            }
        }
        if ($sessDef.ContainsKey("persistentBrowser")) {
            $sessionControls["persistentBrowser"] = @{
                mode = $sessDef["persistentBrowser"]["mode"]
                isEnabled = [bool]$sessDef["persistentBrowser"]["isEnabled"]
            }
        }
    }

    $payload = @{
        displayName = $name
        state = $state
        conditions = $conditions
        grantControls = $grantControls
    }
    if ($description) { $payload["description"] = $description }
    if ($sessionControls) { $payload["sessionControls"] = $sessionControls }

    return $payload
}
#endregion


#region Process Tenant Config

function Invoke-WithErrorHandling {
    param(
        [string]$Workload,
        [string]$Control,
        [scriptblock]$Action,
        [string]$Remediation = '',
        [int]$MaxRetries = 3
    )
    $attempt = 0
    do {
        $attempt++
        try {
            & $Action
            return
        } catch {
            $is429 = ($_.Exception.Response.StatusCode -eq 429)
            if ($is429 -and $attempt -lt $MaxRetries) {
                $retryAfterSec = 10
                try {
                    $ra = $_.Exception.Response.Headers['Retry-After']
                    if ($ra) { $retryAfterSec = [int]$ra }
                } catch { }
                $sleepSec = [Math]::Min($retryAfterSec * [Math]::Pow(2, $attempt - 1), 120)
                Write-Warning "[$Workload/$Control] 429 throttled. Retry $attempt/$MaxRetries after $sleepSec s."
                Start-Sleep -Seconds $sleepSec
            } else {
                Add-Result -Workload $Workload -Control $Control -Status 'Error' -Message $_.Exception.Message -Remediation $Remediation
                Write-Warning "[$Workload/$Control] ERROR: $_"
                return
            }
        }
    } while ($attempt -lt $MaxRetries)
}

# =====================================================================
# Admin Center
# =====================================================================
if ($Workloads -contains 'AdminCenter' -and $tenantConfig.ContainsKey('adminCenter')) {
    Write-SectionHeader "M365 Admin Center"
    $ac = $tenantConfig["adminCenter"]

    # Password expiration
    if ($ac.ContainsKey('passwordExpiration')) {
        Invoke-WithErrorHandling -Workload 'AdminCenter' -Control '1.3.1-PasswordExpiration' -Action {
            $org = Get-MgOrganization
            $desired = if ($ac['passwordExpiration'] -eq 'NeverExpire') { 'DisablePasswordExpiration' } else { 'PasswordExpiration' }
            $current = $org.PasswordPolicies
            $pass = ($desired -eq 'DisablePasswordExpiration' -and $current -contains 'DisablePasswordExpiration')
            if ($Mode -eq 'Assess') {
                Add-Result -Workload 'AdminCenter' -Control '1.3.1-PasswordExpiration' `
                    -Status $(if ($pass) { 'Pass' } else { 'Fail' }) `
                    -Message "Current: $current | Desired: $($ac['passwordExpiration'])" `
                    -Remediation "Update-MgOrganization -PasswordPolicies 'DisablePasswordExpiration'"
            } else {
                if (-not $pass -and $PSCmdlet.ShouldProcess($org.DisplayName, "Set password expiration to $($ac['passwordExpiration'])")) {
                    Update-MgOrganization -OrganizationId $org.Id -PasswordPolicies 'DisablePasswordExpiration'
                    Add-Result -Workload 'AdminCenter' -Control '1.3.1-PasswordExpiration' -Status 'Fixed' -Message "Set to NeverExpire"
                } else {
                    Add-Result -Workload 'AdminCenter' -Control '1.3.1-PasswordExpiration' -Status $(if ($pass) { 'Pass' } else { 'Skipped' }) -Message $(if ($pass) { 'Already correct' } else { 'WhatIf/Confirm declined' })
                }
            }
        }
    }
}

# =====================================================================
# Entra ID
# =====================================================================
if ($Workloads -contains 'EntraID' -and $tenantConfig.ContainsKey('entraId')) {
    Write-SectionHeader "Entra ID"
    $entra = $tenantConfig["entraId"]

    # Block tenant creation
    if ($entra.ContainsKey('blockTenantCreation')) {
        Invoke-WithErrorHandling -Workload 'EntraID' -Control '5.1.2.3-BlockTenantCreation' -Action {
            $policy = Get-MgPolicyAuthorizationPolicy
            $current = $policy.DefaultUserRolePermissions.AllowedToCreateTenants
            $desired = -not [bool]$entra['blockTenantCreation']
            if ($Mode -eq 'Assess') {
                Add-Result -Workload 'EntraID' -Control '5.1.2.3-BlockTenantCreation' `
                    -Status $(if ($current -eq $desired) { 'Pass' } else { 'Fail' }) `
                    -Message "AllowedToCreateTenants = $current | Desired = $desired"
            } else {
                if ($current -ne $desired -and $PSCmdlet.ShouldProcess('Authorization Policy', "Set AllowedToCreateTenants = $desired")) {
                    Update-MgPolicyAuthorizationPolicy -DefaultUserRolePermissions @{ AllowedToCreateTenants = $desired }
                    Add-Result -Workload 'EntraID' -Control '5.1.2.3-BlockTenantCreation' -Status 'Fixed' -Message "Set to $desired"
                } else {
                    Add-Result -Workload 'EntraID' -Control '5.1.2.3-BlockTenantCreation' -Status $(if ($current -eq $desired) { 'Pass' } else { 'Skipped' }) -Message $(if ($current -eq $desired) { 'Already correct' } else { 'Declined' })
                }
            }
        }
    }

    # Block user consent / app registration
    if ($entra.ContainsKey('blockUserConsent')) {
        Invoke-WithErrorHandling -Workload 'EntraID' -Control '5.1.2.2-BlockUserConsent' -Action {
            $policy = Get-MgPolicyAuthorizationPolicy
            $current = $policy.DefaultUserRolePermissions.AllowedToCreateApps
            $desired = -not [bool]$entra['blockUserConsent']
            if ($Mode -eq 'Assess') {
                Add-Result -Workload 'EntraID' -Control '5.1.2.2-BlockUserConsent' `
                    -Status $(if ($current -eq $desired) { 'Pass' } else { 'Fail' }) `
                    -Message "AllowedToCreateApps = $current | Desired = $desired"
            } else {
                if ($current -ne $desired -and $PSCmdlet.ShouldProcess('Authorization Policy', "Set AllowedToCreateApps = $desired")) {
                    Update-MgPolicyAuthorizationPolicy -DefaultUserRolePermissions @{ AllowedToCreateApps = $desired }
                    Add-Result -Workload 'EntraID' -Control '5.1.2.2-BlockUserConsent' -Status 'Fixed' -Message "Set to $desired"
                } else {
                    Add-Result -Workload 'EntraID' -Control '5.1.2.2-BlockUserConsent' -Status $(if ($current -eq $desired) { 'Pass' } else { 'Skipped' }) -Message $(if ($current -eq $desired) { 'Already correct' } else { 'Declined' })
                }
            }
        }
    }

    # Max devices per user
    if ($entra.ContainsKey('maxDevicesPerUser')) {
        Invoke-WithErrorHandling -Workload 'EntraID' -Control '5.1.4.2-MaxDevicesPerUser' -Action {
            $regPolicy = Get-MgPolicyDeviceRegistrationPolicy
            $current = $regPolicy.UserDeviceQuota
            $desired = [int]$entra['maxDevicesPerUser']
            if ($Mode -eq 'Assess') {
                Add-Result -Workload 'EntraID' -Control '5.1.4.2-MaxDevicesPerUser' `
                    -Status $(if ($current -le $desired) { 'Pass' } else { 'Fail' }) `
                    -Message "Current quota: $current | Desired max: $desired"
            } else {
                if ($current -ne $desired -and $PSCmdlet.ShouldProcess('Device Registration Policy', "Set max devices to $desired")) {
                    Update-MgPolicyDeviceRegistrationPolicy -UserDeviceQuota $desired
                    Add-Result -Workload 'EntraID' -Control '5.1.4.2-MaxDevicesPerUser' -Status 'Fixed' -Message "Set to $desired"
                } else {
                    Add-Result -Workload 'EntraID' -Control '5.1.4.2-MaxDevicesPerUser' -Status $(if ($current -eq $desired) { 'Pass' } else { 'Skipped' }) -Message $(if ($current -eq $desired) { 'Already correct' } else { 'Declined' })
                }
            }
        }
    }

    # Banned passwords (supports inline list and/or external file)
    $bannedPasswords = [System.Collections.Generic.List[string]]::new()
    if ($entra.ContainsKey('bannedPasswords') -and $entra['bannedPasswords']) {
        foreach ($p in $entra['bannedPasswords']) { $bannedPasswords.Add($p) }
    }
    if ($entra.ContainsKey('bannedPasswordsFile') -and $entra['bannedPasswordsFile']) {
        $pwFile = $entra['bannedPasswordsFile']
        if (-not [System.IO.Path]::IsPathRooted($pwFile)) {
            $pwFile = Join-Path $baselineDir $pwFile
        }
        if (Test-Path $pwFile) {
            $filePasswords = Get-Content $pwFile | ForEach-Object { $_.Trim() } | Where-Object { $_ -and -not $_.StartsWith('#') }
            foreach ($p in $filePasswords) { $bannedPasswords.Add($p) }
        } else {
            Add-Result -Workload 'EntraID' -Control '5.2.3.2-BannedPasswords' -Status 'Error' -Message "Banned passwords file not found: $pwFile"
        }
    }

    if ($bannedPasswords.Count -gt 0) {
        $bannedPasswords = $bannedPasswords | Select-Object -Unique
        Invoke-WithErrorHandling -Workload 'EntraID' -Control '5.2.3.2-BannedPasswords' -Action {
            $settings = Get-MgDirectorySetting | Where-Object { $_.DisplayName -eq 'Password Rule Settings' }
            if (-not $settings) {
                $template = Get-MgDirectorySettingTemplate | Where-Object { $_.DisplayName -eq 'Password Rule Settings' }
                if (-not $template) {
                    Add-Result -Workload 'EntraID' -Control '5.2.3.2-BannedPasswords' -Status 'Error' -Message "Password Rule Settings template not found."
                    return
                }
                $settings = New-MgDirectorySetting -TemplateId $template.Id
            }
            $currentList = ($settings.Values | Where-Object { $_.Name -eq 'BannedPasswordList' }).Value
            $desiredList = ($bannedPasswords -join ', ')
            if ($Mode -eq 'Assess') {
                $hasAll = ($bannedPasswords | ForEach-Object { $currentList -contains $_ }) -notcontains $false
                Add-Result -Workload 'EntraID' -Control '5.2.3.2-BannedPasswords' `
                    -Status $(if ($hasAll) { 'Pass' } else { 'Fail' }) `
                    -Message "Current: $currentList | Desired: $desiredList"
            } else {
                if ($PSCmdlet.ShouldProcess('Directory Setting', "Update banned password list")) {
                    Update-MgDirectorySetting -DirectorySettingId $settings.Id -Values @{
                        BannedPasswordList = $desiredList
                        EnableBannedPasswordCheck = $true
                    }
                    Add-Result -Workload 'EntraID' -Control '5.2.3.2-BannedPasswords' -Status 'Fixed' -Message "Updated banned password list"
                } else {
                    Add-Result -Workload 'EntraID' -Control '5.2.3.2-BannedPasswords' -Status 'Skipped' -Message "Declined"
                }
            }
        }
    }
}

# =====================================================================
# Conditional Access
# =====================================================================
if ($Workloads -contains 'ConditionalAccess' -and $tenantConfig.ContainsKey('conditionalAccess')) {
    Write-SectionHeader "Conditional Access"
    $caConfig = $tenantConfig["conditionalAccess"]
    $reportOnly = if ($caConfig.ContainsKey('reportOnly')) { [bool]$caConfig['reportOnly'] } else { $true }
    $breakGlassGroupName = $caConfig['breakGlassGroup']
    $breakGlassGroupId = $null
    if ($breakGlassGroupName) {
        $breakGlassGroupId = Get-OrCreateGroup -DisplayName $breakGlassGroupName -MailNickname ($breakGlassGroupName -replace "\s","")
    }

    Write-Host "CA Report-Only mode : $reportOnly" -ForegroundColor $(if ($reportOnly) { 'Yellow' } else { 'Green' })
    Write-Host "Break-glass group   : $breakGlassGroupName ($breakGlassGroupId)" -ForegroundColor Cyan

    # Named locations
    if ($caConfig.ContainsKey('namedLocations') -and $caConfig['namedLocations']) {
        Write-Host "`nNamed locations:" -ForegroundColor Cyan
        foreach ($nlDef in $caConfig['namedLocations']) {
            $nlName = $nlDef['displayName']
            $nlType = $nlDef['type']
            $nlCountries = if ($nlDef.ContainsKey('countriesAndRegions')) { $nlDef['countriesAndRegions'] } else { $null }
            $nlUnknown = if ($nlDef.ContainsKey('includeUnknownCountriesAndRegions')) { [bool]$nlDef['includeUnknownCountriesAndRegions'] } else { $false }
            $nlTrusted = if ($nlDef.ContainsKey('isTrusted')) { [bool]$nlDef['isTrusted'] } else { $false }
            $nlRanges = if ($nlDef.ContainsKey('ipRanges')) { $nlDef['ipRanges'] } else { $null }
            $null = Get-OrCreateNamedLocation -DisplayName $nlName -Type $nlType -CountriesAndRegions $nlCountries -IncludeUnknownCountriesAndRegions:$nlUnknown -IsTrusted:$nlTrusted -IpRanges $nlRanges
        }
    }

    $allCAPolicies = $null
    try {
        $allCAPolicies = Get-MgIdentityConditionalAccessPolicy -All
    } catch {
        Write-Warning "Could not enumerate existing CA policies: $_"
    }

    foreach ($caPolicyDef in $caConfig['policies']) {
        $originalName = $caPolicyDef["name"]
        $policyName = Invoke-ApplyMutation -Name $originalName -Mutation $globalMutation
        $cisControl = $caPolicyDef["cisControl"]

        Invoke-WithErrorHandling -Workload 'ConditionalAccess' -Control "$cisControl-$policyName" -Action {
            # Override state if global reportOnly is true
            $effectiveState = $caPolicyDef["state"]
            if ($reportOnly -and $effectiveState -eq 'enabled') {
                $effectiveState = 'enabledForReportingButNotEnforced'
            }

            # Build payload
            $payload = ConvertTo-CAPolicyPayload -PolicyDef $caPolicyDef -BreakGlassGroupId $breakGlassGroupId -Mutation $globalMutation
            $payload['state'] = $effectiveState
            $payload['displayName'] = $policyName

            # Check for existing policy
            $existing = $null
            if ($allCAPolicies) {
                $existing = $allCAPolicies | Where-Object { $_.DisplayName -eq $policyName } | Select-Object -First 1
            }

            if ($Mode -eq 'Assess') {
                if ($existing) {
                    $stateMatch = ($existing.State -eq $effectiveState)
                    Add-Result -Workload 'ConditionalAccess' -Control "$cisControl-$originalName" `
                        -Status $(if ($stateMatch) { 'Pass' } else { 'Fail' }) `
                        -Message "Policy exists. State: $($existing.State) | Desired: $effectiveState" `
                        -Remediation "Update-MgIdentityConditionalAccessPolicy -State '$effectiveState'"
                } else {
                    Add-Result -Workload 'ConditionalAccess' -Control "$cisControl-$originalName" -Status 'Fail' `
                        -Message "Policy does not exist." `
                        -Remediation "Create via Entra admin center or Graph API"
                }
            } else {
                if ($existing) {
                    if ($PSCmdlet.ShouldProcess($policyName, "Update CA policy state to $effectiveState")) {
                        $updatePayload = @{ state = $effectiveState }
                        if ($payload.ContainsKey('description')) { $updatePayload['description'] = $payload['description'] }
                        Update-MgIdentityConditionalAccessPolicy -ConditionalAccessPolicyId $existing.Id -BodyParameter $updatePayload
                        Add-Result -Workload 'ConditionalAccess' -Control "$cisControl-$originalName" -Status 'Fixed' -Message "Updated state to $effectiveState"
                    } else {
                        Add-Result -Workload 'ConditionalAccess' -Control "$cisControl-$originalName" -Status 'Skipped' -Message "Declined"
                    }
                } else {
                    if ($PSCmdlet.ShouldProcess($policyName, "Create CA policy (STATE=$effectiveState)")) {
                        $newPolicy = New-MgIdentityConditionalAccessPolicy -BodyParameter $payload
                        $allCAPolicies = @($allCAPolicies) + $newPolicy
                        Add-Result -Workload 'ConditionalAccess' -Control "$cisControl-$originalName" -Status 'Fixed' `
                            -Message "Created policy '$policyName' in state '$effectiveState' (ID: $($newPolicy.Id))" `
                            -Remediation "Review in Entra admin center > Protection > Conditional Access"
                    } else {
                        Add-Result -Workload 'ConditionalAccess' -Control "$cisControl-$originalName" -Status 'Skipped' -Message "Declined"
                    }
                }
            }
        }
    }
}

# =====================================================================
# Defender
# =====================================================================
if ($Workloads -contains 'Defender' -and $tenantConfig.ContainsKey('defender')) {
    Write-SectionHeader "Defender for Office 365"
    $defender = $tenantConfig["defender"]

    $script:DefenderAcceptedDomains = $null

    # Safe Links
    if ($defender.ContainsKey('safeLinks')) {
        foreach ($sl in $defender['safeLinks']) {
            $name = Invoke-ApplyMutation -Name $sl['name'] -Mutation $globalMutation
            $cis = $sl['cisControl']
            Invoke-WithErrorHandling -Workload 'Defender' -Control "$cis-$name" -Action {
                $policy = Get-SafeLinksPolicy -Identity $name -ErrorAction SilentlyContinue
                if ($Mode -eq 'Assess') {
                    if ($policy) {
                        $pass = $policy.EnableSafeLinksForEmail -and -not $policy.AllowClickThrough
                        Add-Result -Workload 'Defender' -Control "$cis-$name" -Status $(if ($pass) { 'Pass' } else { 'Fail' }) `
                            -Message "SafeLinks exists. Email=$($policy.EnableSafeLinksForEmail) ClickThrough=$($policy.AllowClickThrough)"
                    } else {
                        Add-Result -Workload 'Defender' -Control "$cis-$name" -Status 'Fail' -Message "Policy not found."
                    }
                } else {
                    if ($policy) {
                        if ($PSCmdlet.ShouldProcess($name, 'Update Safe Links')) {
                            Set-SafeLinksPolicy -Identity $name -EnableSafeLinksForEmail $sl['enabled'] -AllowClickThrough $sl['allowClickThrough'] -TrackClicks $sl['trackClicks'] -ScanUrls $sl['scanUrls'] -EnableForInternalSenders $sl['enableForInternalSenders']
                            Add-Result -Workload 'Defender' -Control "$cis-$name" -Status 'Fixed' -Message "Updated Safe Links policy"
                        } else {
                            Add-Result -Workload 'Defender' -Control "$cis-$name" -Status 'Skipped' -Message "Declined"
                        }
                    } else {
                        if ($PSCmdlet.ShouldProcess($name, 'Create Safe Links policy')) {
                            if (-not $script:DefenderAcceptedDomains)
                            {
                                $script:DefenderAcceptedDomains = @((Get-AcceptedDomain -ErrorAction Stop).Name)
                                if (-not $script:DefenderAcceptedDomains) { throw "No accepted domains returned from Exchange Online." }
                            }
                            New-SafeLinksPolicy -Name $name -EnableSafeLinksForEmail $sl['enabled'] -AllowClickThrough $sl['allowClickThrough'] -TrackClicks $sl['trackClicks'] -ScanUrls $sl['scanUrls'] -EnableForInternalSenders $sl['enableForInternalSenders']
                            New-SafeLinksRule -Name "$name-Rule" -SafeLinksPolicy $name -RecipientDomainIs $script:DefenderAcceptedDomains
                            Add-Result -Workload 'Defender' -Control "$cis-$name" -Status 'Fixed' -Message "Created Safe Links policy + rule"
                        } else {
                            Add-Result -Workload 'Defender' -Control "$cis-$name" -Status 'Skipped' -Message "Declined"
                        }
                    }
                }
            }
        }
    }

    # Safe Attachments
    if ($defender.ContainsKey('safeAttachments')) {
        foreach ($sa in $defender['safeAttachments']) {
            $name = Invoke-ApplyMutation -Name $sa['name'] -Mutation $globalMutation
            $cis = $sa['cisControl']
            Invoke-WithErrorHandling -Workload 'Defender' -Control "$cis-$name" -Action {
                $policy = Get-SafeAttachmentPolicy -Identity $name -ErrorAction SilentlyContinue
                if ($Mode -eq 'Assess') {
                    if ($policy) {
                        $pass = $policy.Enable -and ($policy.Action -eq $sa['action'])
                        Add-Result -Workload 'Defender' -Control "$cis-$name" -Status $(if ($pass) { 'Pass' } else { 'Fail' }) `
                            -Message "SafeAttachments exists. Enabled=$($policy.Enable) Action=$($policy.Action)"
                    } else {
                        Add-Result -Workload 'Defender' -Control "$cis-$name" -Status 'Fail' -Message "Policy not found."
                    }
                } else {
                    if ($policy) {
                        if ($PSCmdlet.ShouldProcess($name, 'Update Safe Attachments')) {
                            Set-SafeAttachmentPolicy -Identity $name -Enable $sa['enabled'] -Action $sa['action']
                            Add-Result -Workload 'Defender' -Control "$cis-$name" -Status 'Fixed' -Message "Updated Safe Attachments policy"
                        } else {
                            Add-Result -Workload 'Defender' -Control "$cis-$name" -Status 'Skipped' -Message "Declined"
                        }
                    } else {
                        if ($PSCmdlet.ShouldProcess($name, 'Create Safe Attachments policy')) {
                            if (-not $script:DefenderAcceptedDomains)
                            {
                                $script:DefenderAcceptedDomains = @((Get-AcceptedDomain -ErrorAction Stop).Name)
                                if (-not $script:DefenderAcceptedDomains) { throw "No accepted domains returned from Exchange Online." }
                            }
                            New-SafeAttachmentPolicy -Name $name -Enable $sa['enabled'] -Action $sa['action']
                            New-SafeAttachmentRule -Name "$name-Rule" -SafeAttachmentPolicy $name -RecipientDomainIs $script:DefenderAcceptedDomains
                            Add-Result -Workload 'Defender' -Control "$cis-$name" -Status 'Fixed' -Message "Created Safe Attachments policy + rule"
                        } else {
                            Add-Result -Workload 'Defender' -Control "$cis-$name" -Status 'Skipped' -Message "Declined"
                        }
                    }
                }
            }
        }
    }

    # Anti-Malware
    if ($defender.ContainsKey('antiMalware')) {
        foreach ($am in $defender['antiMalware']) {
            $name = Invoke-ApplyMutation -Name $am['name'] -Mutation $globalMutation
            $cis = $am['cisControl']
            Invoke-WithErrorHandling -Workload 'Defender' -Control "$cis-$name" -Action {
                $policy = Get-MalwareFilterPolicy -Identity $name -ErrorAction SilentlyContinue
                if ($Mode -eq 'Assess') {
                    if ($policy) {
                        Add-Result -Workload 'Defender' -Control "$cis-$name" -Status $(if ($policy.EnableInternalSenderAdminNotifications) { 'Pass' } else { 'Fail' }) `
                            -Message "AntiMalware exists. InternalNotifications=$($policy.EnableInternalSenderAdminNotifications)"
                    } else {
                        Add-Result -Workload 'Defender' -Control "$cis-$name" -Status 'Fail' -Message "Policy not found."
                    }
                } else {
                    if ($policy) {
                        if ($PSCmdlet.ShouldProcess($name, 'Update anti-malware policy')) {
                            Set-MalwareFilterPolicy -Identity $name -EnableInternalSenderAdminNotifications $am['enableInternalNotifications']
                            Add-Result -Workload 'Defender' -Control "$cis-$name" -Status 'Fixed' -Message "Updated anti-malware policy"
                        } else {
                            Add-Result -Workload 'Defender' -Control "$cis-$name" -Status 'Skipped' -Message "Declined"
                        }
                    } else {
                        if ($PSCmdlet.ShouldProcess($name, 'Create anti-malware policy')) {
                            New-MalwareFilterPolicy -Name $name -EnableInternalSenderAdminNotifications $am['enableInternalNotifications']
                            Add-Result -Workload 'Defender' -Control "$cis-$name" -Status 'Fixed' -Message "Created anti-malware policy"
                        } else {
                            Add-Result -Workload 'Defender' -Control "$cis-$name" -Status 'Skipped' -Message "Declined"
                        }
                    }
                }
            }
        }
    }
}

# =====================================================================
# Exchange
# =====================================================================
if ($Workloads -contains 'Exchange' -and $tenantConfig.ContainsKey('exchange')) {
    Write-SectionHeader "Exchange Online"
    $ex = $tenantConfig["exchange"]

    # Mailbox audit org-wide
    if ($ex.ContainsKey('enableMailboxAuditOrgWide')) {
        Invoke-WithErrorHandling -Workload 'Exchange' -Control '6.1.1-MailboxAudit' -Action {
            $orgConfig = Get-OrganizationConfig
            if ($Mode -eq 'Assess') {
                Add-Result -Workload 'Exchange' -Control '6.1.1-MailboxAudit' `
                    -Status $(if ($orgConfig.AuditDisabled -eq $false) { 'Pass' } else { 'Fail' }) `
                    -Message "AuditDisabled = $($orgConfig.AuditDisabled)"
            } else {
                if ($orgConfig.AuditDisabled -ne $false -and $PSCmdlet.ShouldProcess('Organization Config', 'Enable mailbox auditing')) {
                    Set-OrganizationConfig -AuditDisabled $false
                    Add-Result -Workload 'Exchange' -Control '6.1.1-MailboxAudit' -Status 'Fixed' -Message "Enabled org-wide mailbox auditing"
                } else {
                    Add-Result -Workload 'Exchange' -Control '6.1.1-MailboxAudit' -Status $(if ($orgConfig.AuditDisabled -eq $false) { 'Pass' } else { 'Skipped' }) -Message $(if ($orgConfig.AuditDisabled -eq $false) { 'Already enabled' } else { 'Declined' })
                }
            }
        }
    }

    # Block external forwarding
    if ($ex.ContainsKey('blockExternalForwarding') -and $ex['blockExternalForwarding']) {
        Invoke-WithErrorHandling -Workload 'Exchange' -Control '6.2.1-BlockExternalForwarding' -Action {
            $rule = Get-TransportRule | Where-Object { $_.Name -eq 'CIS-Block-External-Forwarding' }
            if ($Mode -eq 'Assess') {
                Add-Result -Workload 'Exchange' -Control '6.2.1-BlockExternalForwarding' `
                    -Status $(if ($rule) { 'Pass' } else { 'Fail' }) `
                    -Message $(if ($rule) { "Rule exists: $($rule.Name)" } else { "No blocking rule found." })
            } else {
                if (-not $rule -and $PSCmdlet.ShouldProcess('Transport Rules', 'Create external forwarding block')) {
                    New-TransportRule -Name 'CIS-Block-External-Forwarding' -FromScope 'InOrganization' -SentToScope 'NotInOrganization' -RejectMessageReasonText 'External forwarding is disabled per security policy.' -RejectMessageEnhancedStatusCode '5.7.1'
                    Add-Result -Workload 'Exchange' -Control '6.2.1-BlockExternalForwarding' -Status 'Fixed' -Message "Created transport rule"
                } else {
                    Add-Result -Workload 'Exchange' -Control '6.2.1-BlockExternalForwarding' -Status $(if ($rule) { 'Pass' } else { 'Skipped' }) -Message $(if ($rule) { 'Already exists' } else { 'Declined' })
                }
            }
        }
    }
}

# =====================================================================
# SharePoint
# =====================================================================
if ($Workloads -contains 'SharePoint' -and $tenantConfig.ContainsKey('sharePoint')) {
    Write-SectionHeader "SharePoint / OneDrive"
    $spo = $tenantConfig["sharePoint"]

    # Connect now that we have the admin URL from YAML
    $spoAdminUrl = $spo['adminUrl']
    if (-not $spoAdminUrl) {
        # Try to infer from tenant
        $domains = Get-MgDomain
        $defaultDomain = $domains | Where-Object { $_.IsInitial } | Select-Object -First 1
        if ($defaultDomain) {
            $spoAdminUrl = "https://$($defaultDomain.Id -replace '\.onmicrosoft\.com','')-admin.sharepoint.com"
        }
    }
    if ($spoAdminUrl) {
        Connect-PnPOnline -Url $spoAdminUrl -Interactive
        Write-Host "Connected to SharePoint admin: $spoAdminUrl" -ForegroundColor Green
    }

    # SharePoint external sharing
    if ($spo.ContainsKey('sharePointExternalSharing')) {
        Invoke-WithErrorHandling -Workload 'SharePoint' -Control '7.x-SharePointExternalSharing' -Action {
            $tenant = Get-PnPTenant
            $desired = $spo['sharePointExternalSharing']
            if ($Mode -eq 'Assess') {
                Add-Result -Workload 'SharePoint' -Control '7.x-SharePointExternalSharing' `
                    -Status $(if ($tenant.SharingCapability -eq $desired) { 'Pass' } else { 'Fail' }) `
                    -Message "Current: $($tenant.SharingCapability) | Desired: $desired"
            } else {
                if ($tenant.SharingCapability -ne $desired -and $PSCmdlet.ShouldProcess('SharePoint Tenant', "Set sharing to $desired")) {
                    Set-PnPTenant -SharingCapability $desired
                    Add-Result -Workload 'SharePoint' -Control '7.x-SharePointExternalSharing' -Status 'Fixed' -Message "Set to $desired"
                } else {
                    Add-Result -Workload 'SharePoint' -Control '7.x-SharePointExternalSharing' -Status $(if ($tenant.SharingCapability -eq $desired) { 'Pass' } else { 'Skipped' }) -Message $(if ($tenant.SharingCapability -eq $desired) { 'Already set' } else { 'Declined' })
                }
            }
        }
    }

    # OneDrive external sharing
    if ($spo.ContainsKey('oneDriveExternalSharing')) {
        Invoke-WithErrorHandling -Workload 'SharePoint' -Control '7.x-OneDriveExternalSharing' -Action {
            $tenant = Get-PnPTenant
            $desired = $spo['oneDriveExternalSharing']
            if ($Mode -eq 'Assess') {
                Add-Result -Workload 'SharePoint' -Control '7.x-OneDriveExternalSharing' `
                    -Status $(if ($tenant.OneDriveSharingCapability -eq $desired) { 'Pass' } else { 'Fail' }) `
                    -Message "Current: $($tenant.OneDriveSharingCapability) | Desired: $desired"
            } else {
                if ($tenant.OneDriveSharingCapability -ne $desired -and $PSCmdlet.ShouldProcess('OneDrive Tenant', "Set sharing to $desired")) {
                    Set-PnPTenant -OneDriveSharingCapability $desired
                    Add-Result -Workload 'SharePoint' -Control '7.x-OneDriveExternalSharing' -Status 'Fixed' -Message "Set to $desired"
                } else {
                    Add-Result -Workload 'SharePoint' -Control '7.x-OneDriveExternalSharing' -Status $(if ($tenant.OneDriveSharingCapability -eq $desired) { 'Pass' } else { 'Skipped' }) -Message $(if ($tenant.OneDriveSharingCapability -eq $desired) { 'Already set' } else { 'Declined' })
                }
            }
        }
    }

    # Default sharing link type
    if ($spo.ContainsKey('defaultSharingLinkType')) {
        Invoke-WithErrorHandling -Workload 'SharePoint' -Control '7.x-DefaultSharingLinkType' -Action {
            $tenant = Get-PnPTenant
            $desired = $spo['defaultSharingLinkType']
            if ($Mode -eq 'Assess') {
                Add-Result -Workload 'SharePoint' -Control '7.x-DefaultSharingLinkType' `
                    -Status $(if ($tenant.DefaultSharingLinkType -eq $desired) { 'Pass' } else { 'Fail' }) `
                    -Message "Current: $($tenant.DefaultSharingLinkType) | Desired: $desired"
            } else {
                if ($tenant.DefaultSharingLinkType -ne $desired -and $PSCmdlet.ShouldProcess('SharePoint Tenant', "Set default link type to $desired")) {
                    Set-PnPTenant -DefaultSharingLinkType $desired
                    Add-Result -Workload 'SharePoint' -Control '7.x-DefaultSharingLinkType' -Status 'Fixed' -Message "Set to $desired"
                } else {
                    Add-Result -Workload 'SharePoint' -Control '7.x-DefaultSharingLinkType' -Status $(if ($tenant.DefaultSharingLinkType -eq $desired) { 'Pass' } else { 'Skipped' }) -Message $(if ($tenant.DefaultSharingLinkType -eq $desired) { 'Already set' } else { 'Declined' })
                }
            }
        }
    }

    # Deny custom scripts
    if ($spo.ContainsKey('denyCustomScripts') -and $spo['denyCustomScripts']) {
        Invoke-WithErrorHandling -Workload 'SharePoint' -Control '7.x-DenyCustomScripts' -Action {
            $tenant = Get-PnPTenant
            if ($Mode -eq 'Assess') {
                Add-Result -Workload 'SharePoint' -Control '7.x-DenyCustomScripts' `
                    -Status $(if ($tenant.DenyAddAndCustomizePages -eq 1) { 'Pass' } else { 'Fail' }) `
                    -Message "DenyAddAndCustomizePages = $($tenant.DenyAddAndCustomizePages)"
            } else {
                if ($tenant.DenyAddAndCustomizePages -ne 1 -and $PSCmdlet.ShouldProcess('SharePoint Tenant', 'Deny custom scripts')) {
                    Set-PnPTenant -DenyAddAndCustomizePages 1
                    Add-Result -Workload 'SharePoint' -Control '7.x-DenyCustomScripts' -Status 'Fixed' -Message "Denied custom scripts"
                } else {
                    Add-Result -Workload 'SharePoint' -Control '7.x-DenyCustomScripts' -Status $(if ($tenant.DenyAddAndCustomizePages -eq 1) { 'Pass' } else { 'Skipped' }) -Message $(if ($tenant.DenyAddAndCustomizePages -eq 1) { 'Already denied' } else { 'Declined' })
                }
            }
        }
    }
}

# =====================================================================
# Teams
# =====================================================================
if ($Workloads -contains 'Teams' -and $tenantConfig.ContainsKey('teams')) {
    Write-SectionHeader "Microsoft Teams"
    $tm = $tenantConfig["teams"]

    # Anonymous meeting join
    if ($tm.ContainsKey('allowAnonymousUsersToJoinMeeting')) {
        Invoke-WithErrorHandling -Workload 'Teams' -Control '8.x-AnonymousMeetingJoin' -Action {
            $policy = Get-CsTeamsMeetingPolicy -Identity Global
            $desired = [bool]$tm['allowAnonymousUsersToJoinMeeting']
            if ($Mode -eq 'Assess') {
                Add-Result -Workload 'Teams' -Control '8.x-AnonymousMeetingJoin' `
                    -Status $(if ($policy.AllowAnonymousUsersToJoinMeeting -eq $desired) { 'Pass' } else { 'Fail' }) `
                    -Message "AllowAnonymousUsersToJoinMeeting = $($policy.AllowAnonymousUsersToJoinMeeting) | Desired = $desired"
            } else {
                if ($policy.AllowAnonymousUsersToJoinMeeting -ne $desired -and $PSCmdlet.ShouldProcess('Teams Global Meeting Policy', "Set anonymous join = $desired")) {
                    Set-CsTeamsMeetingPolicy -Identity Global -AllowAnonymousUsersToJoinMeeting $desired
                    Add-Result -Workload 'Teams' -Control '8.x-AnonymousMeetingJoin' -Status 'Fixed' -Message "Set to $desired"
                } else {
                    Add-Result -Workload 'Teams' -Control '8.x-AnonymousMeetingJoin' -Status $(if ($policy.AllowAnonymousUsersToJoinMeeting -eq $desired) { 'Pass' } else { 'Skipped' }) -Message $(if ($policy.AllowAnonymousUsersToJoinMeeting -eq $desired) { 'Already set' } else { 'Declined' })
                }
            }
        }
    }

    # Anonymous meeting start
    if ($tm.ContainsKey('allowAnonymousUsersToStartMeeting')) {
        Invoke-WithErrorHandling -Workload 'Teams' -Control '8.x-AnonymousMeetingStart' -Action {
            $policy = Get-CsTeamsMeetingPolicy -Identity Global
            $desired = [bool]$tm['allowAnonymousUsersToStartMeeting']
            if ($Mode -eq 'Assess') {
                Add-Result -Workload 'Teams' -Control '8.x-AnonymousMeetingStart' `
                    -Status $(if ($policy.AllowAnonymousUsersToStartMeeting -eq $desired) { 'Pass' } else { 'Fail' }) `
                    -Message "AllowAnonymousUsersToStartMeeting = $($policy.AllowAnonymousUsersToStartMeeting) | Desired = $desired"
            } else {
                if ($policy.AllowAnonymousUsersToStartMeeting -ne $desired -and $PSCmdlet.ShouldProcess('Teams Global Meeting Policy', "Set anonymous start = $desired")) {
                    Set-CsTeamsMeetingPolicy -Identity Global -AllowAnonymousUsersToStartMeeting $desired
                    Add-Result -Workload 'Teams' -Control '8.x-AnonymousMeetingStart' -Status 'Fixed' -Message "Set to $desired"
                } else {
                    Add-Result -Workload 'Teams' -Control '8.x-AnonymousMeetingStart' -Status $(if ($policy.AllowAnonymousUsersToStartMeeting -eq $desired) { 'Pass' } else { 'Skipped' }) -Message $(if ($policy.AllowAnonymousUsersToStartMeeting -eq $desired) { 'Already set' } else { 'Declined' })
                }
            }
        }
    }

    # Federation
    if ($tm.ContainsKey('allowFederatedUsers')) {
        Invoke-WithErrorHandling -Workload 'Teams' -Control '8.x-Federation' -Action {
            $fedConfig = Get-CsTenantFederationConfiguration
            $desired = [bool]$tm['allowFederatedUsers']
            if ($Mode -eq 'Assess') {
                Add-Result -Workload 'Teams' -Control '8.x-Federation' `
                    -Status $(if ($fedConfig.AllowFederatedUsers -eq $desired) { 'Pass' } else { 'Fail' }) `
                    -Message "AllowFederatedUsers = $($fedConfig.AllowFederatedUsers) | Desired = $desired"
            } else {
                if ($fedConfig.AllowFederatedUsers -ne $desired -and $PSCmdlet.ShouldProcess('Teams Federation', "Set AllowFederatedUsers = $desired")) {
                    Set-CsTenantFederationConfiguration -AllowFederatedUsers $desired
                    Add-Result -Workload 'Teams' -Control '8.x-Federation' -Status 'Fixed' -Message "Set to $desired"
                } else {
                    Add-Result -Workload 'Teams' -Control '8.x-Federation' -Status $(if ($fedConfig.AllowFederatedUsers -eq $desired) { 'Pass' } else { 'Skipped' }) -Message $(if ($fedConfig.AllowFederatedUsers -eq $desired) { 'Already set' } else { 'Declined' })
                }
            }
        }
    }
}
# =====================================================================
# Purview
# =====================================================================
if ($Workloads -contains 'Purview' -and $tenantConfig.ContainsKey('purview')) {
    Write-SectionHeader "Microsoft Purview"
    $pv = $tenantConfig["purview"]

    # Audit log search
    if ($pv.ContainsKey('enableAuditLogSearch')) {
        Invoke-WithErrorHandling -Workload 'Purview' -Control '3.1.1-AuditLogSearch' -Action {
            $org = Get-MgOrganization
            $desired = [bool]$pv['enableAuditLogSearch']
            $current = $org.AuditLogEnabled
            if ($Mode -eq 'Assess') {
                Add-Result -Workload 'Purview' -Control '3.1.1-AuditLogSearch' `
                    -Status $(if ($current -eq $desired) { 'Pass' } else { 'Fail' }) `
                    -Message "AuditLogEnabled = $current | Desired = $desired"
            } else {
                if ($PSCmdlet.ShouldProcess('Organization', "Set AuditLogEnabled = $desired")) {
                    # Audit log is typically enabled via Exchange Online, not directly via Graph org
                    Add-Result -Workload 'Purview' -Control '3.1.1-AuditLogSearch' -Status 'Skipped' -Message "Enable via Exchange Admin Center or Set-AdminAuditLogConfig"
                } else {
                    Add-Result -Workload 'Purview' -Control '3.1.1-AuditLogSearch' -Status 'Skipped' -Message "Declined"
                }
            }
        }
    }

    # DLP policies (draft — warn if uncommented)
    if ($pv.ContainsKey('dlpPolicies')) {
        Write-Host "  NOTE: DLP policies found in baseline but require tenant-specific customization." -ForegroundColor Yellow
        Write-Host "        Review and edit the policies in the YAML before deploying." -ForegroundColor DarkGray
        Add-Result -Workload 'Purview' -Control '3.2.x-DLPPolicies' -Status 'Manual' -Message "Draft DLP policies require customization before deployment"
    }

    # Sensitivity labels (draft — warn if uncommented)
    if ($pv.ContainsKey('sensitivityLabels') -or $pv.ContainsKey('sensitivityLabelPolicies')) {
        Write-Host "  NOTE: Sensitivity labels found in baseline but require tenant-specific customization." -ForegroundColor Yellow
        Write-Host "        Review and edit the labels in the YAML before deploying." -ForegroundColor DarkGray
        Add-Result -Workload 'Purview' -Control '3.3.x-SensitivityLabels' -Status 'Manual' -Message "Draft sensitivity labels require customization before deployment"
    }
}

# =====================================================================
# Power BI
# =====================================================================
if ($Workloads -contains 'PowerBI' -and $tenantConfig.ContainsKey('powerBI')) {
    Write-SectionHeader "Power BI"
    $pbi = $tenantConfig["powerBI"]
    Write-Host "  NOTE: Power BI tenant settings are not yet auto-deployed." -ForegroundColor Yellow
    Write-Host "        Review the YAML and configure via Power BI Admin portal or Microsoft365DSC." -ForegroundColor DarkGray
    Add-Result -Workload 'PowerBI' -Control '9.x-PowerBI' -Status 'Manual' -Message "Power BI settings require manual deployment via Admin portal or Microsoft365DSC"
}
#endregion

#region Summary Report
Write-SectionHeader "Summary Report"

$passCount = ($script:Results | Where-Object { $_.Status -eq 'Pass' }).Count
$failCount = ($script:Results | Where-Object { $_.Status -eq 'Fail' }).Count
$manualCount = ($script:Results | Where-Object { $_.Status -eq 'Manual' }).Count
$fixedCount = $script:ChangesMade
$skippedCount = $script:ChangesSkipped
$errorCount = $script:Errors

Write-Host "Mode:        $Mode" -ForegroundColor $(if ($Mode -eq 'Assess') { 'Green' } else { 'Yellow' })
Write-Host "Workloads:   $($Workloads -join ', ')"
if ($script:EffectiveWhatIf) { Write-Host "*** DRY-RUN / WHATIF MODE ***" -ForegroundColor Magenta }
Write-Host ""
Write-Host "Results:"
Write-Host "  Pass:      $passCount" -ForegroundColor Green
Write-Host "  Fail:      $failCount" -ForegroundColor Red
Write-Host "  Manual:    $manualCount" -ForegroundColor Yellow
if ($Mode -eq 'Deploy') {
    Write-Host "  Fixed:     $fixedCount" -ForegroundColor Cyan
    Write-Host "  Skipped:   $skippedCount" -ForegroundColor Yellow
}
Write-Host "  Errors:    $errorCount" -ForegroundColor $(if ($errorCount -gt 0) { 'Red' } else { 'Gray' })
Write-Host ""

# Export results
$timestamp = Get-Date -Format 'yyyyMMdd_HHmmss'
$reportPath = Join-Path $PSScriptRoot "CISM365-Baseline-Report_${Mode}_${timestamp}.csv"
$script:Results | Export-Csv -Path $reportPath -NoTypeInformation -Force
Write-Host "Report saved to: $reportPath" -ForegroundColor Green

# Show failures if in Assess mode
if ($Mode -eq 'Assess' -and $failCount -gt 0) {
    Write-Host "`nFailed checks:" -ForegroundColor Red
    $script:Results | Where-Object { $_.Status -eq 'Fail' } | ForEach-Object {
        Write-Host "  [$($_.Workload)] $($_.Control): $($_.Message)" -ForegroundColor Red
        if ($_.Remediation) { Write-Host "    Remediation: $($_.Remediation)" -ForegroundColor DarkGray }
    }
}

if ($errorCount -gt 0) {
    Write-Host "`nErrors encountered:" -ForegroundColor Red
    $script:Results | Where-Object { $_.Status -eq 'Error' } | ForEach-Object {
        Write-Host "  [$($_.Workload)] $($_.Control): $($_.Message)" -ForegroundColor Red
    }
}

Write-Host "`nDone." -ForegroundColor Green
#endregion