param
(
    [Parameter(Mandatory=$true, HelpMessage='The Subscription Ids for all the labs & VMs to check for status, like ("<subscriptionId1>", "<subscriptionId2>")')]
    [string[]] $SubscriptionIds,

    [Parameter(Mandatory=$false, HelpMessage='The name mask for labs, including wildcards - like "WMA*" or "*WMA" ')]
    [string] $DevTestLabNameMask = "*",

    [Parameter(Mandatory=$false, HelpMessage="How many days back to check for 'start' events, the larger number the longer the script takes to run...")]
    [int] $DaysToCheck =  7,

    [Parameter(Mandatory=$false, HelpMessage="If true, output ALL VMs instead of just the ones that haven't been started since DaysToCheck")]
    [bool] $PrintAllVMs = $true
)

# ------------------------------------------------------------
# DTL Module dependency
$AzDtlModuleName = "Az.DevTestLabs2.psm1"
$AzDtlModuleSource = "https://raw.githubusercontent.com/petehauge/DTL-VM-Generator/v2-fixes/Az.DevTestLabs2.psm1"

$modulePath = Join-Path -Path (Resolve-Path ./) -ChildPath $AzDtlModuleName

  if (Test-Path -Path $modulePath) {
    # if the file exists, delete it - just in case there's a newer version, we always download the latest
    Remove-Item -Path $modulePath
}

# Download the powershell module and load it
$WebClient = New-Object System.Net.WebClient
$WebClient.DownloadFile($AzDtlModuleSource, $modulePath)

Import-Module $modulePath -Force

# ------------------------------------------------------------

# iterate through all the subscriptions
foreach ($subId in $SubscriptionIds) {

    $sub = Select-AzSubscription -Subscription $subId

    if (-not $sub) {
        Write-Error "Unable to access Subscription Id '$subId'" 
    }

    # We track the activity log per resource group - this way we don't have to 
    # requery every time, we use this to cache the results
    $activityLogs = @{}

    # Get all the labs & iterate on each one
    Get-AzDtlLab -Name $DevTestLabNameMask | ForEach-Object {
        Write-Host "----------------------------------------------------------" -ForegroundColor Green
        Write-Host " Subscription: $subId , DevTest Lab: $($_.Name)" -ForegroundColor Green

        $vms = Get-AzDtlVm -Lab @{Name = $_.Name; ResourceGroupName = $_.ResourceGroupName}

        # Foreach VM, let's add a property to the object based on if the VM had 'start' events
        foreach ($vm in $vms) {

            # Lab's Resource Group activity log
            if (-not $activityLogs.ContainsKey($vm.ResourceGroupName)) {
                # We haven't seen this resource group yet, let's read the activity logs
                $activityLogs.Add($vm.ResourceGroupName, (Get-AzLog -ResourceGroup $vm.ResourceGroupName -StartTime (Get-Date).AddDays($DaysToCheck * -1) -WarningAction SilentlyContinue -MaxRecord 100000 ))
            }

            # Resource group is always the 4th item with an Azure Resource Id
            $vmRg = $vm.Properties.computeId.Split('/')[4]

            # VM Compute Resource Group activity log
            if (-not $activityLogs.ContainsKey($vmRg)) {
                # We haven't seen this resource group yet, let's read the activity logs
                $activityLogs.Add($vmRg, (Get-AzLog -ResourceGroup $vmRg -StartTime (Get-Date).AddDays($DaysToCheck * -1) -WarningAction SilentlyContinue -MaxRecord 100000))
            }


            if ($vmRg -ne $vm.ResourceGroupName) {
                $combinedActivityLogs = & {
                    $activityLogs[$vm.ResourceGroupName]
                    $activityLogs[$vmRg]
                }
            }
            else {
                $combinedActivityLogs = $activityLogs[$vm.ResourceGroupName]
            }

            # See if we have any DTL "Start" events for this VM
            $lastStartEvent = $combinedActivityLogs | Where-Object { ($_.Authorization.Action -eq "microsoft.devtestlab/labs/virtualmachines/start/action" -and $_.ResourceId -eq $vm.ResourceId) -or  ($_.ResourceId -eq $vm.Properties.computeId -and $_.Authorization.Action -eq "Microsoft.Compute/virtualMachines/start/action")} | Sort-Object -Property @{Expression = {$_.EventTimestamp}; Descending = $True} | Sort-Object -Property @{Expression = {$_.EventTimestamp}; Descending = $True} | Select -First 1
            if ($lastStartEvent) {
                # We found a start event for this VM, let's add the property
                Add-Member -InputObject $vm.Properties -MemberType NoteProperty -Name LastStartDateUTC -Value $lastStartEvent.EventTimestamp
            }

            if ($vm.Properties.AllowClaim) {
                $owner = "*[ CLAIMABLE ]*"
            }
            else {
                if ($vm.Properties.OwnerUserPrincipalName) {
                    $owner = $vm.Properties.OwnerUserPrincipalName
                }
                else {
                    $owner = $vm.Properties.ownerObjectId
                }
            }

            Add-Member -InputObject $vm.Properties -MemberType NoteProperty -Name Owner -Value $owner
        }

        if (-not $PrintAllVMs) {
            $vms = $vms | Where-Object {-not ($_.Properties.LastStartDateUTC -eq $null)}
        }
    
        if ($vms) {

            $vms | Select-Object `
                @{label="Name";expression={$_.Name}},
                @{label="Owner";expression={$_.Properties.Owner}},
                @{label="ProvisioningState";expression={$_.Properties.provisioningState}},
                @{label="PowerState";expression={$_.Properties.lastKnownPowerState}},
                @{label="ArtifactStatus";expression={$_.Properties.ArtifactDeploymentStatus.deploymentStatus}},
                @{label="LastStartDateUTC";expression={$_.Properties.LastStartDateUTC}} `
             | Format-Table
        }
        else {
            Write-Host "$([Environment]::NewLine)   No VM details to output in this lab...$([Environment]::NewLine)"
        }
    }
}


# WORKAROUND:  Delete the azure context that we saved to disk
# Can remove the workaround once this is fixed: https://github.com/Azure/azure-powershell/issues/9448
if (Test-Path (Join-Path "." "AzContext.json")) {
    Remove-Item (Join-Path "." "AzContext.json") -Force
}

# Remove the module just in case so it doesn't conflict with future scripts
Remove-Module -Name "Az.DevTesTLabs2" -ErrorAction SilentlyContinue