<#
.SYNOPSIS
    Deploy scheduled task to monitor AVD session hosts and trigger Entra ID sync
.DESCRIPTION
    This script creates the monitoring script and scheduled task to detect new AVD session hosts
    in the specified OU and trigger AAD Connect delta sync to expedite Entra Hybrid Join.
    It also configures the service account with appropriate batch job logon permissions.
    
    NEW: Now supports querying specific domain controllers to handle replication delays
.NOTES
    FileName:    Deploy-AVDSyncScheduledTask.ps1
    Author:      Modified for AVD environment
    Created:     13/01/2026
    Version:     1.6.0 (Fixed device filtering to query by name first, then filter by creation time)
    
    Requirements:
    - Run on AAD Connect server
    - Service account with appropriate permissions
    - ActiveDirectory PowerShell module
    
    CHANGE LOG:
    v1.6.0 - Fixed device filtering: Now queries by name pattern first, then filters by creation time in PowerShell (resolves issue where whenCreated LDAP filter was unreliable)
#>

#Requires -RunAsAdministrator

[CmdletBinding()]
param(
    [Parameter(Mandatory = $false)]
    [string]$ScriptPath = "C:\Scripts\AVD",
    
    [Parameter(Mandatory = $false)]
    [string]$ServiceAccountName
)

# Ensure script directory exists
if (-not (Test-Path $ScriptPath)) {
    New-Item -Path $ScriptPath -ItemType Directory -Force | Out-Null
    Write-Host "Created script directory: $ScriptPath" -ForegroundColor Green
}

# Prompt for monitored OU
Write-Host "`n====== AVD Sync Task Configuration ======" -ForegroundColor Cyan
Write-Host ""
Write-Host "Please provide the Distinguished Name of the OU to monitor for new AVD session hosts." -ForegroundColor Yellow
Write-Host "Example: OU=Production,OU=Desktop,OU=Pooled,OU=Host Pools,OU=AVD,OU=Azure,DC=test,DC=co,DC=uk" -ForegroundColor Gray
Write-Host ""
Write-Host "====== OU Search Scope Options ======" -ForegroundColor Cyan
Write-Host "The script can search in three ways:" -ForegroundColor White
Write-Host "  1. Subtree   - Searches the OU and ALL nested sub-OUs recursively [RECOMMENDED]" -ForegroundColor Green
Write-Host "  2. Base      - Searches ONLY the specified OU (no sub-OUs)" -ForegroundColor White
Write-Host "  3. OneLevel  - Searches only immediate child OUs (one level deep)" -ForegroundColor White
Write-Host ""
Write-Host "Which search scope would you like? (1/2/3) [Default: 1]" -ForegroundColor Yellow
$scopeChoice = Read-Host "Enter choice"

if ([string]::IsNullOrWhiteSpace($scopeChoice)) {
    $scopeChoice = "1"
}

switch ($scopeChoice) {
    "1" { 
        $searchScope = "Subtree"
        $scopeDescription = "the specified OU and all nested sub-OUs"
    }
    "2" { 
        $searchScope = "Base"
        $scopeDescription = "ONLY the specified OU (no sub-OUs)"
    }
    "3" { 
        $searchScope = "OneLevel"
        $scopeDescription = "only immediate child OUs"
    }
    default { 
        $searchScope = "Subtree"
        $scopeDescription = "the specified OU and all nested sub-OUs"
    }
}

Write-Host "Selected scope: $searchScope - will search $scopeDescription" -ForegroundColor Green
Write-Host ""
$searchBase = Read-Host "Monitored OU (Distinguished Name)"

# Validate OU exists
Write-Host "Validating OU..." -ForegroundColor Cyan
try {
    $ouExists = Get-ADOrganizationalUnit -Identity $searchBase -ErrorAction Stop
    Write-Host "✓ OU validated successfully" -ForegroundColor Green
}
catch {
    Write-Host "✗ Error: Unable to find OU: $searchBase" -ForegroundColor Red
    Write-Host "  $($_.Exception.Message)" -ForegroundColor Red
    exit 1
}

# Prompt for device pattern
Write-Host ""
Write-Host "Please provide the device naming pattern to monitor (supports wildcards)." -ForegroundColor Yellow
Write-Host "Examples: vm-sh-*, vm-*, session-host-*" -ForegroundColor Gray
Write-Host ""
$devicePrefix = Read-Host "Device naming pattern"

if ([string]::IsNullOrWhiteSpace($devicePrefix)) {
    $devicePrefix = "vm-sh-*"
    Write-Host "Using default pattern: $devicePrefix" -ForegroundColor Cyan
}

