#Requires -Version 5.1

[CmdletBinding(SupportsShouldProcess = $true)]
param(
    [Parameter(Mandatory = $true)]
    [ValidatePattern('^[A-Za-z0-9._:-]+$')]
    [string] $ServerHost,

    [ValidateRange(1, 65535)]
    [int] $ServerPort = 22,

    [ValidatePattern('^[A-Za-z0-9._-]+$')]
    [string] $RemoteUser = 'caiwu-tunnel',

    [Parameter(Mandatory = $true)]
    [string] $PrivateKeyPath,

    [Parameter(Mandatory = $true)]
    [ValidatePattern('^[A-Za-z0-9._:-]+$')]
    [string] $RemoteBindAddress,

    [Parameter(Mandatory = $true)]
    [ValidateRange(1, 65535)]
    [int] $RemoteBindPort,

    [Parameter(Mandatory = $true)]
    [ValidatePattern('^[A-Za-z0-9._:-]+$')]
    [string] $UpstreamHost,

    [ValidateRange(1, 65535)]
    [int] $UpstreamPort = 1433,

    [ValidateRange(10, 300)]
    [int] $RetrySeconds = 30,

    [ValidatePattern('^[a-z0-9_-]+$')]
    [string] $SourceKey = 'default',

    [ValidatePattern('^[A-Za-z0-9._-]+$')]
    [string] $TaskName = 'CaiWuXiTong-UpstreamTunnel',

    [string] $ConfigPath
)

Set-StrictMode -Version Latest
$ErrorActionPreference = 'Stop'

function ConvertTo-TaskArgument {
    param([Parameter(Mandatory = $true)][string] $Value)

    return '"' + $Value.Replace('"', '\"') + '"'
}

$tunnelScript = (Resolve-Path -LiteralPath (Join-Path $PSScriptRoot 'keep-upstream-tunnel.ps1') -ErrorAction Stop).Path
$resolvedKeyPath = (Resolve-Path -LiteralPath $PrivateKeyPath -ErrorAction Stop).Path
$windowsPowerShell = Join-Path $env:SystemRoot 'System32\WindowsPowerShell\v1.0\powershell.exe'

if ($SourceKey -ne 'default' -and -not $PSBoundParameters.ContainsKey('TaskName')) {
    $TaskName = 'CaiWuXiTong-UpstreamTunnel-{0}' -f $SourceKey
}

$privateConfigPath = if ([string]::IsNullOrWhiteSpace($ConfigPath)) {
    $configFileName = if ($SourceKey -eq 'default') {
        'upstream-tunnel.json'
    }
    else {
        'upstream-tunnel.{0}.json' -f $SourceKey
    }

    Join-Path (Join-Path $env:LOCALAPPDATA 'CaiWuXiTong') $configFileName
}
else {
    [System.IO.Path]::GetFullPath($ConfigPath)
}

if (-not (Test-Path -LiteralPath $windowsPowerShell -PathType Leaf)) {
    throw 'Windows PowerShell was not found; the tunnel task cannot be installed.'
}

$argumentList = @(
    '-NoProfile',
    '-ExecutionPolicy', 'Bypass',
    '-WindowStyle', 'Hidden',
    '-File', (ConvertTo-TaskArgument $tunnelScript),
    '-ConfigPath', (ConvertTo-TaskArgument $privateConfigPath)
)
$taskArguments = $argumentList -join ' '
$currentUser = [System.Security.Principal.WindowsIdentity]::GetCurrent().Name

