function Get-OrCreateGroup {
    param(
        [string]$DisplayName,
        [string]$MailNickname
    )
    $group = Get-MgGroup -Filter "displayName eq '$DisplayName'" -ConsistencyLevel eventual -CountVariable count
    if ($group) {
        Write-Host "🚫 Group '$DisplayName' already exists. Using existing GroupId: $($group.Id)"
        return $group
    }
    else {
        Write-Host "✅ Creating group '$DisplayName'..."
        return New-MgGroup -BodyParameter @{
            DisplayName     = $DisplayName
            MailEnabled     = $false
            MailNickname    = $MailNickname
            SecurityEnabled = $true
        }
    }
}

# Function to get or create CA policy
function Get-OrCreatePolicy {
    param(
        [string]$DisplayName,
        [hashtable]$BodyParams
    )
    $policy = Get-MgIdentityConditionalAccessPolicy -Filter "displayName eq '$DisplayName'"
    if ($policy) {
        Write-Host "🚫 Policy '$DisplayName' already exists. Skipping creation."
        return $policy
    }
    else {
        Write-Host "✅ Creating policy '$DisplayName'..."
        return New-MgIdentityConditionalAccessPolicy -BodyParameter $BodyParams
    }
}

function Get-OrCreateCountryNamedLocations {
    param(
        [string[]]$Countries  # ISO codes
    )

    $namedLocationIds = @()

    foreach ($c in $Countries) {
        $displayName = "Country - $c"

        # Check if it exists
        $existing = Get-MgIdentityConditionalAccessNamedLocation -Filter "displayName eq '$displayName'" -ErrorAction SilentlyContinue
        if ($existing) {
            Write-Host "🚫 Named Location for $c already exists: $($existing.Id)"
            $namedLocationIds += $existing.Id
        }
        else {
            # Create new Named Location
            $newLoc = New-MgIdentityConditionalAccessNamedLocation -BodyParameter @{
                "@odata.type" = "#microsoft.graph.countryNamedLocation"
                DisplayName = $displayName
                CountriesAndRegions = @($c)
                IncludeUnknownCountriesAndRegions = $false
            }
            Write-Host "✅ Created Named Location for $c $($newLoc.Id)"
            $namedLocationIds += $newLoc.Id
        }
    }

    return $namedLocationIds
}