# NEW: Domain Controller configuration
Write-Host ""
Write-Host "====== Domain Controller Query Configuration ======" -ForegroundColor Cyan
Write-Host ""
Write-Host "To handle AD replication delays, the script can query specific domain controllers." -ForegroundColor Yellow
Write-Host "This ensures newly created computer objects are detected even before full replication." -ForegroundColor White
Write-Host ""
Write-Host "Options:" -ForegroundColor Cyan
Write-Host "  1. Query ALL domain controllers (comprehensive but slower)" -ForegroundColor White
Write-Host "  2. Query SPECIFIC domain controllers (recommended - faster and targeted)" -ForegroundColor Green
Write-Host "  3. Query DEFAULT domain controller only (fastest but may miss new objects)" -ForegroundColor White
Write-Host ""
Write-Host "Which option would you like? (1/2/3) [Default: 2]" -ForegroundColor Yellow
$dcChoice = Read-Host "Enter choice"

if ([string]::IsNullOrWhiteSpace($dcChoice)) {
    $dcChoice = "2"
}

$domainControllers = @()
$dcQueryMode = ""

switch ($dcChoice) {
    "1" {
        $dcQueryMode = "All"
        Write-Host "Will query ALL domain controllers" -ForegroundColor Green
    }
    "2" {
        $dcQueryMode = "Specific"
        Write-Host ""
        Write-Host "Enter domain controller names (one per line, press Enter twice when done):" -ForegroundColor Cyan
        Write-Host "Example: DC01.contoso.local" -ForegroundColor Gray
        Write-Host ""
        
        do {
            $dc = Read-Host "Domain Controller"
            if (-not [string]::IsNullOrWhiteSpace($dc)) {
                # Validate DC exists
                try {
                    $testDC = Get-ADDomainController -Identity $dc -ErrorAction Stop
                    $domainControllers += $testDC.HostName
                    Write-Host "  ✓ Validated: $($testDC.HostName)" -ForegroundColor Green
                }
                catch {
                    Write-Host "  ✗ Warning: Could not validate DC '$dc'" -ForegroundColor Yellow
                    $addAnyway = Read-Host "  Add anyway? (yes/no)"
                    if ($addAnyway -eq "yes") {
                        $domainControllers += $dc
                    }
                }
            }
        } while (-not [string]::IsNullOrWhiteSpace($dc))
        
        if ($domainControllers.Count -eq 0) {
            Write-Host "No domain controllers specified. Falling back to default DC query." -ForegroundColor Yellow
            $dcQueryMode = "Default"
        }
        else {
            Write-Host ""
            Write-Host "Will query these domain controllers:" -ForegroundColor Green
            $domainControllers | ForEach-Object { Write-Host "  • $_" -ForegroundColor White }
        }
    }
    "3" {
        $dcQueryMode = "Default"
        Write-Host "Will query default domain controller only" -ForegroundColor Green
    }
    default {
        $dcQueryMode = "Specific"
        Write-Host "Invalid choice. Defaulting to Specific DC query mode." -ForegroundColor Yellow
    }
}

Write-Host ""
Write-Host "Configuration Summary:" -ForegroundColor Green
Write-Host "Monitored OU: $searchBase" -ForegroundColor White
Write-Host "Search Scope: $searchScope ($scopeDescription)" -ForegroundColor White
Write-Host "Device Pattern: $devicePrefix" -ForegroundColor White
Write-Host "DC Query Mode: $dcQueryMode" -ForegroundColor White
if ($dcQueryMode -eq "Specific" -and $domainControllers.Count -gt 0) {
    Write-Host "Target DCs: $($domainControllers -join ', ')" -ForegroundColor White
}
Write-Host ""

# Convert DC array to PowerShell array syntax for script
$dcArrayString = if ($domainControllers.Count -gt 0) {
    "@(`"$($domainControllers -join '", "')`")"
} else {
    "@()"
}

# Define the monitoring script content
$monitoringScriptContent = @"
<#
.SYNOPSIS
    Monitor AVD session hosts and trigger AAD Connect sync for new devices
.DESCRIPTION
    This script monitors the AVD OU for session hosts matching the specified pattern that were created
    in the last 15 minutes. When detected, it triggers an AAD Connect delta sync to expedite
    Entra Hybrid Join for Azure Virtual Desktop session hosts.
    
    Search Scope: $searchScope - searches $scopeDescription
    DC Query Mode: $dcQueryMode
.NOTES
    FileName:    Sync-NewAVDSessionHostsToEntraID.ps1
    Author:      Modified for AVD environment
    Created:     14/01/2026
    Version:     1.6.0
#>

#Requires -Modules ActiveDirectory

[CmdletBinding()]
param()

# Configuration
`$searchBase = "$searchBase"
`$searchScope = "$searchScope"
`$devicePrefix = "$devicePrefix"
`$minutesLookback = 15
`$replicationWaitSeconds = 30
`$logPath = "C:\Scripts\AVD\Logs"
`$dcQueryMode = "$dcQueryMode"
`$targetDomainControllers = $dcArrayString

# Ensure log directory exists
if (-not (Test-Path `$logPath)) {
    New-Item -Path `$logPath -ItemType Directory -Force | Out-Null
}

# Clean up old log files (older than 14 days)
try {
    `$cutoffDate = (Get-Date).AddDays(-14)
    `$oldLogs = Get-ChildItem -Path `$logPath -Filter "AVDSync_*.log" -File | Where-Object { `$_.LastWriteTime -lt `$cutoffDate }
    
    if (`$oldLogs) {
        foreach (`$log in `$oldLogs) {
            Remove-Item `$log.FullName -Force
        }
    }
}
catch {
    # Silently continue if log cleanup fails - don't want to stop the monitoring
}

