function Show-Header {
    Clear-Host
    Write-Host "====================================" -ForegroundColor Cyan
    Write-Host "      UNIFIED SHEEPIT LAUNCHER" -ForegroundColor Cyan
    Write-Host "====================================" -ForegroundColor Cyan
    Write-Host
}

$SheepItJarUrls = @(
    'https://www.sheepit-renderfarm.com/media/applet/client-latest.php',
    'https://www.sheepit-renderfarm.com/media/applet/client-latest.jar'
)
$script:SheepItUsername = $null
$script:SheepItRenderKey = $null

function Initialize-SheepItCredentials {
    if (-not $script:SheepItUsername) {
        $script:SheepItUsername = "RaincloudTheDragon"
    }

    if (-not $script:SheepItRenderKey) {
        $script:SheepItRenderKey = "IfCOWBHFQpceG0601DmyrwOOJOAp2UJAQ0O0X0jF"
    }
}

$workers = @(
    @{ ID = 1; Name = "i9kf"; SSHArgs = "-t i9kf"; Enabled = $true },
    @{ ID = 2; Name = "blender-boss"; SSHArgs = "-t blender-boss"; Enabled = $true },
    @{ ID = 3; Name = "max"; SSHArgs = "-t max"; Enabled = $true },
    @{ ID = 4; Name = "masterbox"; SSHArgs = "-t masterbox"; Enabled = $true },
    @{ ID = 5; Name = "echo"; SSHArgs = "-t echo"; Enabled = $true },
    @{ ID = 6; Name = "i9-13ks"; SSHArgs = "-t -p 22146 i9-13ks"; Enabled = $true }
)

$global:UnifiedWorkerDataRoot = 'C:\ProgramData\UnifiedWorkers'
$script:ControllerScriptBase64 = $null
$script:AttachHelperScriptBase64 = $null

function Get-FileBase64 {
    param([string]$Path)
    [Convert]::ToBase64String([IO.File]::ReadAllBytes($Path))
}

function Get-ControllerScriptBase64 {
    if (-not $script:ControllerScriptBase64) {
        $controllerPath = Join-Path $PSScriptRoot 'remote_worker_controller.ps1'
        $script:ControllerScriptBase64 = Get-FileBase64 -Path $controllerPath
    }
    return $script:ControllerScriptBase64
}

function Get-AttachHelperScriptBase64 {
    if (-not $script:AttachHelperScriptBase64) {
        $helperPath = Join-Path $PSScriptRoot 'remote_worker_attach.ps1'
        $script:AttachHelperScriptBase64 = Get-FileBase64 -Path $helperPath
    }
    return $script:AttachHelperScriptBase64
}

function ConvertTo-Base64Unicode {
    param([string]$Content)
    [Convert]::ToBase64String([Text.Encoding]::Unicode.GetBytes($Content))
}

function Get-WorkerSshArgs {
    param([string]$RawArgs)
    if ([string]::IsNullOrWhiteSpace($RawArgs)) {
        return @()
    }
    return $RawArgs -split '\s+' | Where-Object { $_.Trim().Length -gt 0 }
}

function Invoke-RemotePowerShell {
    param(
        [object]$Worker,
        [string]$Script
    )

    $bytes = [Text.Encoding]::Unicode.GetBytes($Script)
    $encoded = [Convert]::ToBase64String($bytes)
    $sshArgs = @('-o','ServerAliveInterval=60','-o','ServerAliveCountMax=30')
    $sshArgs += (Get-WorkerSshArgs -RawArgs $Worker.SSHArgs)
    $sshArgs += "powershell -NoLogo -NoProfile -ExecutionPolicy Bypass -EncodedCommand $encoded"

    & ssh @sshArgs
    return $LASTEXITCODE
}

function Ensure-ControllerDeployed {
    param([object]$Worker)
    $controllerBase64 = Get-ControllerScriptBase64
    $script = @'
`$dataRoot = '{0}'
New-Item -ItemType Directory -Path `$dataRoot -Force | Out-Null
`$controllerPath = Join-Path `$dataRoot 'controller.ps1'
[IO.File]::WriteAllBytes(`$controllerPath, [Convert]::FromBase64String('{1}'))
'@ -f $global:UnifiedWorkerDataRoot, $controllerBase64
    Invoke-RemotePowerShell -Worker $Worker -Script $script | Out-Null
}

function Ensure-AttachHelperDeployed {
    param([object]$Worker)
    $helperBase64 = Get-AttachHelperScriptBase64
    $script = @'
`$dataRoot = '{0}'
New-Item -ItemType Directory -Path `$dataRoot -Force | Out-Null
`$attachPath = Join-Path `$dataRoot 'attach-helper.ps1'
[IO.File]::WriteAllBytes(`$attachPath, [Convert]::FromBase64String('{1}'))
'@ -f $global:UnifiedWorkerDataRoot, $helperBase64
    Invoke-RemotePowerShell -Worker $Worker -Script $script | Out-Null
}

