[string] $CollectionUri = "$env:SYSTEM_TEAMFOUNDATIONCOLLECTIONURI"
[string] $project = "$env:SYSTEM_TEAMPROJECT"
[string] $BuildId = "$env:BUILD_BUILDID"
[string] $BuildNumber = "$env:BUILD_BUILDNUMBER"
[string] $AuthType = "Bearer"
[string] $ReleaseVersionField = "MyProject.ReleaseVersion"

Add-Type -TypeDefinition @"
   public enum VersionType
   {
      Unknown = 0,
      Alpha,
      Beta,
      ReleaseCandidate,
      Release
   }
"@

function Get-VersionType([string] $version)
{
    if ([string]::IsNullOrEmpty($version))
    {
        return [VersionType]::Unknown
    }

    $releaseExpression = "^\d+(\.\d+){2}$"

    if ($version -match $releaseExpression)
    {
        return [VersionType]::Release
    }
    
    $releaseCandidateExpression = "^\d+(\.\d+){2}-beta.+$"
     
    if ($version -match $releaseCandidateExpression)
    {
        return [VersionType]::ReleaseCandidate
    }

    $betaExpression = "^\d+(\.\d+){2}-unstable.+$"
     
    if ($version -match $betaExpression)
    {
        return [VersionType]::Beta
    }

    $alphaExpression = "^\d+(\.\d+){2}-.+$"
     
    if ($version -match $alphaExpression)
    {
        return [VersionType]::Alpha
    }
     
    return [VersionType]::Unknown
}

function Get-Headers
{
    $authorization = "$AuthType $env:SYSTEM_ACCESSTOKEN"
    
    return @{
        Authorization = $authorization
    }
}

function Get-BuildWorkItemIds([string] $collectionUri, [string] $projectName, [string] $buildId)
{
    $headers = Get-Headers
    $restUri = $collectionUri + $projectName + "/_apis/build/builds/" + $buildId + "/workitems?api-version=2.0"
    [string[]] $buildWorkItemIds = @()
 
    $response = Invoke-RestMethod -Uri $restUri -ContentType "application/json" -headers $headers -Method GET
 
    Write-Debug (ConvertTo-Json $response -Depth 100)
 
    $itemCount = $response.count
    
    for($index = 0; $index -lt $itemCount ; $index++)
    {
        $buildWorkItem = $response.value[$index]
        $workItemId = $buildWorkItem.id

        Write-Verbose "Found work item $workItemId linked to build $buildId"

        $buildWorkItemIds += $workItemId
    }

    return $buildWorkItemIds
}

function Get-BuildWorkItems([string] $collectionUri, [string] $projectName, [string] $buildId)
{
    $buildWorkItemIds = @(Get-BuildWorkItemIds $collectionUri $projectName $buildId)
    
    if ($buildWorkItemIds.Length -eq 0)
    {
        return @()
    }

    return Get-WorkItems $collectionUri $buildWorkItemIds
}

function Get-WorkItems([string] $collectionUri, [string[]] $workItemIds)
{
    $headers = Get-Headers
    $ids = ($workItemIds -join ",")
    $restUri = $collectionUri + "_apis/wit/workitems?ids=" + $ids + "&api-version=1.0&`$expand=relations"

    $response = Invoke-RestMethod -Uri $restUri -ContentType "application/json" -headers $headers -Method GET
 
    Write-Debug (ConvertTo-Json $response -Depth 100)

    if ($response.value -eq $null)
    {
        return @()
    }

    return $response.value;
}

function Get-WorkItem([string] $collectionUri, [string] $workItemId)
{
    Write-Debug "Getting work item $workItemId"

    $headers = Get-Headers
    $restUri = $collectionUri + "_apis/wit/workitems/" + $workItemId + "?api-version=1.0&`$expand=relations"

    $response = Invoke-RestMethod -Uri $restUri -ContentType "application/json" -headers $headers -Method GET
 
    Write-Debug (ConvertTo-Json $response -Depth 100)

    return $response;
}

function Get-ParentWorkItemId([object] $workItem)
{
    if ($workItem -eq $null)
    {
        return $null;
    }
    
    if ($workItem.relations -eq $null)
    {
        Write-Debug "$($workItem.fields."System.WorkItemType") #$($workItem.id) does not have any relationships"

        return $null;
    }
    
    $relation = $workItem.relations | Where-Object {$_.rel -eq "System.LinkTypes.Hierarchy-Reverse" }

    if ($relation -eq $null)
    {
        Write-Debug "$($workItem.fields."System.WorkItemType") #$($workItem.id) does not have a parent work item"

        return $null
    }

    $workItemUri = $relation.url
    
    $found = $workItemUri -match "\d+$"
    
    if ($found) 
    {
        $parentWorkItemId = $matches[0]

        Write-Verbose "Found work item $parentWorkItemId as a parent of $($workItem.fields."System.WorkItemType") #$($workItem.id)"
        
        return $parentWorkItemId
    }

    Write-Debug "No parent work item found for $($workItem.fields."System.WorkItemType") #$($workItem.id)"

    return $null
}

function Get-ParentWorkItem([string] $collectionUri, [object] $workItem)
{
    $parentWorkItemId = Get-ParentWorkItemId $workItem

    if ($parentWorkItemId -eq $null)
    {
        return $null
    }

    return Get-WorkItem $collectionUri $parentWorkItemId
}