if ($PSCmdlet.ShouldProcess($TaskName, 'register and start the read-only upstream tunnel task')) {
    $privateConfigDirectory = Split-Path -Parent $privateConfigPath
    $null = New-Item -ItemType Directory -Path $privateConfigDirectory -Force

    $privateSettings = [ordered] @{
        SourceKey = $SourceKey
        ServerHost = $ServerHost
        ServerPort = $ServerPort
        RemoteUser = $RemoteUser
        PrivateKeyPath = $resolvedKeyPath
        RemoteBindAddress = $RemoteBindAddress
        RemoteBindPort = $RemoteBindPort
        UpstreamHost = $UpstreamHost
        UpstreamPort = $UpstreamPort
        RetrySeconds = $RetrySeconds
    }
    $json = $privateSettings | ConvertTo-Json
    [System.IO.File]::WriteAllText($privateConfigPath, $json, (New-Object System.Text.UTF8Encoding($false)))

    $currentSid = [System.Security.Principal.WindowsIdentity]::GetCurrent().User.Value
    & icacls.exe $privateConfigPath '/inheritance:r' '/grant:r' "*$currentSid`:(F)" '*S-1-5-18:(F)' | Out-Null
    if ($LASTEXITCODE -ne 0) {
        throw 'The private tunnel configuration permissions could not be restricted.'
    }

    # The ScheduledTasks module can report false RPC errors on trimmed Windows
    # installations. The built-in COM API works in Windows PowerShell and pwsh.
    $taskService = New-Object -ComObject 'Schedule.Service'
    $rootFolder = $null
    $definition = $null
    $registeredTask = $null

    try {
        $taskService.Connect()
        $rootFolder = $taskService.GetFolder('\')
        $definition = $taskService.NewTask(0)

        $definition.RegistrationInfo.Description = 'CaiWuXiTong restricted read-only upstream tunnel with automatic reconnect.'
        $definition.Principal.UserId = $currentUser
        $definition.Principal.LogonType = 3 # TASK_LOGON_INTERACTIVE_TOKEN
        $definition.Principal.RunLevel = 0 # TASK_RUNLEVEL_LUA

        $definition.Settings.Enabled = $true
        $definition.Settings.StartWhenAvailable = $true
        $definition.Settings.DisallowStartIfOnBatteries = $false
        $definition.Settings.StopIfGoingOnBatteries = $false
        $definition.Settings.ExecutionTimeLimit = 'PT0S'
        $definition.Settings.MultipleInstances = 2 # TASK_INSTANCES_IGNORE_NEW
        $definition.Settings.RestartCount = 999
        $definition.Settings.RestartInterval = 'PT1M'

        $trigger = $definition.Triggers.Create(9) # TASK_TRIGGER_LOGON
        $trigger.UserId = $currentUser
        $trigger.Enabled = $true

        # Restart-on-failure does not cover tasks that Windows reports as
        # externally terminated (for example 0xC000013A).  Re-run the task
        # every five minutes as a watchdog; IgnoreNew keeps the healthy
        # long-running wrapper as the only instance.
        $watchdogTrigger = $definition.Triggers.Create(2) # TASK_TRIGGER_DAILY
        $watchdogTrigger.StartBoundary = (Get-Date).AddMinutes(1).ToString("yyyy-MM-dd'T'HH:mm:ss")
        $watchdogTrigger.DaysInterval = 1
        $watchdogTrigger.Repetition.Interval = 'PT5M'
        $watchdogTrigger.Repetition.Duration = 'P1D'
        $watchdogTrigger.Repetition.StopAtDurationEnd = $false
        $watchdogTrigger.Enabled = $true

        $action = $definition.Actions.Create(0) # TASK_ACTION_EXEC
        $action.Path = $windowsPowerShell
        $action.Arguments = $taskArguments
        $action.WorkingDirectory = $PSScriptRoot

        $existingTask = $null
        $previousWrapperPid = $null
        try {
            $existingTask = $rootFolder.GetTask($TaskName)
        }
        catch [System.Runtime.InteropServices.COMException] {
            $existingTask = $null
        }
        catch [System.IO.FileNotFoundException] {
            # PowerShell 7 can translate Task Scheduler's missing-task HRESULT
            # to FileNotFoundException instead of COMException.
            $existingTask = $null
        }

        if ($null -ne $existingTask -and $existingTask.State -eq 4) {
            $statePath = $privateConfigPath + '.pid'
            if (-not (Test-Path -LiteralPath $statePath -PathType Leaf)) {
                throw 'The existing tunnel task is running without a managed PID file; refusing an unsafe restart.'
            }

            $wrapperPid = [int] (Get-Content -LiteralPath $statePath -Raw)
            $previousWrapperPid = $wrapperPid
            $wrapperProcess = Get-Process -Id $wrapperPid -ErrorAction Stop
            if ($wrapperProcess.ProcessName -notin @('powershell', 'pwsh')) {
                throw 'The managed tunnel PID does not belong to PowerShell; refusing an unsafe restart.'
            }

            $sshChildren = @(
                Get-CimInstance -ClassName Win32_Process -Filter "ParentProcessId = $wrapperPid AND Name = 'ssh.exe'" -ErrorAction Stop
            )
            $existingTask.Stop(0)

            foreach ($sshChild in $sshChildren) {
                $sshProcess = Get-Process -Id ([int] $sshChild.ProcessId) -ErrorAction SilentlyContinue
                if ($null -ne $sshProcess -and $sshProcess.ProcessName -eq 'ssh') {
                    Stop-Process -Id $sshProcess.Id -Force
                }
            }

            for ($attempt = 0; $attempt -lt 20 -and $existingTask.State -eq 4; $attempt++) {
                Start-Sleep -Milliseconds 100
            }

            if ($existingTask.State -eq 4) {
                throw 'The existing tunnel task did not stop cleanly.'
            }

            $remainingWrapper = Get-Process -Id $wrapperPid -ErrorAction SilentlyContinue
            if ($null -ne $remainingWrapper) {
                if ($remainingWrapper.ProcessName -notin @('powershell', 'pwsh')) {
                    throw 'The stopped tunnel PID was reused by another process; refusing an unsafe cleanup.'
                }
                Stop-Process -Id $wrapperPid -Force
            }

            [System.IO.File]::Delete($statePath)
            Start-Sleep -Milliseconds 500
        }
        else {
            $inactiveStatePath = $privateConfigPath + '.pid'
            if (Test-Path -LiteralPath $inactiveStatePath -PathType Leaf) {
                $inactivePid = [int] (Get-Content -LiteralPath $inactiveStatePath -Raw)
                $inactiveProcess = Get-Process -Id $inactivePid -ErrorAction SilentlyContinue
                if ($null -ne $inactiveProcess -and $inactiveProcess.ProcessName -in @('powershell', 'pwsh')) {
                    throw 'A managed tunnel wrapper is running outside the scheduled task; refusing an unsafe replacement.'
                }
                [System.IO.File]::Delete($inactiveStatePath)
            }
        }

        if ($null -ne $existingTask) {
            [void] [System.Runtime.InteropServices.Marshal]::FinalReleaseComObject($existingTask)
        }

        $registeredTask = $rootFolder.RegisterTaskDefinition(
            $TaskName,
            $definition,
            6, # TASK_CREATE_OR_UPDATE
            $null,
            $null,
            3, # TASK_LOGON_INTERACTIVE_TOKEN
            $null
        )

        $taskStarted = $false
        for ($startAttempt = 0; $startAttempt -lt 3 -and -not $taskStarted; $startAttempt++) {
            $runningInstance = $registeredTask.Run($null)
            if ($null -ne $runningInstance -and [System.Runtime.InteropServices.Marshal]::IsComObject($runningInstance)) {
                [void] [System.Runtime.InteropServices.Marshal]::FinalReleaseComObject($runningInstance)
            }

            for ($stateAttempt = 0; $stateAttempt -lt 20; $stateAttempt++) {
                Start-Sleep -Milliseconds 100
                $taskProbe = $rootFolder.GetTask($TaskName)
                $taskState = [int] $taskProbe.State
                [void] [System.Runtime.InteropServices.Marshal]::FinalReleaseComObject($taskProbe)

                if ($taskState -eq 4) {
                    $taskStarted = $true
                    break
                }
            }

            if (-not $taskStarted) {
                Start-Sleep -Milliseconds 500
            }
        }

        if (-not $taskStarted) {
            throw 'The tunnel task was registered but did not enter the running state.'
        }

        $statePath = $privateConfigPath + '.pid'
        $managedPid = $null
        for ($stateAttempt = 0; $stateAttempt -lt 30; $stateAttempt++) {
            Start-Sleep -Milliseconds 100
            if (Test-Path -LiteralPath $statePath -PathType Leaf) {
                $candidatePid = [int] (Get-Content -LiteralPath $statePath -Raw)
                $candidateProcess = Get-Process -Id $candidatePid -ErrorAction SilentlyContinue
                if ($null -ne $candidateProcess -and $candidateProcess.ProcessName -in @('powershell', 'pwsh')) {
                    if ($null -eq $previousWrapperPid -or $candidatePid -ne $previousWrapperPid) {
                        $managedPid = $candidatePid
                        break
                    }
                }
            }
        }

        if ($null -eq $managedPid) {
            throw 'The tunnel task started without creating a valid managed PID file.'
        }

        Write-Output ('Scheduled task {0} was registered and started.' -f $TaskName)
    }
    finally {
        foreach ($comObject in @($registeredTask, $definition, $rootFolder, $taskService)) {
            if ($null -ne $comObject -and [System.Runtime.InteropServices.Marshal]::IsComObject($comObject)) {
                [void] [System.Runtime.InteropServices.Marshal]::FinalReleaseComObject($comObject)
            }
        }
    }
}