# Check if required Microsoft.Graph modules are installed
$requiredModules = @(
    "Microsoft.Graph.Authentication",
    "Microsoft.Graph.Users"
)

foreach ($module in $requiredModules) {
    if (!(Get-Module -ListAvailable -Name $module)) {
        Write-Host "$module module not found. Installing..."
        Install-Module $module -Force -AllowClobber -Scope CurrentUser
    } else {
        Write-Host "$module module is already installed."
    }
}

# Import only the required modules
Import-Module Microsoft.Graph.Authentication
Import-Module Microsoft.Graph.Users

# Connect to Microsoft Graph with required permissions
try {
    Connect-MgGraph -Scopes @(
        "User.Read.All",
        "UserAuthenticationMethod.Read.All",
        "Directory.Read.All",
        "User.ReadBasic.All",
        "Policy.Read.All"
    ) -UseDeviceAuthentication:$false
} catch {
    Write-Error "Failed to connect to Microsoft Graph: $($_.Exception.Message)"
    exit 1
}

# Get all enabled users with @domain.onmicrosoft.com domain
try {
    Write-Host "Fetching users..."
    $users = Get-MgUser -All -Property UserPrincipalName,Id,AccountEnabled -ErrorAction Stop | 
             Where-Object { $_.AccountEnabled -eq $true -and $_.UserPrincipalName -like "*@domain.onmicrosoft.com" }
    
    if ($null -eq $users -or $users.Count -eq 0) {
        Write-Warning "No enabled users found with @domain.onmicrosoft.com domain"
        Disconnect-MgGraph
        exit 0
    }
    
    Write-Host "Found $($users.Count) enabled users"
} catch {
    Write-Error "Failed to get users: $($_.Exception.Message)"
    Write-Error "Please ensure you have the correct permissions: User.Read.All, UserAuthenticationMethod.Read.All, Directory.Read.All, User.ReadBasic.All, Policy.Read.All"
    Disconnect-MgGraph
    exit 1
}

$results = @()

foreach ($user in $users) {
    try {
        Write-Host "Checking MFA status for $($user.UserPrincipalName)..."
        
        # Try using the beta endpoint for authentication methods
        $authMethods = Invoke-MgGraphRequest -Method GET -Uri "https://graph.microsoft.com/beta/users/$($user.Id)/authentication/methods" -ErrorAction Stop
        
        # Check if MFA methods are registered (exclude password-only methods)
        $mfaRegistered = $authMethods.value | Where-Object { $_.'@odata.type' -ne "#microsoft.graph.passwordAuthenticationMethod" }
        
        # Determine MFA status
        $mfaStatus = if ($mfaRegistered) { "Enabled" } else { "Disabled" }

        # Add user details to results
        $results += [PSCustomObject]@{
            UserPrincipalName = $user.UserPrincipalName
            MFAStatus         = $mfaStatus
            RegisteredMethods = ($mfaRegistered | ForEach-Object { $_.'@odata.type' -replace "#microsoft.graph.", "" }) -join ", "
        }
    } catch {
        Write-Warning "Failed to get MFA data for $($user.UserPrincipalName): $($_.Exception.Message)"
        
        # Try another approach using conditional access policies
        try {
            Write-Host "Trying conditional access approach for $($user.UserPrincipalName)..."
            
            # Check if user is subject to MFA policies
            $userPolicies = Invoke-MgGraphRequest -Method GET -Uri "https://graph.microsoft.com/beta/users/$($user.Id)/authentication/conditionalAccessPolicies" -ErrorAction Stop
            
            # If user has conditional access policies, assume MFA is required
            $hasMfaPolicy = $userPolicies.value.Count -gt 0
            
            # Add user details to results with conditional access approach
            $results += [PSCustomObject]@{
                UserPrincipalName = $user.UserPrincipalName
                MFAStatus         = if ($hasMfaPolicy) { "Enabled (Policy)" } else { "Unknown" }
                RegisteredMethods = "Conditional Access Check"
            }
        } catch {
            Write-Warning "Conditional access approach also failed for $($user.UserPrincipalName): $($_.Exception.Message)"
            
            # Add user with unknown status
            $results += [PSCustomObject]@{
                UserPrincipalName = $user.UserPrincipalName
                MFAStatus         = "Unknown"
                RegisteredMethods = "Check Failed"
            }
        }
    }
}

# Display results
if ($results.Count -gt 0) {
    $results | Format-Table -AutoSize
} else {
    Write-Host "No MFA data was collected for any users."
}

# Disconnect from Microsoft Graph
Disconnect-MgGraph