function Get-RemoteWorkerPaths {
    param([string]$WorkerType, [string]$WorkerName)
    $base = Join-Path $global:UnifiedWorkerDataRoot $WorkerType
    $instance = Join-Path $base $WorkerName
    [pscustomobject]@{
        InstanceRoot = $instance
        MetadataPath = Join-Path $instance 'state\worker-info.json'
        CommandPath  = Join-Path $instance 'state\commands.txt'
        LogPath      = Join-Path $instance 'logs\worker.log'
    }
}

function Get-EnsureWorkerScript {
    param(
        [string]$WorkerName,
        [string]$WorkerType,
        [string]$PayloadBase64
    )

    $paths = Get-RemoteWorkerPaths -WorkerType $WorkerType -WorkerName $WorkerName
    $controllerPath = Join-Path $global:UnifiedWorkerDataRoot 'controller.ps1'

    @'
`$controllerPath = '{0}'
`$metaPath = '{1}'
`$workerName = '{2}'
`$workerType = '{3}'
`$payloadBase64 = '{4}'

if (-not (Test-Path `$controllerPath)) {
    throw "Controller missing at `$controllerPath"
}

`$shouldStart = `$true
if (Test-Path `$metaPath) {
    try {
        `$meta = Get-Content `$metaPath -Raw | ConvertFrom-Json
        if (`$meta.Status -eq 'running' -and `$meta.WorkerPid) {
            if (Get-Process -Id `$meta.WorkerPid -ErrorAction SilentlyContinue) {
                Write-Host "Worker `$workerName already running (PID `$($meta.WorkerPid))."
                `$shouldStart = `$false
            }
        }
    } catch {
        Write-Host "Failed to read metadata. Controller will restart worker." -ForegroundColor Yellow
    }
}

if (`$shouldStart) {
    `$psExe = (Get-Command pwsh -ErrorAction SilentlyContinue)?.Source
    if (-not `$psExe) {
        `$psExe = (Get-Command powershell -ErrorAction Stop).Source
    }

    `$controllerArgs = @(
        '-NoLogo','-NoProfile','-ExecutionPolicy','Bypass',
        '-File',"`$controllerPath",
        '-WorkerName',"`$workerName",
        '-WorkerType',"`$workerType",
        '-PayloadBase64',"`$payloadBase64"
    )

    Start-Process -FilePath `$psExe -ArgumentList `$controllerArgs -WindowStyle Hidden | Out-Null
    Write-Host "Worker `$workerName started under controller." -ForegroundColor Green
}
'@ -f $controllerPath, $paths.MetadataPath, $WorkerName, $WorkerType, $PayloadBase64
}

function Ensure-PersistentWorker {
    param(
        [object]$Worker,
        [string]$WorkerType,
        [string]$PayloadScript
    )

    Ensure-ControllerDeployed -Worker $Worker
    $payloadBase64 = ConvertTo-Base64Unicode -Content $PayloadScript
    $ensureScript = Get-EnsureWorkerScript -WorkerName $Worker.Name -WorkerType $WorkerType -PayloadBase64 $payloadBase64
    Invoke-RemotePowerShell -Worker $Worker -Script $ensureScript | Out-Null
}

function Invoke-WorkerAttach {
    param(
        [object]$Worker,
        [string]$WorkerType,
        [switch]$CommandOnly,
        [string]$Command
    )

    Ensure-AttachHelperDeployed -Worker $Worker
    $paramsBlock = @("-WorkerName","$($Worker.Name)","-WorkerType","$WorkerType")
    if ($CommandOnly) {
        $paramsBlock += "-CommandOnly"
    }
    if ($Command) {
        $paramsBlock += "-Command"
        $paramsBlock += $Command
    }

    $paramsLiteral = ($paramsBlock | ForEach-Object { "'$_'" }) -join ', '
    $script = @'
`$helperPath = Join-Path '{0}' 'attach-helper.ps1'
if (-not (Test-Path `$helperPath)) {
    throw "Attach helper missing at `$helperPath"
}
`$arguments = @({1})
& `$helperPath @arguments
'@ -f $global:UnifiedWorkerDataRoot, $paramsLiteral

    Invoke-RemotePowerShell -Worker $Worker -Script $script | Out-Null
}

