param(
    [Parameter(Mandatory = $true)]
    [string]$WorkerName,

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

    [string]$PayloadBase64,

    [string]$PayloadBase64Path,

[string]$DataRoot = (Join-Path ([Environment]::GetFolderPath('LocalApplicationData')) 'UnifiedWorkers'),

    [int]$MaxRestarts = 5,

    [int]$RestartDelaySeconds = 10
)

$ErrorActionPreference = 'Stop'
[System.Management.Automation.Runspaces.Runspace]::DefaultRunspace = $Host.Runspace

# region Path setup
$workerRoot = Join-Path -Path $DataRoot -ChildPath $WorkerType
$instanceRoot = Join-Path -Path $workerRoot -ChildPath $WorkerName
New-Item -ItemType Directory -Path $instanceRoot -Force | Out-Null

$logsRoot = Join-Path -Path $instanceRoot -ChildPath 'logs'
$stateRoot = Join-Path -Path $instanceRoot -ChildPath 'state'
New-Item -ItemType Directory -Path $logsRoot -Force | Out-Null
New-Item -ItemType Directory -Path $stateRoot -Force | Out-Null

$logPath = Join-Path -Path $logsRoot -ChildPath 'worker.log'
$metaPath = Join-Path -Path $stateRoot -ChildPath 'worker-info.json'
$commandPath = Join-Path -Path $stateRoot -ChildPath 'commands.txt'
$payloadPath = Join-Path -Path $stateRoot -ChildPath "payload.ps1"
# endregion

# region Logging
$logStream = [System.IO.FileStream]::new(
    $logPath,
    [System.IO.FileMode]::Append,
    [System.IO.FileAccess]::Write,
    [System.IO.FileShare]::ReadWrite
)
$logWriter = [System.IO.StreamWriter]::new($logStream, [System.Text.Encoding]::UTF8)
$logWriter.AutoFlush = $true

# Create C# event handler class that doesn't require PowerShell runspace
if (-not ("UnifiedWorkers.ProcessLogHandler" -as [type])) {
    $csharpCode = @'
using System;
using System.Diagnostics;
using System.IO;

namespace UnifiedWorkers
{
    public sealed class ProcessLogHandler
    {
        private readonly TextWriter _writer;
        private readonly string _prefix;
        private readonly object _lock = new object();

        public ProcessLogHandler(TextWriter writer, string prefix)
        {
            _writer = writer ?? throw new ArgumentNullException("writer");
            _prefix = prefix ?? throw new ArgumentNullException("prefix");
        }

        public void OnDataReceived(object sender, DataReceivedEventArgs e)
        {
            if (e == null || string.IsNullOrEmpty(e.Data))
            {
                return;
            }

            lock (_lock)
            {
                try
                {
                    var timestamp = DateTime.UtcNow.ToString("u");
                    _writer.WriteLine(string.Format("[{0} {1}] {2}", _prefix, timestamp, e.Data));
                    _writer.Flush();
                }
                catch
                {
                    // Ignore write errors to prevent cascading failures
                }
            }
        }
    }
}
'@
    Add-Type -TypeDefinition $csharpCode -ErrorAction Stop
}

function Write-LogLine {
    param(
        [string]$Prefix,
        [string]$Message
    )

    if (-not $logWriter) { return }
    $timestamp = (Get-Date).ToString('u')
    $logWriter.WriteLine("[$Prefix $timestamp] $Message")
}

function Write-ControllerLog {
    param([string]$Message)
    Write-LogLine -Prefix 'CTRL' -Message $Message
}

function Write-FatalLog {
    param([string]$Message)

    try {
        Write-ControllerLog $Message
    }
    catch {
        $timestamp = (Get-Date).ToString('u')
        $fallback = "[CTRL $timestamp] $Message"
        try {
            [System.IO.File]::AppendAllText($logPath, $fallback + [Environment]::NewLine, [System.Text.Encoding]::UTF8)
        }
        catch {
            # last resort: write to host
            Write-Error $fallback
        }
    }
}
# endregion

# region Helpers

function Resolve-PayloadBase64 {
    if ($PayloadBase64) {
        return $PayloadBase64.Trim()
    }

    if ($PayloadBase64Path) {
        if (-not (Test-Path $PayloadBase64Path)) {
            throw "Payload file '$PayloadBase64Path' not found."
        }

        $content = Get-Content -Path $PayloadBase64Path -Raw
        if ([string]::IsNullOrWhiteSpace($content)) {
            throw "Payload file '$PayloadBase64Path' is empty."
        }
        
        return $content.Trim()
    }

    throw "No payload data provided to controller."
}

function Write-Metadata {
    param(
        [string]$Status,
        [nullable[int]]$WorkerPid = $null,
        [nullable[int]]$ControllerPid = $PID,
        [int]$Restarts = 0,
        [nullable[int]]$LastExitCode = $null
    )

    $payload = [pscustomobject]@{
        WorkerName     = $WorkerName
        WorkerType     = $WorkerType
        Status         = $Status
        ControllerPid  = $ControllerPid
        WorkerPid      = $WorkerPid
        Restarts       = $Restarts
        LastExitCode   = $LastExitCode
        LogPath        = $logPath
        CommandPath    = $commandPath
        PayloadPath    = $payloadPath
        UpdatedAtUtc   = (Get-Date).ToUniversalTime()
    }

    $payload | ConvertTo-Json -Depth 5 | Set-Content -Path $metaPath -Encoding UTF8
}

