$ErrorActionPreference = 'Stop';
Set-StrictMode -Version 2.0

[System.Reflection.Assembly]::LoadWithPartialName("System.ServiceProcess") | Out-Null
$ServiceControllerStatus = [System.ServiceProcess.ServiceControllerStatus];

function Get-BatchServiceState { 
	param (
		[Parameter(ValueFromPipeline = $true)]
		[ValidateSet("BatchProcessor", "BatchProcessScheduler")]
		[string]
		$ServiceName
	)

	process {
		if (-not $ServiceName) {
			return
		}

		$service = Get-Service -Name $ServiceName

		if (-not $service) {
			Write-Error "Service not found: $ServiceName"
		}

		return (
			New-Object PSObject -Property  @{
				Service = $service;
				StatusSnapshot = $null;
			});
	}
}

<#
.Synopsis
Get the number of tasks currently in progress in the batch processor.
#>
function Get-ProcessingTasksCount {
	$bpStateDirectory = Get-ProcessorStateDirectory
	$numTasks = @(Get-ChildItem $bpStateDirectory -filter "*.task").Count

	return $numTasks;
}

<#
.Synopsis
Allow the dispatch of processes to the batch processor service to be paused or resumed.
#>
function Set-ProcessDispatchState {
	param (
		[bool]$DispatchTasks
	)
	
	$dispatchSwitch = '-p' # pause
	$requiredState = 'pause'
	if ($DispatchTasks) {
		$dispatchSwitch = '-r'
		$requiredState = 'resume'
	}
	
	$command = Join-Path ( Get-InstallationDirectory ) 'DotNet\bin\BatchProcessTaskConsole.exe';
	
	& $command 'process-dispatch' $dispatchSwitch 
	if ($LastExitCode -ne 0) {
		Write-Error -Message "Could not $requiredState the dispatch of processing tasks to the batch processor; utility exited with code $LastExitCode. See output of command for details."
	}
	
	if (-not $DispatchTasks) {
		Start-Sleep -s 2 # wait a teeny bit to allow the message a chance to get to the scheduler
	}
}

<#
.Synopsis
Wait for currently processing tasks to complete, up to a given number
of minutes.
#>
function Complete-ProcessorTasks {
	param(
		[int]$TimeoutMinutes = 360 # default 6 hours
	)
	
	$numTasks = Get-ProcessingTasksCount;
	$notifiedNumTasks = 0;
	
	$timeoutEnabled = ($TimeoutMinutes -gt 0);
	$maxEndTime = (Get-Date).AddMinutes($TimeoutMinutes);
	
	while ($numTasks -gt 0) {
		if ($numTasks -ne $notifiedNumTasks) {
			Write-Host "    * $($numTasks) task(s) currently processing.";
			$notifiedNumTasks = $numTasks;
		}
		
		Start-Sleep -s 15 # wait for 15 seconds between checks
		
		$numTasks = Get-ProcessingTasksCount
		
		if ($timeoutEnabled -and (Get-Date) -gt $maxEndTime) {
			# deal with timeout separately as this is an error
			Write-Error "Timeout ($TimeoutMinutes minute(s)) expired with $($numTasks) task(s) currently processing. Processing will be resumed, hotfix deployment has failed."
		}
	}
}

function Start-BatchProcessor {
	Write-Host '- Starting Batch Processor service'
	Start-Service 'BatchProcessor'
	
	Write-Host '- Allow tasks to be sent to Batch Processor service for processing'
	Set-ProcessDispatchState -DispatchTasks $true;
}

<#
.Synopsis
Safely stop the batch processor service. This process pauses the dispatch of new tasks
and then waits until currently pending tasks have completed execution.

.Parameter TimeoutMinutes
The number of minutes that the function should wait to allow the running process queue to
empty.

.Parameter Service
The batch processor service state object

