# Master Unified Flamenco Launcher Script
Write-Host "==========================================" -ForegroundColor Cyan
Write-Host "    UNIFIED FLAMENCO WORKER LAUNCHER" -ForegroundColor Cyan
Write-Host "==========================================" -ForegroundColor Cyan
Write-Host

# Define worker-specific configuration
$workers = @(
    @{
        ID = 1
        Name = "i9kf"
        SSHHost = "i9kf"
        SSHPort = 22
        SSHArgs = "-t i9kf"
    },
    @{
        ID = 2
        Name = "blender-boss"
        SSHHost = "blender-boss"
        SSHPort = 22
        SSHArgs = "-t blender-boss"
    },
    @{
        ID = 3
        Name = "max"
        SSHHost = "max"
        SSHPort = 22
        SSHArgs = "-t max"
    },
    @{
        ID = 4
        Name = "masterbox"
        SSHHost = "masterbox"
        SSHPort = 22
        SSHArgs = "-t masterbox"
    },
    @{
        ID = 5
        Name = "echo"
        SSHHost = "echo"
        SSHPort = 22
        SSHArgs = "-t echo"
    },
    @{
        ID = 6
        Name = "i9-13ks"
        SSHHost = "i9-13ks"
        SSHPort = 22146
        SSHArgs = "-t -p 22146 i9-13ks"
    }
)

$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(
        [Parameter(Mandatory = $true)][object]$Worker,
        [Parameter(Mandatory = $true)][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}'))
Write-Host "Controller ready at `$controllerPath"
'@ -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
    return [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'

    return @'
`$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 "Unable to parse metadata for `$workerName. Starting new controller." -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
    )

    Write-Host "[$($Worker.Name)] Ensuring worker is running under controller..." -ForegroundColor Cyan
    Ensure-ControllerDeployed -Worker $Worker
    $payloadBase64 = ConvertTo-Base64Unicode -Content $PayloadScript
    $ensureScript = Get-EnsureWorkerScript -WorkerName $Worker.Name -WorkerType $WorkerType -PayloadBase64 $payloadBase64
    $result = Invoke-RemotePowerShell -Worker $Worker -Script $ensureScript
    if ($result -ne 0) {
        Write-Host "[$($Worker.Name)] Remote ensure command exited with code $result." -ForegroundColor Yellow
    }
}

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
}

# FUNCTIONS

# This function generates the standard PowerShell remote command
function Get-RemoteStandardWorkerCommand {
    @'
Write-Host "Setting up network connections..." -ForegroundColor Cyan

# Define arrays of drives and network paths
$drives = @('A:', 'F:', 'N:', 'P:')
$networkPaths = @(
    '\\NEXUS\amazon',
    '\\NEXUS\flamenco',
    '\\NEXUS\proj',
    '\\NAS\amazon'
)

# Disconnect all existing connections
Write-Host "Disconnecting existing network connections..." -ForegroundColor Yellow
foreach ($path in $networkPaths) { net use $path /delete /y 2>$null }
foreach ($drive in $drives) { net use $drive /delete /y 2>$null }
Write-Host "All network connections cleared." -ForegroundColor Green

# Check if any workers are running
$workerProcesses = Get-Process -Name "flamenco-worker" -ErrorAction SilentlyContinue
if ($workerProcesses) {
    Write-Host "Found $(($workerProcesses | Measure-Object).Count) running Flamenco workers." -ForegroundColor Yellow
    Write-Host "Running workers will NOT be stopped." -ForegroundColor Yellow
} else {
    Write-Host "No running Flamenco workers found." -ForegroundColor Green
}

# Connect to network shares
Write-Host "Establishing network connections..." -ForegroundColor Cyan

# Connect to NEXUS with password automatically supplied
net use \\NEXUS\amazon /user:Nathan HeadsTalk1ng! /persistent:yes
if ($LASTEXITCODE -eq 0) {
    # Map all NEXUS drives
    net use A: \\NEXUS\amazon /persistent:yes
    net use F: \\NEXUS\flamenco /persistent:yes
    net use P: \\NEXUS\proj /persistent:yes
} else {
    Write-Host "Failed to connect to NEXUS" -ForegroundColor Red
    exit 1  # Exit with error code to trigger restart
}

# Connect to NAS with password automatically supplied
net use N: \\NAS\amazon /user:Nathan HeadsTalk1ng! /persistent:yes
if ($LASTEXITCODE -ne 0) {
    Write-Host "Failed to connect to NAS" -ForegroundColor Red
    exit 1  # Exit with error code to trigger restart
}

# Verify connections
Write-Host "Current network connections:" -ForegroundColor Cyan
net use

# Start worker
Write-Host "Starting Flamenco worker..." -ForegroundColor Cyan
Set-Location 'F:\software\Flamenco 3.7'
if (Test-Path 'flamenco-worker.exe') {
    Write-Host "Running flamenco-worker.exe..." -ForegroundColor Green
    # Run the worker and capture its exit code
    $workerProcess = Start-Process -FilePath '.\flamenco-worker.exe' -NoNewWindow -PassThru -Wait
    $exitCode = $workerProcess.ExitCode
    Write-Host "Flamenco worker process has terminated with exit code: $exitCode" -ForegroundColor Yellow
    exit $exitCode  # Exit with the worker's exit code to trigger restart if needed
} else {
    Write-Host "Error: flamenco-worker.exe not found in F:\software\Flamenco 3.7" -ForegroundColor Red
    exit 1  # Exit with error code to trigger restart
}
'@
}