function Get-RemoteSheepItCommand {
    param(
        [string]$RenderKey,
        [string]$Username
    )

    $safeKey = $RenderKey -replace "'", "''"
    $safeUser = $Username -replace "'", "''"

    $urlLiteral = '@(' + (($SheepItJarUrls | ForEach-Object { "'$_'" }) -join ', ') + ')'

@"
`$ErrorActionPreference = 'Stop'

try {
    `$appData = [Environment]::GetFolderPath('ApplicationData')
    `$sheepDir = Join-Path `$appData 'sheepit'
    if (-not (Test-Path `$sheepDir)) {
        New-Item -Path `$sheepDir -ItemType Directory -Force | Out-Null
    }

    `$jarPath = Join-Path `$sheepDir 'sheepit-client.jar'
    `$urls = $urlLiteral
    `$headers = @{ 'User-Agent' = 'Mozilla/5.0' }

    if (Test-Path `$jarPath) {
        Write-Host "SheepIt client already present at `$jarPath. Skipping download." -ForegroundColor Green
    }
    else {
        `$downloaded = `$false

        foreach (`$url in `$urls) {
            Write-Host "Downloading SheepIt client from `$url..." -ForegroundColor Cyan
            try {
                Invoke-WebRequest -Uri `$url -OutFile `$jarPath -UseBasicParsing -Headers `$headers
                `$downloaded = `$true
                Write-Host "Download complete." -ForegroundColor Green
                break
            }
            catch {
                Write-Host ("Download failed from {0}: {1}" -f `$url, `$_.Exception.Message) -ForegroundColor Yellow
            }
        }

        if (-not `$downloaded) {
            throw 'Unable to download SheepIt client from any known URL.'
        }
    }

    Write-Host "Starting SheepIt client..." -ForegroundColor Cyan
    
    # Check and fix problematic environment variables that can cause boot class path errors
    `$envVarsFixed = `$false
    
    # Check JAVA_HOME - invalid values like '\' or empty can cause issues
    if (`$env:JAVA_HOME) {
        if (`$env:JAVA_HOME -eq '\' -or `$env:JAVA_HOME.Trim() -eq '' -or -not (Test-Path `$env:JAVA_HOME)) {
            Write-Host "Warning: Invalid JAVA_HOME detected ('`$env:JAVA_HOME'). Temporarily unsetting..." -ForegroundColor Yellow
            Remove-Item Env:\JAVA_HOME -ErrorAction SilentlyContinue
            `$envVarsFixed = `$true
        }
    }
    
    # Check JAVA_TOOL_OPTIONS - invalid values can cause boot class path errors
    if (`$env:JAVA_TOOL_OPTIONS) {
        if (`$env:JAVA_TOOL_OPTIONS -eq '\' -or `$env:JAVA_TOOL_OPTIONS.Trim() -eq '') {
            Write-Host "Warning: Invalid JAVA_TOOL_OPTIONS detected ('`$env:JAVA_TOOL_OPTIONS'). Temporarily unsetting..." -ForegroundColor Yellow
            Remove-Item Env:\JAVA_TOOL_OPTIONS -ErrorAction SilentlyContinue
            `$envVarsFixed = `$true
        }
    }
    
    if (`$envVarsFixed) {
        Write-Host "Environment variables fixed. Proceeding with Java launch..." -ForegroundColor Green
    }
    
    # Check Java version
    try {
        `$javaOutput = java -version 2>&1
        `$javaVersion = `$javaOutput | Select-Object -First 1
        Write-Host "Java version: `$javaVersion" -ForegroundColor Gray
    }
    catch {
        Write-Host "Warning: Could not determine Java version" -ForegroundColor Yellow
    }
    
    Set-Location `$sheepDir
    
    # Use -XX:+IgnoreUnrecognizedVMOptions to handle any incompatible JVM args
    # This flag helps with Java 9+ compatibility issues
    `$javaArgs = @('-XX:+IgnoreUnrecognizedVMOptions', '-jar', `$jarPath, 
                   '-ui', 'text', '--log-stdout', '--verbose', 
                   '-gpu', 'OPTIX_0', '-login', '${safeUser}', '-password', '${safeKey}')
    
    try {
        & java @javaArgs
    }
    catch {
        Write-Host ('Java execution error: {0}' -f `$_.Exception.Message) -ForegroundColor Red
        Write-Host "If the error persists, try reinstalling Java (Temurin 21 recommended)." -ForegroundColor Yellow
        throw
    }
}

