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

$SheepItJarUrl = '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 "'", "''"
    $userLine = if ([string]::IsNullOrWhiteSpace($Username)) {
        ''
    } else {
        "    Write-Host 'Authenticating as ${safeUser} via render key.' -ForegroundColor Yellow`n"
    }

@"
`$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'
    Write-Host "Downloading SheepIt client..." -ForegroundColor Cyan
    Invoke-WebRequest -Uri '$SheepItJarUrl' -OutFile `$jarPath -UseBasicParsing
    Write-Host "Download complete." -ForegroundColor Green

    Write-Host "Starting SheepIt client..." -ForegroundColor Cyan
${userLine}    Set-Location `$sheepDir
    & java -jar `$jarPath --ui text --stdout --verbose -gpu optix 0 --renderkey '${safeKey}'
}
catch {
    Write-Host "Error: `$($_.Exception.Message)" -ForegroundColor Red
}
"@
}

function New-SheepItSessionScript {
    param(
        [object]$Worker,
        [string]$EncodedCommand
    )

@"
Write-Host 'Connecting to $($Worker.Name)...' -ForegroundColor Cyan
ssh -o ServerAliveInterval=60 -o ServerAliveCountMax=30 $($Worker.SSHArgs) `"powershell -NoLogo -NoProfile -ExecutionPolicy Bypass -EncodedCommand $EncodedCommand`"
Write-Host "`nSSH session ended." -ForegroundColor Yellow
Read-Host "Press Enter to close"
"@
}

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
    $encoded = [Convert]::ToBase64String([System.Text.Encoding]::Unicode.GetBytes($remoteCommand))
    $sessionScript = New-SheepItSessionScript -Worker $Worker -EncodedCommand $encoded
    $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