.Outputs
Boolean. True when the Batch Processor service was stopped, and can be restarted.
If this method returns false, the Batch Processor service must not be started,
to avoid causing problems in a DR environment.
#>
function Stop-BatchProcessor {
	param(
		[Parameter(Mandatory=$true, ValueFromPipeline=$true)]
		[PSObject]$Service,

		[int]$TimeoutMinutes = 360 # default 6 hours		
	)
	# stop process dispatch, wait for the queue to clear, then stop the service
	Write-Host "Stopping the Batch Processor service safely. The process will wait up to $TimeoutMinutes minute(s) for processing tasks to complete."
		
	Write-Host ' - Pausing the dispatch of new tasks to the batch processor...'
	Set-ProcessDispatchState -DispatchTasks $false;

	try {
		Write-Host ' - Waiting for tasks to complete (this may take hours, so check periodically)...'
		Complete-ProcessorTasks($TimeoutMinutes);
	} catch {
		Write-Warning -Message 'Failure occurred while waiting for processing completion. Restarting batch processing.'
			
		try {
			Set-ProcessDispatchState -DispatchTasks $true;
		} catch {
			Write-Warning "Could not resume processing: $($_.Message)"
		}
			
			
		throw; # rethrow failure
	}
		
	Stop-ServiceSync $Service
}

function Stop-ServiceSync {
	param (
		[Parameter(Mandatory = $true, ValueFromPipeline=$true)]
		[PSObject]
		$Service,
		
		[switch]
		$PassThru
	)

	process {
		$Service |
			Assert-ServiceStatusForDeployment | 
			Out-Null

		if ($Service.Service.Status -ne $ServiceControllerStatus::Stopped) {
			Write-Host "Stopping '$($Service.Service.Name)' service (waiting up to one minute) ..."
			$Service.Service.Refresh();
			$Service.Service.Stop();
			
			$Service.Service.WaitForStatus($ServiceControllerStatus::Stopped, [System.TimeSpan]::FromMinutes(1)) # this will fail with an exception if it doesn't stop in time.
			
			Start-Sleep -s 10 # wait a teeny bit to allow the service process to fully exit
		}

		if ($PassThru) {
			Write-Output $Service
		}
	}
}

function Assert-ServiceStatusForDeployment {
	param (
		[Parameter(Mandatory = $true, ValueFromPipeline = $true)]
		[PSObject]$Service
	)

	begin {
		$expectedServiceStatus =  @($ServiceControllerStatus::Running, $ServiceControllerStatus::Stopped)
	}

	process {
		if ($expectedServiceStatus -notcontains $Service.Service.Status) {
			# Avoid kicking off deployment when services are in transient status. 
			# Given that this cmdlet is mainly for cloud deployment, when services are in such status, 
			# biggest chance is there's already another software install/deployment in progress. 
			# So, instead of giving this sanity check a grace period, a direct failure is preferred.	
			throw [System.InvalidOperationException] "The state of the service $($Service.Service.Name) was unexpected ($($Service.Service.Status)). Please check the service and retry the hotfix application."
		}

		Write-Output $Service
	}
}

function Set-ServiceStateSnapshot {
	param (
		[Parameter(ValueFromPipeline = $true)]
		[PSObject]$Service
	)
	process {
		$Service.Service.Refresh();
		$Service.StatusSnapshot = $Service.Service.Status;

		Write-Output $Service;
	}
}

function Enter-DeploymentMode {
	param (
		[Parameter(Mandatory = $true, ValueFromPipeline = $true)]
		[PSObject]$Service
	)

	begin {
		$processorService = $null;
		$otherServices = @();
	}
	

	process {

		if (-not $Service) {
			return;
		}

		$Service | 
			Assert-ServiceStatusForDeployment | 
			Set-ServiceStateSnapshot | 
			Out-Null;

		if ($Service.Service.Status -eq $ServiceControllerStatus::Running) {
			if ($Service.Service.Name -eq 'BatchProcessor') {
				$processorService = $Service;
			} else {
				$otherServices += $Service;
			}
		}
	}

	end {
		# The Batch Processor service must be shut down first: this will take some time.
		# In the meantime, the scheduler and other services can be up continuing with normal
		# processing. This way there is less impact on the system as a whole.

		if ($processorService) {
			Stop-BatchProcessor $processorService
		}

		if ($otherServices) {
			$otherServices | Stop-ServiceSync
		}

		$allServices = $otherServices
		if ($processorService) {
			$allServices = @($processorService) + $otherServices
		}

		Write-Output $allServices
	}
}

function Exit-DeploymentMode {
	param (
		[Parameter(Mandatory = $true, ValueFromPipeline = $true)]
		[PSObject]
		$Service
	)

	process {

		if ($Service.StatusSnapshot -eq [System.ServiceProcess.ServiceControllerStatus]::Running) {
			if ($Service.Service.Name -eq 'BatchProcessor') {
				Start-BatchProcessor
			} else {			
				Write-Host '- Starting service', $Service.Service.Name
				Start-Service $Service.Service.Name
			}
		}

		Write-Output $Service
	}
}