# Function to write log entries
function Write-LogEntry {
    param(
        [string]`$Message,
        [ValidateSet('Info', 'Warning', 'Error')]
        [string]`$Level = 'Info'
    )
    
    `$timestamp = Get-Date -Format "yyyy-MM-dd HH:mm:ss"
    `$logFile = Join-Path `$logPath "AVDSync_`$(Get-Date -Format 'yyyyMMdd').log"
    `$logEntry = "`$timestamp [`$Level] `$Message"
    
    Add-Content -Path `$logFile -Value `$logEntry
    
    switch (`$Level) {
        'Warning' { Write-Warning `$Message }
        'Error' { Write-Error `$Message }
        default { Write-Verbose `$Message -Verbose }
    }
}

# Function to query a specific domain controller
function Search-DomainController {
    param(
        [string]`$Server,
        [string]`$SearchBase,
        [string]`$SearchScope,
        [string]`$Filter,
        [string]`$DevicePrefix,
        [datetime]`$TimeThreshold
    )
    
    try {
        Write-LogEntry "Querying DC: `$Server" -Level Info
        
        # First, get all computers matching the name pattern
        `$filterScript = "Name -like '`$DevicePrefix'"
        
        `$computers = Get-ADComputer -Filter `$filterScript ``
            -SearchBase `$SearchBase ``
            -SearchScope `$SearchScope ``
            -Server `$Server ``
            -Properties Created, Modified, Description, OperatingSystem ``
            -ErrorAction Stop
        
        if (`$computers) {
            Write-LogEntry "  Found `$(`$computers.Count) computer(s) matching pattern on `$Server" -Level Info
            
            # Filter by creation time in PowerShell
            `$recentComputers = `$computers | Where-Object { `$_.Created -ge `$TimeThreshold }
            
            if (`$recentComputers) {
                Write-LogEntry "  `$(@(`$recentComputers).Count) computer(s) created within threshold" -Level Info
                return `$recentComputers
            }
            else {
                Write-LogEntry "  No computers created within threshold on `$Server" -Level Info
                return @()
            }
        }
        else {
            Write-LogEntry "  No computers found matching pattern on `$Server" -Level Info
            return @()
        }
    }
    catch {
        Write-LogEntry "  Error querying `${Server}: `$(`$_.Exception.Message)" -Level Warning
        return @()
    }
}

# Start monitoring
Write-LogEntry "========================================" -Level Info
Write-LogEntry "Starting AVD session host monitoring" -Level Info
Write-LogEntry "SearchScope: `$searchScope" -Level Info
Write-LogEntry "DC Query Mode: `$dcQueryMode" -Level Info
Write-LogEntry "========================================" -Level Info

try {
    # Import Active Directory module
    Import-Module ActiveDirectory -ErrorAction Stop
    
    # Calculate time threshold
    `$timeThreshold = [DateTime]::Now.AddMinutes(-`$minutesLookback)
    Write-LogEntry "Checking for devices created after: `$timeThreshold" -Level Info
    
    `$allSessionHosts = @()
    `$queriedDCs = @()
    
    # Determine which DCs to query
    switch (`$dcQueryMode) {
        "All" {
            Write-LogEntry "Retrieving all domain controllers..." -Level Info
            `$dcs = Get-ADDomainController -Filter * | Select-Object -ExpandProperty HostName
            Write-LogEntry "Found `$(`$dcs.Count) domain controller(s)" -Level Info
            `$queriedDCs = `$dcs
        }
        "Specific" {
            if (`$targetDomainControllers.Count -gt 0) {
                Write-LogEntry "Using specified domain controllers: `$(`$targetDomainControllers -join ', ')" -Level Info
                `$queriedDCs = `$targetDomainControllers
            }
            else {
                Write-LogEntry "No target DCs specified, falling back to default" -Level Warning
                `$queriedDCs = @(`$null)  # null will use default DC
            }
        }
        "Default" {
            Write-LogEntry "Using default domain controller" -Level Info
            `$queriedDCs = @(`$null)  # null will use default DC
        }
    }
    
    # Query each DC
    foreach (`$dc in `$queriedDCs) {
        if (`$null -eq `$dc) {
            # Query default DC
            Write-LogEntry "Querying default domain controller..." -Level Info
            
            # First, get all computers matching the name pattern
            `$filter = "Name -like '`$devicePrefix'"
            
            `$computers = Get-ADComputer -Filter `$filter ``
                -SearchBase `$searchBase ``
                -SearchScope `$searchScope ``
                -Properties Created, Modified, Description, OperatingSystem ``
                -ErrorAction Stop
            
            if (`$computers) {
                Write-LogEntry "  Found `$(`$computers.Count) computer(s) matching pattern" -Level Info
                
                # Filter by creation time in PowerShell
                `$recentComputers = `$computers | Where-Object { `$_.Created -ge `$timeThreshold }
                
                if (`$recentComputers) {
                    Write-LogEntry "  `$(@(`$recentComputers).Count) computer(s) created within threshold" -Level Info
                    `$allSessionHosts += `$recentComputers
                }
                else {
                    Write-LogEntry "  No computers created within threshold" -Level Info
                }
            }
            else {
                Write-LogEntry "  No computers found matching pattern" -Level Info
            }
        }
        else {
            # Query specific DC
            `$computers = Search-DomainController ``
                -Server `$dc ``
                -SearchBase `$searchBase ``
                -SearchScope `$searchScope ``
                -DevicePrefix `$devicePrefix ``
                -TimeThreshold `$timeThreshold
            
            if (`$computers) {
                `$allSessionHosts += `$computers
            }
        }
    }
    
    # Remove duplicates (same computer found on multiple DCs)
    `$uniqueSessionHosts = `$allSessionHosts | Sort-Object -Property Name -Unique
    
    if (`$uniqueSessionHosts) {
        Write-LogEntry "Found `$(`$uniqueSessionHosts.Count) unique AVD session host(s) across all queried DCs" -Level Info
        
        `$syncRequired = `$false
        
        foreach (`$sessionHost in `$uniqueSessionHosts) {
            # Calculate age of the computer object
            `$createdAge = ([DateTime]::Now - `$sessionHost.Created).TotalMinutes
            
            Write-LogEntry "Device: `$(`$sessionHost.Name), Created: `$(`$sessionHost.Created), Age: `$([math]::Round(`$createdAge, 2)) minutes" -Level Info
            
            # Verify it's within our threshold and truly new
            if (`$createdAge -le (`$minutesLookback + 1)) {
                Write-LogEntry "Device `$(`$sessionHost.Name) qualifies for sync (created `$([math]::Round(`$createdAge, 2)) minutes ago)" -Level Info
                `$syncRequired = `$true
            }
            else {
                Write-LogEntry "Device `$(`$sessionHost.Name) is outside threshold (created `$([math]::Round(`$createdAge, 2)) minutes ago)" -Level Warning
            }
        }
        
        if (`$syncRequired) {
            # Wait for AD replication
            Write-LogEntry "Waiting `$replicationWaitSeconds seconds for AD replication..." -Level Info
            Start-Sleep -Seconds `$replicationWaitSeconds
            
            # Trigger AAD Connect sync
            Write-LogEntry "Triggering AAD Connect Delta Sync" -Level Info
            
            try {
                Start-ADSyncSyncCycle -PolicyType Delta -ErrorAction Stop
                Write-LogEntry "AAD Connect Delta Sync triggered successfully" -Level Info
            }
            catch {
                Write-LogEntry "Failed to trigger AAD Connect sync: `$(`$_.Exception.Message)" -Level Error
                throw
            }
        }
        else {
            Write-LogEntry "No devices qualify for sync at this time" -Level Info
        }
    }
    else {
        Write-LogEntry "No new AVD session hosts found matching pattern '`$devicePrefix' in the last `$minutesLookback minutes" -Level Info
    }
}
catch {
    Write-LogEntry "Error during monitoring: `$(`$_.Exception.Message)" -Level Error
    Write-LogEntry "Stack Trace: `$(`$_.ScriptStackTrace)" -Level Error
    exit 1
}

Write-LogEntry "Monitoring cycle completed" -Level Info
Write-LogEntry "========================================" -Level Info
"@

# Save the monitoring script
$monitoringScriptPath = Join-Path $ScriptPath "Sync-NewAVDSessionHostsToEntraID.ps1"
$monitoringScriptContent | Out-File -FilePath $monitoringScriptPath -Encoding UTF8 -Force

Write-Host "Monitoring script created: $monitoringScriptPath" -ForegroundColor Green

# Get service account credentials
if ([string]::IsNullOrEmpty($ServiceAccountName)) {
    Write-Host "`nPlease enter the service account credentials" -ForegroundColor Cyan
    Write-Host "Recommended: svc-avd-entra-sync or svc-aadconnect-avd" -ForegroundColor Yellow
    $credential = Get-Credential -Message "Enter service account credentials"
    $serviceAccountUser = $credential.UserName
}
else {
    $credential = Get-Credential -UserName $ServiceAccountName -Message "Enter password for $ServiceAccountName"
    $serviceAccountUser = $ServiceAccountName
}

# Extract domain and username, handle different formats
Write-Host "`nResolving account format..." -ForegroundColor Gray
$accountParts = $serviceAccountUser -split '\\'

if ($accountParts.Count -eq 2) {
    # Domain\Username format
    $domain = $accountParts[0]
    $username = $accountParts[1]
    $fullyQualifiedUser = $serviceAccountUser
    Write-Host "Detected format: Domain\Username" -ForegroundColor Gray
}
elseif ($serviceAccountUser -like "*@*") {
    # UPN format (user@domain.com)
    Write-Host "Detected format: UPN (User Principal Name)" -ForegroundColor Gray
    $upnParts = $serviceAccountUser.Split('@')
    $username = $upnParts[0]
    $upnDomain = $upnParts[1]
    
    # Try to resolve UPN to domain\username format
    Write-Host "Attempting to resolve UPN to domain account..." -ForegroundColor Gray
    try {
        # Try using ADSI to find the account
        $searcher = New-Object System.DirectoryServices.DirectorySearcher
        $searcher.Filter = "(&(objectCategory=User)(userPrincipalName=$serviceAccountUser))"
        $searcher.PropertiesToLoad.Add("sAMAccountName") | Out-Null
        $searcher.PropertiesToLoad.Add("distinguishedName") | Out-Null
        
        $result = $searcher.FindOne()
        
        if ($result) {
            $samAccountName = $result.Properties["samaccountname"][0]
            $distinguishedName = $result.Properties["distinguishedname"][0]
            
            # Extract domain from DN
            $dnParts = $distinguishedName -split ','
            $dcParts = $dnParts | Where-Object { $_ -like "DC=*" }
            $domain = ($dcParts -join '.' -replace 'DC=', '')
            
            # Get NetBIOS domain name
            try {
                $rootDSE = [ADSI]"LDAP://RootDSE"
                $configNC = $rootDSE.configurationNamingContext
                $searcher2 = New-Object System.DirectoryServices.DirectorySearcher
                $searcher2.SearchRoot = [ADSI]"LDAP://CN=Partitions,$configNC"
                $searcher2.Filter = "(&(objectCategory=crossRef)(nCName=$($rootDSE.defaultNamingContext)))"
                $searcher2.PropertiesToLoad.Add("nETBIOSName") | Out-Null
                $partition = $searcher2.FindOne()
                
                if ($partition) {
                    $domain = $partition.Properties["netbiosname"][0]
                }
            }
            catch {
                # If we can't get NetBIOS name, try current domain
                Write-Host "  Could not resolve NetBIOS domain name, using current domain" -ForegroundColor DarkGray
                $domain = $env:USERDOMAIN
            }
            
            $fullyQualifiedUser = "$domain\$samAccountName"
            Write-Host "  ✓ Resolved UPN to: $fullyQualifiedUser" -ForegroundColor Green
        }
        else {
            Write-Host "  ✗ Could not resolve UPN via ADSI" -ForegroundColor Yellow
            $fullyQualifiedUser = "$env:USERDOMAIN\$username"
            Write-Host "  Using fallback: $fullyQualifiedUser" -ForegroundColor Gray
        }
    }
    catch {
        Write-Host "  ✗ Error resolving UPN: $($_.Exception.Message)" -ForegroundColor Yellow
        $fullyQualifiedUser = "$env:USERDOMAIN\$username"
        Write-Host "  Using fallback: $fullyQualifiedUser" -ForegroundColor Gray
    }
}
else {
    # Just username
    Write-Host "Detected format: Username only" -ForegroundColor Gray
    $username = $serviceAccountUser
    $fullyQualifiedUser = "$env:USERDOMAIN\$username"
}

Write-Host "Final account format: $fullyQualifiedUser" -ForegroundColor Cyan

# Prompt for batch logon rights configuration
Write-Host "`n====== Batch Logon Rights Configuration ======" -ForegroundColor Cyan
Write-Host ""
Write-Host "The service account requires 'Log on as a batch job' rights on this local server." -ForegroundColor Yellow
Write-Host "This permission allows the scheduled task to run under the service account." -ForegroundColor White
Write-Host ""
Write-Host "Would you like to grant '$serviceAccountUser' batch logon rights now? (yes/no)" -ForegroundColor Cyan
$grantBatchRights = Read-Host

if ($grantBatchRights -eq "yes") {
    Write-Host ""
    Write-Host "Configuring batch logon rights..." -ForegroundColor Cyan
    
    try {
        # First, validate and resolve the account to get its SID
        Write-Host "Validating account: $fullyQualifiedUser" -ForegroundColor Gray
        
        $accountSID = $null
        $resolvedAccount = $null
        
        # Try different account formats in order of likelihood
        $formatsToTry = @(
            $fullyQualifiedUser,           # Domain\Username (most likely to work)
            $serviceAccountUser,            # Original format entered (could be UPN)
            $username,                      # Just the username
            "$env:COMPUTERNAME\$username"   # Local computer account (unlikely but possible)
        )
        
        foreach ($format in $formatsToTry) {
            if ([string]::IsNullOrWhiteSpace($format)) { continue }
            
            try {
                Write-Host "  Trying format: $format" -ForegroundColor Gray
                $ntAccount = New-Object System.Security.Principal.NTAccount($format)
                $accountSID = $ntAccount.Translate([System.Security.Principal.SecurityIdentifier])
                $resolvedAccount = $format
                Write-Host "  ✓ Account resolved successfully!" -ForegroundColor Green
                Write-Host "  Account SID: $($accountSID.Value)" -ForegroundColor Gray
                break
            }
            catch {
                Write-Host "  ✗ Failed: $($_.Exception.Message.Split("`n")[0])" -ForegroundColor DarkGray
                continue
            }
        }
        
        if (-not $accountSID) {
            throw "Unable to resolve account '$serviceAccountUser' to a valid SID. Please check the account name format."
        }
        
        Write-Host ""
        Write-Host "Resolved account: $resolvedAccount" -ForegroundColor Green
        
        # Define the C# code for LSA functions
        $lsaDefinition = @"
using System;
using System.Runtime.InteropServices;
using System.Text;

namespace LSA
{
    using LSA_HANDLE = IntPtr;

    [StructLayout(LayoutKind.Sequential)]
    public struct LSA_OBJECT_ATTRIBUTES
    {
        public int Length;
        public IntPtr RootDirectory;
        public IntPtr ObjectName;
        public int Attributes;
        public IntPtr SecurityDescriptor;
        public IntPtr SecurityQualityOfService;
    }

    [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Unicode)]
    public struct LSA_UNICODE_STRING
    {
        public ushort Length;
        public ushort MaximumLength;
        public string Buffer;
    }

    public class LSAWrapper
    {
        [DllImport("advapi32.dll", CharSet = CharSet.Unicode, SetLastError = true)]
        public static extern uint LsaOpenPolicy(
            IntPtr SystemName,
            ref LSA_OBJECT_ATTRIBUTES ObjectAttributes,
            uint AccessMask,
            out IntPtr PolicyHandle);

        [DllImport("advapi32.dll", CharSet = CharSet.Unicode, SetLastError = true)]
        public static extern uint LsaAddAccountRights(
            LSA_HANDLE PolicyHandle,
            IntPtr AccountSid,
            LSA_UNICODE_STRING[] UserRights,
            uint CountOfRights);

        [DllImport("advapi32.dll")]
        public static extern uint LsaClose(IntPtr ObjectHandle);

        [DllImport("advapi32.dll")]
        public static extern int LsaNtStatusToWinError(uint Status);

        public const uint POLICY_CREATE_ACCOUNT = 0x00000010;
        public const uint POLICY_LOOKUP_NAMES = 0x00000800;
    }
}
"@

        # Add the type if not already added
        if (-not ([System.Management.Automation.PSTypeName]'LSA.LSAWrapper').Type) {
            Write-Host "Loading LSA functions..." -ForegroundColor Gray
            Add-Type -TypeDefinition $lsaDefinition
        }

        # Open LSA Policy
        Write-Host "Opening LSA policy..." -ForegroundColor Gray
        $objectAttributes = New-Object LSA.LSA_OBJECT_ATTRIBUTES
        $objectAttributes.Length = 0
        $objectAttributes.RootDirectory = [IntPtr]::Zero
        $objectAttributes.ObjectName = [IntPtr]::Zero
        $objectAttributes.Attributes = 0
        $objectAttributes.SecurityDescriptor = [IntPtr]::Zero
        $objectAttributes.SecurityQualityOfService = [IntPtr]::Zero

        $policyHandle = [IntPtr]::Zero
        $result = [LSA.LSAWrapper]::LsaOpenPolicy(
            [IntPtr]::Zero,
            [ref]$objectAttributes,
            0x00000810,  # POLICY_CREATE_ACCOUNT | POLICY_LOOKUP_NAMES
            [ref]$policyHandle
        )

        if ($result -ne 0) {
            $winError = [LSA.LSAWrapper]::LsaNtStatusToWinError($result)
            throw "LsaOpenPolicy failed with error code: $winError"
        }

        Write-Host "Adding batch logon right for SID: $($accountSID.Value)" -ForegroundColor Gray
        
        try {
            # Convert SID to byte array
            $sidBytes = New-Object byte[] $accountSID.BinaryLength
            $accountSID.GetBinaryForm($sidBytes, 0)
            
            # Allocate memory for SID
            $pSid = [System.Runtime.InteropServices.Marshal]::AllocHGlobal($sidBytes.Length)
            [System.Runtime.InteropServices.Marshal]::Copy($sidBytes, 0, $pSid, $sidBytes.Length)

            # Create LSA_UNICODE_STRING for the right
            $right = New-Object LSA.LSA_UNICODE_STRING
            $right.Buffer = "SeBatchLogonRight"
            $right.Length = [UInt16]($right.Buffer.Length * 2)
            $right.MaximumLength = [UInt16](($right.Buffer.Length + 1) * 2)

            $rights = @($right)

            # Add the right
            $result = [LSA.LSAWrapper]::LsaAddAccountRights($policyHandle, $pSid, $rights, 1)
            
            if ($result -ne 0) {
                $winError = [LSA.LSAWrapper]::LsaNtStatusToWinError($result)
                throw "LsaAddAccountRights failed with error code: $winError"
            }

            Write-Host ""
            Write-Host "✓ Successfully granted batch logon rights to: $resolvedAccount" -ForegroundColor Green
            Write-Host "  Account SID: $($accountSID.Value)" -ForegroundColor Gray
        }
        finally {
            # Cleanup
            if ($pSid -ne [IntPtr]::Zero) {
                [System.Runtime.InteropServices.Marshal]::FreeHGlobal($pSid)
            }
            [LSA.LSAWrapper]::LsaClose($policyHandle) | Out-Null
        }
    }
    catch {
        Write-Host ""
        Write-Host "✗ Error configuring batch logon rights: $($_.Exception.Message)" -ForegroundColor Red
        Write-Host ""
        Write-Host "Please configure this manually before the scheduled task will work:" -ForegroundColor Yellow
        Write-Host "  1. Open Local Security Policy (secpol.msc)" -ForegroundColor White
        Write-Host "  2. Navigate to: Security Settings > Local Policies > User Rights Assignment" -ForegroundColor White
        Write-Host "  3. Open 'Log on as a batch job'" -ForegroundColor White
        Write-Host "  4. Click 'Add User or Group'" -ForegroundColor White
        Write-Host "  5. Enter: $serviceAccountUser" -ForegroundColor White
        Write-Host "  6. Click OK and apply the changes" -ForegroundColor White
        Write-Host ""
        Write-Host "Press any key to continue with scheduled task creation..." -ForegroundColor Cyan
        $null = $Host.UI.RawUI.ReadKey("NoEcho,IncludeKeyDown")
    }
}
else {
    Write-Host "Skipping batch logon rights configuration." -ForegroundColor Yellow
    Write-Host ""
    Write-Host "IMPORTANT: You must configure batch logon rights manually before the task will run:" -ForegroundColor Red
    Write-Host "  1. Open Local Security Policy (secpol.msc)" -ForegroundColor White
    Write-Host "  2. Navigate to: Security Settings > Local Policies > User Rights Assignment" -ForegroundColor White
    Write-Host "  3. Open 'Log on as a batch job'" -ForegroundColor White
    Write-Host "  4. Click 'Add User or Group'" -ForegroundColor White
    Write-Host "  5. Enter: $serviceAccountUser" -ForegroundColor White
    Write-Host "  6. Click OK and apply the changes" -ForegroundColor White
    Write-Host ""
}

# Create scheduled task
Write-Host "`nCreating scheduled task..." -ForegroundColor Cyan

$taskName = "Sync-NewAVDSessionHostsToEntraID"
$taskDescription = "Monitors AVD OU ($searchScope scope) for new session hosts matching pattern '$devicePrefix' created in the last 15 minutes. Queries domain controllers ($dcQueryMode mode) and triggers AAD Connect delta sync to expedite Entra Hybrid Join"

# Define task action
$action = New-ScheduledTaskAction `
    -Execute 'Powershell.exe' `
    -Argument "-NoProfile -WindowStyle Hidden -ExecutionPolicy Bypass -File `"$monitoringScriptPath`"" `
    -WorkingDirectory $ScriptPath

# Define task trigger (every 5 minutes indefinitely, starting immediately)
# Use a time 1 minute in the past to ensure it triggers immediately
$triggerTime = (Get-Date).AddMinutes(-1)
$trigger = New-ScheduledTaskTrigger -Once -At $triggerTime -RepetitionInterval (New-TimeSpan -Minutes 5)
# Remove the duration limit to run indefinitely
$trigger.Repetition.Duration = ""

# Define task settings (including run task as soon as possible if missed)
$settings = New-ScheduledTaskSettingsSet `
    -AllowStartIfOnBatteries `
    -DontStopIfGoingOnBatteries `
    -StartWhenAvailable `
    -RunOnlyIfNetworkAvailable `
    -MultipleInstances IgnoreNew `
    -ExecutionTimeLimit (New-TimeSpan -Hours 1)

# Register the task
try {
    # Check if task already exists
    $existingTask = Get-ScheduledTask -TaskName $taskName -ErrorAction SilentlyContinue
    
    if ($existingTask) {
        Write-Host "Scheduled task '$taskName' already exists. Unregistering..." -ForegroundColor Yellow
        Unregister-ScheduledTask -TaskName $taskName -Confirm:$false
    }
    
    $task = Register-ScheduledTask `
        -Action $action `
        -Trigger $trigger `
        -Settings $settings `
        -TaskName $taskName `
        -Description $taskDescription `
        -User $credential.UserName `
        -Password $credential.GetNetworkCredential().Password `
        -RunLevel Highest
    
    Write-Host "`n✓ Scheduled task created successfully!" -ForegroundColor Green
    Write-Host "Task Name: $taskName" -ForegroundColor Cyan
    Write-Host "Frequency: Every 5 minutes, indefinitely" -ForegroundColor Cyan
    Write-Host "Run As: $($credential.UserName)" -ForegroundColor Cyan
    
    # Display next run time
    $taskInfo = Get-ScheduledTaskInfo -TaskName $taskName
    Write-Host "Next Run: $($taskInfo.NextRunTime)" -ForegroundColor Cyan
}
catch {
    Write-Host "✗ Error creating scheduled task: $($_.Exception.Message)" -ForegroundColor Red
    
    # Check if it's a logon rights issue
    if ($_.Exception.Message -like "*logon*" -or $_.Exception.Message -like "*0x80070569*") {
        Write-Host "`nThis appears to be a batch logon rights issue." -ForegroundColor Yellow
        Write-Host "Please configure batch logon rights manually:" -ForegroundColor Red
        Write-Host "  1. Open Local Security Policy (secpol.msc)" -ForegroundColor White
        Write-Host "  2. Navigate to: Security Settings > Local Policies > User Rights Assignment" -ForegroundColor White
        Write-Host "  3. Open 'Log on as a batch job'" -ForegroundColor White
        Write-Host "  4. Click 'Add User or Group'" -ForegroundColor White
        Write-Host "  5. Enter: $serviceAccountUser" -ForegroundColor White
        Write-Host "  6. Click OK and apply the changes" -ForegroundColor White
        Write-Host "  7. Run this script again" -ForegroundColor White
    }
    else {
        Write-Host "`nPlease check the error message above and try again." -ForegroundColor Yellow
    }
    exit 1
}

# Summary
Write-Host "`n====== Deployment Summary ======" -ForegroundColor Green
Write-Host "Script Location: $monitoringScriptPath" -ForegroundColor White
Write-Host "Log Location: C:\Scripts\AVD\Logs" -ForegroundColor White
Write-Host "Scheduled Task: $taskName" -ForegroundColor White
Write-Host "Monitored OU: $searchBase" -ForegroundColor White
Write-Host "Search Scope: $searchScope ($scopeDescription)" -ForegroundColor White
Write-Host "Device Pattern: $devicePrefix" -ForegroundColor White
Write-Host "DC Query Mode: $dcQueryMode" -ForegroundColor White
if ($dcQueryMode -eq "Specific" -and $domainControllers.Count -gt 0) {
    Write-Host "Target DCs: $($domainControllers -join ', ')" -ForegroundColor White
}
Write-Host "Service Account: $serviceAccountUser" -ForegroundColor White
Write-Host "`nThe task will run every 5 minutes and check for new AVD session hosts." -ForegroundColor Yellow
Write-Host "When detected, it will trigger an AAD Connect delta sync automatically." -ForegroundColor Yellow
Write-Host "`nQuerying multiple DCs helps detect new computers before full AD replication completes." -ForegroundColor Cyan
Write-Host "`nYou can manually test the script by running:" -ForegroundColor Cyan
Write-Host "  $monitoringScriptPath" -ForegroundColor White