# This function generates the CMD PowerShell remote command
function Get-RemoteCmdWorkerCommand {
    @'
Write-Host "Setting up network connections..." -ForegroundColor Cyan

# Define arrays of drives and network paths
$drives = @('A:', 'F:', 'N:', 'P:')
$networkPaths = @(
    '\\NEXUS\amazon',
    '\\NEXUS\flamenco',
    '\\NEXUS\proj',
    '\\NAS\amazon'
)

# Disconnect all existing connections
Write-Host "Disconnecting existing network connections..." -ForegroundColor Yellow
foreach ($path in $networkPaths) { net use $path /delete /y 2>$null }
foreach ($drive in $drives) { net use $drive /delete /y 2>$null }
Write-Host "All network connections cleared." -ForegroundColor Green

# Connect to network shares
Write-Host "Establishing network connections..." -ForegroundColor Cyan

# Connect to NEXUS with password automatically supplied
net use \\NEXUS\amazon /user:Nathan HeadsTalk1ng! /persistent:yes
if ($LASTEXITCODE -eq 0) {
    # Map all NEXUS drives
    net use A: \\NEXUS\amazon /persistent:yes
    net use F: \\NEXUS\flamenco /persistent:yes
    net use P: \\NEXUS\proj /persistent:yes
} else {
    Write-Host "Failed to connect to NEXUS" -ForegroundColor Red
}

# Connect to NAS with password automatically supplied
net use N: \\NAS\amazon /user:Nathan HeadsTalk1ng! /persistent:yes

# Verify connections
Write-Host "Current network connections:" -ForegroundColor Cyan
net use

# Start worker via CMD - hardcoded paths 
Write-Host "Running command file..." -ForegroundColor Cyan
$defaultCmdPath = "F:\software\Flamenco 3.7\run-flamenco-worker.cmd"

if (Test-Path $defaultCmdPath) {
    Set-Location "F:\software\Flamenco 3.7"
    Write-Host "Starting worker..." -ForegroundColor Green
    # Use hardcoded path to avoid variable expansion issues
    cmd.exe /c "F:\software\Flamenco 3.7\run-flamenco-worker.cmd"
    Write-Host "Worker process has terminated." -ForegroundColor Yellow
} else {
    Write-Host "Command file not found at default location." -ForegroundColor Red
    $customPath = Read-Host "Enter path to .cmd file"
    
    if (Test-Path $customPath) {
        $customDir = Split-Path -Parent $customPath
        Set-Location $customDir
        Write-Host "Starting worker from custom path..." -ForegroundColor Green
        # For custom path, we need to use the variable but in a different way
        Invoke-Expression "cmd.exe /c `"$customPath`""
        Write-Host "Worker process has terminated." -ForegroundColor Yellow
    } else {
        Write-Host "Custom path not found." -ForegroundColor Red
    }
}
'@
}

# This function generates a simplified CMD worker command specifically for Launch All functionality
function Get-RemoteSimplifiedCmdWorkerCommand {
    @'
Write-Host "Setting up network connections..." -ForegroundColor Cyan

# Define arrays of drives and network paths
$drives = @('A:', 'F:', 'N:', 'P:')
$networkPaths = @(
    '\\NEXUS\amazon',
    '\\NEXUS\flamenco',
    '\\NEXUS\proj',
    '\\NAS\amazon'
)

# Disconnect all existing connections
Write-Host "Disconnecting existing network connections..." -ForegroundColor Yellow
foreach ($path in $networkPaths) { net use $path /delete /y 2>$null }
foreach ($drive in $drives) { net use $drive /delete /y 2>$null }
Write-Host "All network connections cleared." -ForegroundColor Green

# Connect to network shares
Write-Host "Establishing network connections..." -ForegroundColor Cyan

# Connect to NEXUS with password automatically supplied
net use \\NEXUS\amazon /user:Nathan HeadsTalk1ng! /persistent:yes
if ($LASTEXITCODE -eq 0) {
    # Map all NEXUS drives
    net use A: \\NEXUS\amazon /persistent:yes
    net use F: \\NEXUS\flamenco /persistent:yes
    net use P: \\NEXUS\proj /persistent:yes
} else {
    Write-Host "Failed to connect to NEXUS" -ForegroundColor Red
}

# Connect to NAS with password automatically supplied
net use N: \\NAS\amazon /user:Nathan HeadsTalk1ng! /persistent:yes

# Verify connections
Write-Host "Current network connections:" -ForegroundColor Cyan
net use

# Simple direct command execution with automatic "2" input
Write-Host "Running Flamenco worker..." -ForegroundColor Cyan
Set-Location -Path "F:\software\Flamenco 3.7"
if (Test-Path -Path "run-flamenco-worker.cmd") {
    # Create a temporary file to store the "2" input
    $tempInputFile = [System.IO.Path]::GetTempFileName()
    Set-Content -Path $tempInputFile -Value "2"
    # Run the command with input redirected from our temp file
    cmd.exe /c "run-flamenco-worker.cmd < $tempInputFile"
    # Clean up the temp file
    Remove-Item -Path $tempInputFile -Force
    Write-Host "Worker process has terminated." -ForegroundColor Yellow
} else {
    Write-Host "Worker command file not found." -ForegroundColor Red
}
'@
}

# This function launches the standard worker
function Start-StandardWorker {
    param (
        [Parameter(Mandatory = $true)]
        [object]$Worker
    )

    $payloadScript = Get-RemoteStandardWorkerCommand
    Ensure-PersistentWorker -Worker $Worker -WorkerType 'flamenco' -PayloadScript $payloadScript
    Invoke-WorkerAttach -Worker $Worker -WorkerType 'flamenco'
}

# This function launches the CMD worker
function Start-CmdWorker {
    param (
        [Parameter(Mandatory = $true)]
        [object]$Worker
    )

    $payloadScript = Get-RemoteCmdWorkerCommand
    Ensure-PersistentWorker -Worker $Worker -WorkerType 'flamenco' -PayloadScript $payloadScript
    Invoke-WorkerAttach -Worker $Worker -WorkerType 'flamenco'
}

# This function launches ALL workers in Windows Terminal tabs
function Start-AllWorkers {
    param (
        [Parameter(Mandatory = $true)]
        [string]$WorkerType
    )

    Write-Host "Ensuring ALL $WorkerType workers are running under controllers..." -ForegroundColor Cyan

    $payloadScript = if ($WorkerType -eq 'CMD') {
        Get-RemoteCmdWorkerCommand
    } else {
        Get-RemoteStandardWorkerCommand
    }

    foreach ($worker in $workers) {
        Ensure-PersistentWorker -Worker $worker -WorkerType 'flamenco' -PayloadScript $payloadScript
    }

    Write-Host "All workers processed. Attach to any worker from the menu to monitor output." -ForegroundColor Green
    Write-Host "Press Enter to return to the menu..." -ForegroundColor Green
    Read-Host | Out-Null
}

# Main menu loop
$exitRequested = $false
while (-not $exitRequested) {
    Clear-Host
    # Display main menu
    Write-Host "==========================================" -ForegroundColor Cyan
    Write-Host "    UNIFIED FLAMENCO WORKER LAUNCHER" -ForegroundColor Cyan
    Write-Host "==========================================" -ForegroundColor Cyan
    Write-Host
    Write-Host "Main Menu:" -ForegroundColor Magenta
    Write-Host "1. Launch Standard Worker (direct worker execution)" -ForegroundColor Yellow
    Write-Host "2. Launch CMD Worker (uses .cmd file - more resilient)" -ForegroundColor Yellow
    Write-Host "3. Exit" -ForegroundColor Yellow
    Write-Host
    
    $mainSelection = Read-Host "Select option (1-3)"
    
    switch ($mainSelection) {
        "1" {
            # Standard Worker Menu
            Clear-Host
            Write-Host "==========================================" -ForegroundColor Cyan
            Write-Host "    STANDARD WORKER SELECTION" -ForegroundColor Cyan
            Write-Host "==========================================" -ForegroundColor Cyan
            Write-Host
            Write-Host "Select a system to connect to:" -ForegroundColor Yellow
            
            foreach ($worker in $workers) {
                Write-Host "$($worker.ID). $($worker.Name)" -ForegroundColor Green
            }
            
            Write-Host "0. Launch ALL workers (separate windows)" -ForegroundColor Magenta
            Write-Host "B. Back to main menu" -ForegroundColor Yellow
            Write-Host
            
            $workerSelection = Read-Host "Enter your selection"
            
            if ($workerSelection -eq "B" -or $workerSelection -eq "b") {
                # Go back to main menu
                continue
            }
            elseif ($workerSelection -eq "0") {
                # Launch all standard workers - handle this specifically to avoid confusion with worker IDs
                Write-Host "Selected: Launch ALL workers" -ForegroundColor Cyan
                Start-AllWorkers -WorkerType "Standard"
            }
            else {
                # Try to convert to integer and launch selected worker
                try {
                    $selectedID = [int]$workerSelection
                    $selectedWorker = $workers | Where-Object { $_.ID -eq $selectedID }
                    
                    if ($selectedWorker) {
                        Start-StandardWorker -Worker $selectedWorker
                    }
                    else {
                        Write-Host "Invalid selection. Worker ID $selectedID not found." -ForegroundColor Red
                        Start-Sleep -Seconds 2
                    }
                }
                catch {
                    Write-Host "Invalid selection. Please try again." -ForegroundColor Red
                    Start-Sleep -Seconds 2
                }
            }
        }
        "2" {
            # CMD Worker Menu
            Clear-Host
            Write-Host "==========================================" -ForegroundColor Cyan
            Write-Host "    CMD WORKER SELECTION" -ForegroundColor Cyan
            Write-Host "==========================================" -ForegroundColor Cyan
            Write-Host
            Write-Host "Select a system to connect to:" -ForegroundColor Yellow
            
            foreach ($worker in $workers) {
                Write-Host "$($worker.ID). $($worker.Name)" -ForegroundColor Green
            }
            
            Write-Host "0. Launch ALL workers (separate windows)" -ForegroundColor Magenta
            Write-Host "B. Back to main menu" -ForegroundColor Yellow
            Write-Host
            
            $cmdSelection = Read-Host "Enter your selection"
            
            if ($cmdSelection -eq "B" -or $cmdSelection -eq "b") {
                # Go back to main menu
                continue
            }
            elseif ($cmdSelection -eq "0") {
                # Launch all CMD workers - handle this specifically to avoid confusion with worker IDs
                Write-Host "Selected: Launch ALL workers" -ForegroundColor Cyan
                Start-AllWorkers -WorkerType "CMD"
            }
            else {
                # Try to convert to integer and launch selected worker
                try {
                    $selectedID = [int]$cmdSelection
                    $selectedWorker = $workers | Where-Object { $_.ID -eq $selectedID }
                    
                    if ($selectedWorker) {
                        Start-CmdWorker -Worker $selectedWorker
                    }
                    else {
                        Write-Host "Invalid selection. Worker ID $selectedID not found." -ForegroundColor Red
                        Start-Sleep -Seconds 2
                    }
                }
                catch {
                    Write-Host "Invalid selection. Please try again." -ForegroundColor Red
                    Start-Sleep -Seconds 2
                }
            }
        }
        "3" {
            # Exit
            $exitRequested = $true
        }
        default {
            Write-Host "Invalid selection. Please try again." -ForegroundColor Red
            Start-Sleep -Seconds 2
        }
    }
}

Write-Host "`nExiting Unified Flamenco Launcher. Goodbye!" -ForegroundColor Cyan