function Get-PendingCommands {
    if (-not (Test-Path $commandPath)) {
        return @()
    }

    try {
        $lines = Get-Content -Path $commandPath -ErrorAction Stop
        Remove-Item -Path $commandPath -Force -ErrorAction SilentlyContinue
        return $lines | Where-Object { -not [string]::IsNullOrWhiteSpace($_) }
    }
    catch {
        return @()
    }
}
# endregion

try {
    # record initial state before launching worker
    Write-Metadata -Status 'initializing' -WorkerPid $null -ControllerPid $PID -Restarts 0

    $resolvedPayloadBase64 = Resolve-PayloadBase64
    $PayloadBase64 = $resolvedPayloadBase64

    try {
        # Write payload script to disk
        # The payload is base64-encoded UTF-16 (Unicode), so decode it properly
        Write-ControllerLog "Decoding payload base64 (length: $($resolvedPayloadBase64.Length))"
        $payloadBytes = [Convert]::FromBase64String($resolvedPayloadBase64)
        Write-ControllerLog "Decoded payload to $($payloadBytes.Length) bytes"
        
        # Convert UTF-16 bytes back to string, then write as UTF-8 (PowerShell's preferred encoding)
        $payloadText = [Text.Encoding]::Unicode.GetString($payloadBytes)
        [IO.File]::WriteAllText($payloadPath, $payloadText, [Text.Encoding]::UTF8)
        Write-ControllerLog "Payload written to $payloadPath ($($payloadText.Length) characters)"
    }
    catch {
        Write-FatalLog "Unable to write payload: $($_.Exception.Message)"
        if ($_.Exception.InnerException) {
            Write-FatalLog "Inner exception: $($_.Exception.InnerException.Message)"
        }
        throw
    }

    $restartCount = 0
    $controllerPid = $PID

    while ($restartCount -le $MaxRestarts) {
        try {
            # Initialize worker process
            $psi = [System.Diagnostics.ProcessStartInfo]::new()
            $pwsh = Get-Command pwsh -ErrorAction SilentlyContinue
            if ($pwsh) {
                $psi.FileName = $pwsh.Source
            }
            else {
                $psi.FileName = (Get-Command powershell -ErrorAction Stop).Source
            }

            $psi.Arguments = "-NoLogo -NoProfile -ExecutionPolicy Bypass -File `"$payloadPath`""
            $psi.UseShellExecute = $false
            $psi.RedirectStandardInput = $true
            $psi.RedirectStandardOutput = $true
            $psi.RedirectStandardError = $true
            $psi.CreateNoWindow = $true

            $workerProcess = New-Object System.Diagnostics.Process
            $workerProcess.StartInfo = $psi

            if (-not $workerProcess.Start()) {
                throw "Failed to start worker process."
            }

            Write-ControllerLog "Worker process started with PID $($workerProcess.Id)"
            Write-Metadata -Status 'running' -WorkerPid $workerProcess.Id -ControllerPid $controllerPid -Restarts $restartCount

            $stdoutHandler = [UnifiedWorkers.ProcessLogHandler]::new($logWriter, 'OUT')
            $stderrHandler = [UnifiedWorkers.ProcessLogHandler]::new($logWriter, 'ERR')
            
            $outputHandler = [System.Diagnostics.DataReceivedEventHandler]$stdoutHandler.OnDataReceived
            $errorHandler = [System.Diagnostics.DataReceivedEventHandler]$stderrHandler.OnDataReceived

            $workerProcess.add_OutputDataReceived($outputHandler)
            $workerProcess.add_ErrorDataReceived($errorHandler)
            $workerProcess.BeginOutputReadLine()
            $workerProcess.BeginErrorReadLine()

            while (-not $workerProcess.HasExited) {
                $commands = Get-PendingCommands
                foreach ($command in $commands) {
                    $trimmed = $command.Trim()
                    if (-not $trimmed) { continue }

                    Write-ControllerLog "Received command '$trimmed'"
                    try {
                        $workerProcess.StandardInput.WriteLine($trimmed)
                        $workerProcess.StandardInput.Flush()
                    }
                    catch {
                        Write-ControllerLog "Failed to forward command '$trimmed': $($_.Exception.Message)"
                    }

                    if ($trimmed -ieq 'quit') {
                        Write-ControllerLog "Quit command issued. Waiting for worker to exit."
                    }
                }

                Start-Sleep -Milliseconds 500
            }

            $exitCode = $workerProcess.ExitCode
            Write-ControllerLog "Worker exited with code $exitCode"
            Write-Metadata -Status 'stopped' -WorkerPid $null -ControllerPid $controllerPid -Restarts $restartCount -LastExitCode $exitCode

            if ($exitCode -eq 0) {
                break
            }
        }
        catch {
            Write-ControllerLog "Controller error: $($_.Exception.Message)"
        }

        $restartCount++
        if ($restartCount -gt $MaxRestarts) {
            Write-ControllerLog "Maximum restart attempts reached. Controller stopping."
            break
        }

        Write-ControllerLog "Restarting worker in $RestartDelaySeconds seconds (attempt $restartCount of $MaxRestarts)."
        Start-Sleep -Seconds $RestartDelaySeconds
    }

    Write-Metadata -Status 'inactive' -WorkerPid $null -ControllerPid $controllerPid -Restarts $restartCount
    Write-ControllerLog "Controller exiting."
}
catch {
    Write-FatalLog "Fatal controller error: $($_.Exception.Message)"
    if ($_.ScriptStackTrace) {
        Write-FatalLog "Stack: $($_.ScriptStackTrace)"
    }
    throw
}
finally {
    if ($logWriter) {
        $logWriter.Dispose()
    }
    if ($logStream) {
        $logStream.Dispose()
    }
}