Connect-MgGraph -NoWelcome -Scopes `
    "Policy.ReadWrite.SecurityDefaults", `
    "Policy.ReadWrite.ConditionalAccess", `
    "Policy.Read.All", `
    "Group.ReadWrite.All", `
    "Directory.ReadWrite.All", `
    "UserAuthenticationMethod.ReadWrite.All", `
    "Application.ReadWrite.All"
    

Write-Host "Connected to Microsoft Graph"

Write-Host "✅ Disabling Security Defaults..."
Update-MgPolicyIdentitySecurityDefaultEnforcementPolicy -BodyParameter @{ isEnabled = $false }

Write-Host "ℹ️  Checking/Creating Exclusion Groups..."

$teamsRoomsGroup = Get-OrCreateGroup -DisplayName "TeamsRoomsExclusions" -MailNickname "TeamsRoomsExcl"
$osTravelGroup   = Get-OrCreateGroup -DisplayName "OSTravelExclusions" -MailNickname "OSTravelExcl"

Write-Host "ℹ️  Teams Rooms GroupId: $($teamsRoomsGroup.Id)"
Write-Host "ℹ️  OS Travel GroupId: $($osTravelGroup.Id)"

Write-Host "ℹ️  Checking/Creating Conditional Access Policy for Country Blocking..."

$AllowedCountries = @("AU","NZ","US","GB","DE","FR")  # Add any ISO codes here

$allowedLocationIds = Get-OrCreateCountryNamedLocations -Countries $AllowedCountries

# 2. Wait for consistency
foreach ($locId in $allowedLocationIds) {
    $retries = 0
    do {
        Start-Sleep -Seconds 2
        $check = Get-MgIdentityConditionalAccessNamedLocation -Filter "id eq '$locId'" -ErrorAction SilentlyContinue
        $retries++
    } while (-not $check -and $retries -lt 10)

    if (-not $check) { throw "⚠️ Named Location $locId not found in directory." }
}

$countryPolicyParams = @{
    DisplayName = "Block countries except allowed"
    State       = "disabled"
    Conditions  = @{
        Users = @{
            IncludeUsers   = @("All")
            ExcludeGroups  = @($osTravelGroup.Id)
            
        }
        Locations = @{
            IncludeLocations = @("All")
            ExcludeLocations = $allowedLocationIds
        }
        Applications = @{
            IncludeApplications = @("All")
            ExcludeApplications = @()
        }
        ClientAppTypes = @("all")
    }
    GrantControls = @{
        Operator        = "OR"
        BuiltInControls = @("block")
    }
}

$Policy = Get-OrCreatePolicy -DisplayName "Block countries except allowed" -BodyParams $countryPolicyParams

# Define the policy from your JSON
$policyParams = @{
    DisplayName = "Require multifactor authentication for admins"
    State       = "enabled"
    Conditions  = @{
        Users = @{
            IncludeRoles = @(
                "62e90394-69f5-4237-9190-012177145e10" # Global Administrator
                "194ae4cb-b126-40b2-bd5b-6091b380977d" # Privileged Role Administrator
                "f28a1f50-f6e7-4571-818b-6a12f2af6b6c" # Security Reader
                "29232cdf-9323-42fd-ade2-1d097af3e4de" # Exchange Administrator
                "b1be1c3e-b65d-4f19-8427-f6fa0d97feb9" # SharePoint Administrator
                "729827e3-9c14-49f7-bb1b-9608f156bbb8" # Security Administrator
                "b0f54661-2d74-4c50-afa3-1ec803f12efe" # Helpdesk Administrator
                "fe930be7-5e62-47db-91af-98c3a49a38b1" # User Administrator
                "c4e39bd9-1100-46d3-8c65-fb160da0071f" # Authentication Administrator
                "9b895d92-2cd3-44c7-9d02-a6ac2d5ea5c3" # Conditional Access Administrator
                "158c047a-c907-4556-b7ef-446551a6b5f7" # Security Operator
                "966707d0-3269-4727-9be2-8c3a10f19b9d" # Reports Reader
                "7be44c8a-adaf-4e2a-84d6-ab2649e08a13" # Billing Administrator
                "e8611ab8-c189-46e8-94e1-60213ab1f814" # Cloud Application Administrator
            )
            ExcludeRoles = @(
                "d29b2b05-8046-44ba-8758-1e26182fcf32" # Directory Synchronization Accounts
            )
        }
        Applications = @{
            IncludeApplications = @("All")
            ExcludeApplications = @()
        }
        ClientAppTypes = @("all")
    }
    GrantControls = @{
        Operator        = "OR"
        BuiltInControls = @("mfa")
    }
}

# Create the Conditional Access policy
$policy = Get-OrCreatePolicy -DisplayName "Require multifactor authentication for admins" -BodyParams $policyParams

# Define the policy from your JSON
$policyParams = @{
    DisplayName = "Require multifactor authentication for all users"
    State       = "enabled"
    Conditions  = @{
        Users = @{
            IncludeUsers = @("All")
            ExcludeGroups  = @($osTravelGroup.Id)
            ExcludeRoles = @(
                "d29b2b05-8046-44ba-8758-1e26182fcf32" # Directory Synchronization Accounts
            )
        }
        Applications = @{
            IncludeApplications = @("All")
            ExcludeApplications = @()
        }
        ClientAppTypes = @("all")
    }
    GrantControls = @{
        Operator        = "OR"
        BuiltInControls = @("mfa")
    }
}

# Create the Conditional Access policy
$policy = Get-OrCreatePolicy -DisplayName "Require multifactor authentication for all users" -BodyParams $policyParams

# Define the policy from your JSON
$policyParams = @{
    DisplayName = "Require multifactor authentication for Azure management"
    State       = "enabled"
    Conditions  = @{
        Users = @{
            IncludeUsers = @("All")
            ExcludeRoles = @(
                "d29b2b05-8046-44ba-8758-1e26182fcf32" # Directory Synchronization Accounts
            )
        }
        Applications = @{
            IncludeApplications = @(
                "797f4846-ba00-4fd7-ba43-dac1f8f63013" # Azure Management
            )
            ExcludeApplications = @()
        }
        ClientAppTypes = @("all")
    }
    GrantControls = @{
        Operator        = "OR"
        BuiltInControls = @("mfa")
    }
}

# Create the Conditional Access policy
$policy = Get-OrCreatePolicy -DisplayName "Require multifactor authentication for Azure management" -BodyParams $policyParams

# Get Graph token for REST calls
$token   = (Get-MgContext).AccessToken
$headers = @{ Authorization = "Bearer $token"; "Content-Type" = "application/json" }

# Get all users
$users = Get-MgUser -All -Property "id,userPrincipalName,accountEnabled,strongAuthenticationRequirements"

foreach ($u in $users) {
    # Skip disabled accounts
    if (-not $u.AccountEnabled) {
        Write-Host "🚫 Skipping disabled user $($u.UserPrincipalName)"
        continue
    }

    # Skip if MFA already not set
    if (-not $u.StrongAuthenticationRequirements -or $u.StrongAuthenticationRequirements.Count -eq 0) {
        Write-Host "➖ No per-user MFA set for $($u.UserPrincipalName) – skipping"
        continue
    }

    # Clear per-user MFA
    $body = @{ strongAuthenticationRequirements = @() } | ConvertTo-Json -Depth 3
    Invoke-RestMethod -Method PATCH `
        -Headers $headers `
        -Uri "https://graph.microsoft.com/v1.0/users/$($u.Id)" `
        -Body $body

    Write-Host "✅ Disabled per-user MFA for $($u.UserPrincipalName)"
}