function Ensure-SheepItWorkerController {
    param([object]$Worker)
    Initialize-SheepItCredentials
    $payloadScript = Get-RemoteSheepItCommand -RenderKey $script:SheepItRenderKey -Username $script:SheepItUsername
    Ensure-PersistentWorker -Worker $Worker -WorkerType 'sheepit' -PayloadScript $payloadScript
}

function Start-SheepItWorker {
    param([object]$Worker)
    Ensure-SheepItWorkerController -Worker $Worker
    Invoke-WorkerAttach -Worker $Worker -WorkerType 'sheepit'
}

function Start-AllSheepIt {
    Initialize-SheepItCredentials
    foreach ($worker in $workers | Where-Object { $_.Enabled }) {
        Ensure-SheepItWorkerController -Worker $worker
    }
    Write-Host "All enabled SheepIt workers ensured running under controllers." -ForegroundColor Green
    Write-Host "Use the attach option to monitor any worker." -ForegroundColor Cyan
    Read-Host "Press Enter to continue" | Out-Null
}

function Send-SheepItCommandAll {
    param([string]$CommandText)

    foreach ($worker in $workers | Where-Object { $_.Enabled }) {
        Write-Host "[$($worker.Name)] Sending '$CommandText'..." -ForegroundColor Gray
        Invoke-WorkerAttach -Worker $worker -WorkerType 'sheepit' -CommandOnly -Command $CommandText
    }
    Write-Host "Command '$CommandText' dispatched to all enabled workers." -ForegroundColor Green
    Read-Host "Press Enter to continue" | Out-Null
}
catch {
    Write-Host ('Error: {0}' -f `$_.Exception.Message) -ForegroundColor Red
    Write-Host ('Stack trace: {0}' -f `$_.ScriptStackTrace) -ForegroundColor DarkRed
}
"@
}



function Select-SheepItWorker {
    while ($true) {
        Show-Header
        Write-Host "Select a system:" -ForegroundColor Magenta

        foreach ($worker in $workers) {
            $status = if ($worker.Enabled) { "ready" } else { "disabled" }
            $color = if ($worker.Enabled) { "Green" } else { "DarkYellow" }
            Write-Host ("{0}. {1} ({2})" -f $worker.ID, $worker.Name, $status) -ForegroundColor $color
        }

        Write-Host "B. Back" -ForegroundColor Yellow
        $selection = Read-Host "Choose system"

        if ($selection -match '^[Bb]$') {
            return
        }

        if ($selection -match '^\d+$') {
            $id = [int]$selection
            $worker = $workers | Where-Object { $_.ID -eq $id }

            if ($worker) {
                Start-SheepItWorker -Worker $worker
                Write-Host
                Read-Host "Press Enter to return to the main menu" | Out-Null
                return
            }

            Write-Host "Unknown selection." -ForegroundColor Red
            Start-Sleep -Seconds 1
        }
        else {
            Write-Host "Invalid input." -ForegroundColor Red
            Start-Sleep -Seconds 1
        }
    }
}

Initialize-SheepItCredentials

while ($true) {
    Show-Header
    Write-Host "Main Menu:" -ForegroundColor Magenta
    Write-Host "1. Launch/Attach SheepIt on a single system" -ForegroundColor Yellow
    Write-Host "2. Ensure all ready systems are running" -ForegroundColor Yellow
    Write-Host "3. Pause all workers" -ForegroundColor Yellow
    Write-Host "4. Resume all workers" -ForegroundColor Yellow
    Write-Host "5. Quit all workers" -ForegroundColor Yellow
    Write-Host "6. Exit" -ForegroundColor Yellow

    $choice = Read-Host "Select option (1-3)"

    switch ($choice) {
        '1' { Select-SheepItWorker }
        '2' { Start-AllSheepIt }
        '3' { Send-SheepItCommandAll -CommandText 'pause' }
        '4' { Send-SheepItCommandAll -CommandText 'resume' }
        '5' { Send-SheepItCommandAll -CommandText 'quit' }
        '6' { break }
        default {
            Write-Host "Invalid selection." -ForegroundColor Red
            Start-Sleep -Seconds 1
        }
    }
}

Write-Host "`nExiting SheepIt launcher." -ForegroundColor Cyan