function Get-WorkItemFound($workItems, [string] $workItemId)
{
    foreach ($workItem in $workItems)
    {
        if ($workItem.id -eq $workItemId)
        {
            return $true
        }
    }

    return $false
}

function Get-IsSupportedParent($workItem)
{
    if ($workItem -eq $null)
    {
        return $false
    }

    $workItemType = $workItem.fields."System.WorkItemType"

    Write-Debug "Work item type is $workItemType"

    if ($workItemType -eq "Product Backlog Item")
    {
        return $true
    }
    
    if ($workItemType -eq "Bug")
    {
        return $true
    }

    return $false
}

function Get-RelatedWorkItems([string] $collectionUri, [string] $projectName, [string] $buildId)
{
    $workItems = @(Get-BuildWorkItems $collectionUri $projectName $buildId)

    Write-Verbose "Found $($workItems.length) work items directly related to build $buildId"

    foreach ($workItem in $workItems)
    {
        $parentWorkItemId = Get-ParentWorkItemId $workItem

        if ($parentWorkItemId -eq $null)
        {
            continue
        }

        if (Get-WorkItemFound $workItems $parentWorkItemId -eq $true)
        {
            Write-Debug "Skipping $parentWorkItemId because it has already been included"

            continue
        }

        $parentWorkItem = Get-WorkItem $collectionUri $parentWorkItemId
        $isParentWorkItemSupported = Get-IsSupportedParent $parentWorkItem

        if ($isParentWorkItemSupported -eq $false)
        {
            Write-Debug "Skipping $($parentWorkItem.fields."System.WorkItemType") #$($parentWorkItem.id) because it is not a PBI or Bug"

            continue
        }

        Write-Verbose "Adding $($parentWorkItem.fields."System.WorkItemType") #$($parentWorkItem.id) to the related work items to process"
    
        $workItems += $parentWorkItem
    }
    
    if ($workItems.length -gt 0)
    {
        Write-Host "Found the following work items related to build $buildId"

        foreach ($workItem in $workItems)
        {
            Write-Host "`t$($workItem.fields."System.WorkItemType") #$($workItem.id)"
        }
    }

    return $workItems
}

function Set-FieldOperation([string[]] $operations, [string] $buildNumber, [PSObject] $workItem, [string] $fieldName, [boolean] $setReleaseTypeOnly)
{
    if ([string]::IsNullOrEmpty($fieldName))
    {
        return $operations
    }

    $fieldValue = $workItem.fields.$fieldName
    $fieldVersion = Get-VersionType $fieldValue
    $buildVersion = Get-VersionType $buildNumber
    
    if ($fieldVersion -ge $buildVersion)
    {
        Write-Verbose "Skipping $($workItem.fields."System.WorkItemType") #$($workItem.id) $fieldName as it is currently ($fieldVersion) $fieldValue"

        return $operations
    }
        
    if ($setReleaseTypeOnly -eq $true -and $buildVersion -ne [VersionType]::Release)
    {
        Write-Verbose "Skipping $($workItem.fields."System.WorkItemType") #$($workItem.id) $fieldName as it is currently ($fieldVersion) $fieldValue and the field only supports Release versions"

        return $operations
    }

    $operations += "{`"op`": `"add`", `"path`": `"/fields/$fieldName`", `"value`": `"$buildNumber`"}"
    Write-Host "Setting $($workItem.fields."System.WorkItemType") #$($workItem.id) $fieldName from ($fieldVersion) $fieldValue to ($buildVersion) $buildNumber"
        
    return $operations
}

function Set-BuildVersionInfo
{
    $buildVersion = Get-VersionType $BuildNumber

    Write-Host "Build version is ($buildVersion) $BuildNumber"

    $headers = Get-Headers
    $buildWorkItems = Get-RelatedWorkItems $CollectionUri $project $BuildId
    $integratedInField = "Microsoft.VSTS.Build.IntegrationBuild"

    foreach ($workItem in $buildWorkItems)
    {
        [string[]] $operations = @()

        $operations = Set-FieldOperation $operations $BuildNumber $workItem $integratedInField $false
        
        $isValidParentWorkItem = Get-IsSupportedParent $workItem

        if ($isValidParentWorkItem -eq $true)
        {
            # We will only set the release version field for PBI and Bug work item types
            $operations = Set-FieldOperation $operations $BuildNumber $workItem $ReleaseVersionField $true
        }
        
        if ($operations.Length -eq 0)
        {
            Write-Host "No changes are being made to $($workItem.fields."System.WorkItemType") #$($workItem.id)"

            continue
        }

        $body = "[" + ($operations -join ",") + "]"
        $restUri = $workItem.url + "?api-version=1.0"
        
        Write-Debug "Updating $($workItem.fields."System.WorkItemType") #$($workItem.id) [$restUri] with the following operations: $body"
        
        try
        {
            $response = Invoke-RestMethod -Uri $restUri -Body $body -ContentType "application/json-patch+json" -headers $headers -Method Patch
            
            Write-Debug (ConvertTo-Json $response -Depth 100)
        }
        catch
        {    
            Write-Error "StatusCode: $_.Exception.Response.StatusCode.value__"
            Write-Error "StatusDescription: $_.Exception.Response.StatusDescription"
            Write-Error $_
        }
    }
}

#________________________________________
#
#      Run the script
#________________________________________

#$VerbosePreference = "Continue"
#$DebugPreference = "SilentlyContinue"

Set-BuildVersionInfo