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 }
)

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
    }
}
catch {
    Write-Host ('Error: {0}' -f `$_.Exception.Message) -ForegroundColor Red
    Write-Host ('Stack trace: {0}' -f `$_.ScriptStackTrace) -ForegroundColor DarkRed
}
"@
}

function New-SheepItSessionScript {
    param(
        [object]$Worker,
        [string]$RemoteCommand
    )
    
    $rawArgs = @($Worker.SSHArgs -split '\s+' | Where-Object { $_ -and $_.Trim().Length -gt 0 })
    $sshArgsJson = if ($rawArgs) { ($rawArgs | ConvertTo-Json -Compress) } else { '[]' }
    $targetHost = ($rawArgs | Where-Object { $_ -notmatch '^-{1,2}' } | Select-Object -Last 1)
    $hostLiteral = if ($targetHost) { "'" + ($targetHost -replace "'", "''") + "'" } else { '$null' }
    $portValue = $null
    for ($i = 0; $i -lt $rawArgs.Count; $i++) {
        if ($rawArgs[$i] -eq '-p' -and $i -lt ($rawArgs.Count - 1)) {
            $portValue = $rawArgs[$i + 1]
            break
        }
    }
    $portLiteral = if ($portValue) { "'" + ($portValue -replace "'", "''") + "'" } else { '$null' }
    $remoteScriptBase64 = [Convert]::ToBase64String([System.Text.Encoding]::Unicode.GetBytes($RemoteCommand))
    
    $scriptContent = @"
Write-Host 'Connecting to $($Worker.Name)...' -ForegroundColor Cyan
`$sshArgs = ConvertFrom-Json '$sshArgsJson'
`$targetHost = $hostLiteral
`$port = $portLiteral
`$scriptBase64 = '$remoteScriptBase64'

if (-not `$targetHost) {
    Write-Host "Unable to determine SSH host for $($Worker.Name)." -ForegroundColor Red
    return
}

`$localTemp = [System.IO.Path]::GetTempFileName() + '.ps1'
`$remoteScript = [System.Text.Encoding]::Unicode.GetString([Convert]::FromBase64String(`$scriptBase64))
Set-Content -Path `$localTemp -Value `$remoteScript -Encoding UTF8

`$remoteFile = "sheepit-$([System.Guid]::NewGuid().ToString()).ps1"

`$scpArgs = @('-o','ServerAliveInterval=60','-o','ServerAliveCountMax=30')
if (`$port) {
    `$scpArgs += '-P'
    `$scpArgs += `$port
}
`$scpArgs += `$localTemp
`$scpArgs += ("${targetHost}:`$remoteFile")

Write-Host "Transferring SheepIt script to `$targetHost..." -ForegroundColor Gray
& scp @scpArgs
if (`$LASTEXITCODE -ne 0) {
    Write-Host "Failed to copy script to `$targetHost." -ForegroundColor Red
    Remove-Item `$localTemp -ErrorAction SilentlyContinue
    return
}

`$sshBase = @('-o','ServerAliveInterval=60','-o','ServerAliveCountMax=30') + `$sshArgs
`$execArgs = `$sshBase + @('powershell','-NoLogo','-NoProfile','-ExecutionPolicy','Bypass','-File',"`$remoteFile")

Write-Host "Starting SheepIt client on `$targetHost..." -ForegroundColor Cyan
& ssh @execArgs

Write-Host "Cleaning up remote script..." -ForegroundColor Gray
`$cleanupArgs = `$sshBase + @('powershell','-NoLogo','-NoProfile','-ExecutionPolicy','Bypass','-Command',"`$ErrorActionPreference='SilentlyContinue'; Remove-Item -LiteralPath ``"$remoteFile``" -ErrorAction SilentlyContinue")
& ssh @cleanupArgs | Out-Null

Remove-Item `$localTemp -ErrorAction SilentlyContinue
Write-Host "`nSSH session ended." -ForegroundColor Yellow
Read-Host "Press Enter to close"
"@
    
    return $scriptContent
}

function Start-SheepItTab {
    param(
        [string]$Title,
        [string]$Content
    )

    $tempScript = [System.IO.Path]::GetTempFileName() + '.ps1'
    Set-Content -Path $tempScript -Value $Content -Encoding UTF8

    if (Get-Command wt.exe -ErrorAction SilentlyContinue) {
        Start-Process wt.exe -ArgumentList "-w 0 new-tab --title `"$Title`" powershell -NoLogo -NoProfile -ExecutionPolicy Bypass -File `"$tempScript`""
    } else {
        Start-Process powershell -ArgumentList "-NoLogo -NoProfile -ExecutionPolicy Bypass -File `"$tempScript`""
    }
}

function Start-SheepItWorker {
    param([object]$Worker)

    if (-not $Worker.Enabled) {
        Write-Host "$($Worker.Name) is not enabled for SheepIt." -ForegroundColor Yellow
        return
    }

    Initialize-SheepItCredentials

    $remoteCommand = Get-RemoteSheepItCommand -RenderKey $script:SheepItRenderKey -Username $script:SheepItUsername
    $sessionScript = New-SheepItSessionScript -Worker $Worker -RemoteCommand $remoteCommand
    $title = "$($Worker.Name) - SheepIt"

    Start-SheepItTab -Title $title -Content $sessionScript
    Write-Host "Opened SheepIt session for $($Worker.Name) in a new terminal tab." -ForegroundColor Green
}

function Start-AllSheepIt {
    Initialize-SheepItCredentials

    $targets = $workers | Where-Object { $_.Enabled }

    if (-not $targets) {
        Write-Host "No systems are ready for SheepIt." -ForegroundColor Yellow
        return
    }

    foreach ($worker in $targets) {
        Start-SheepItWorker -Worker $worker
        Start-Sleep -Milliseconds 200
    }
}

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 SheepIt on a single system" -ForegroundColor Yellow
    Write-Host "2. Launch SheepIt on all ready systems" -ForegroundColor Yellow
    Write-Host "3. Exit" -ForegroundColor Yellow

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

    switch ($choice) {
        '1' { Select-SheepItWorker }
        '2' {
            Start-AllSheepIt
            Write-Host
            Read-Host "Press Enter to return to the main menu" | Out-Null
        }
        '3' { break }
        default {
            Write-Host "Invalid selection." -ForegroundColor Red
            Start-Sleep -Seconds 1
        }
    }
}

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