#Requires -Version 7.0
# =============================================================================
# batch_rsync.ps1 — Parallel rsync launcher for Windows PowerShell
# =============================================================================
#
# KEY IMPLEMENTATION NOTES
# ------------------------
#
# Parallelism
#   Uses RunspacePool + [PowerShell]::Create() so multiple rsync processes
#   run concurrently without blocking the main thread. All runspace script
#   blocks are self-contained (no closures over outer variables).
#
# As-completed result collection
#   WaitHandle.WaitAny() polls the async handles so results are printed as
#   soon as each task finishes, matching the behaviour of Python's as_completed().
#
# --files-from (recursive-files mode)
#   rsync is fed a NUL-delimited temp file via --from0 --files-from=<tmpfile>
#   rather than stdin, because stdin is not reliably available inside a
#   runspace. The temp file is always cleaned up in a finally block.
#
# Path handling
#   All source paths are converted from Windows backslashes to forward slashes
#   before being passed to rsync, which is required for MSYS2/Cygwin rsync
#   to interpret them correctly (e.g. C:/data/... not C:\data\...).
#
# Argument passing to rsync
#   Uses ProcessStartInfo.ArgumentList (System.Collections.Generic.List<string>)
#   rather than a single argument string. This avoids quoting/escaping issues
#   with paths that contain spaces or special characters.
#   Requires .NET 5+ — hence the #Requires -Version 7.0 constraint.
#
# Auto-tune worker selection
#   Workers are derived from average file size and destination type (local vs
#   remote). Smaller files favour higher concurrency; large files and remote
#   destinations are throttled to avoid saturating the link or dest I/O.
#   The -Workers value is always treated as an upper bound.
#
# Size-balanced chunk partitioning (recursive-files mode)
#   Files are sorted largest-first and distributed across worker buckets with
#   a greedy least-loaded algorithm to minimise the slowest-worker stall.
#
# Remote destination detection
#   Heuristics check for rsync://, :: (daemon), Windows drive letters, UNC
#   paths, and user@host:/path patterns to decide whether the destination is
#   local or remote without requiring an actual connection.
#
# Verbose output
#   Uses PowerShell's built-in -Verbose switch (CmdletBinding) so there is no
#   custom flag. Pass -Verbose on the command line to enable detailed output.
#
# Requirements
#   - PowerShell 7.0+ (for ProcessStartInfo.ArgumentList via .NET 5+)
#   - rsync on PATH (MSYS2, Cygwin, or WSL interop on Windows)
#
# =============================================================================

<#
.SYNOPSIS
    Parallel rsync launcher for Windows PowerShell.

.DESCRIPTION
    Runs multiple rsync invocations in parallel using either top-level entries
    or balanced recursive file chunks. Requires rsync on PATH (e.g. via MSYS2
    or Cygwin on Windows).

    Paths are converted to forward-slash notation for rsync compatibility.
    MSYS2 rsync accepts Windows absolute paths using forward slashes (C:/path/...).

.EXAMPLE
    # Basic sync (recursive-files mode, auto worker count):
    .\batch_rsync.ps1 C:\data user@host:/dest/

.EXAMPLE
    # Auto-tune workers based on file size and local vs remote destination:
    .\batch_rsync.ps1 C:\data user@host:/dest/ -AutoTune

.EXAMPLE
    # Limit to a specific number of parallel workers:
    .\batch_rsync.ps1 C:\data D:\backup -Workers 8

.EXAMPLE
    # Use top-level mode (one rsync process per immediate source entry):
    .\batch_rsync.ps1 C:\data D:\backup -Mode top-level

.EXAMPLE
    # Dry run (show what would be transferred without doing it):
    .\batch_rsync.ps1 C:\data user@host:/dest/ -DryRun

.EXAMPLE
    # Verbose output (command templates, task previews, elapsed time):
    .\batch_rsync.ps1 C:\data D:\backup -Verbose

.EXAMPLE
    # Show rsync output for all tasks:
    .\batch_rsync.ps1 C:\data D:\backup -ShowRsyncOutput

.EXAMPLE
    # Show rsync output only when a task fails:
    .\batch_rsync.ps1 C:\data user@host:/dest/ -ShowRsyncOutputOnFailure

.EXAMPLE
    # Full-featured (auto-tune, verbose, show on failure):
    .\batch_rsync.ps1 C:\data user@host:/dest/ -AutoTune -Verbose -ShowRsyncOutputOnFailure

.EXAMPLE
    # Override rsync flags entirely:
    .\batch_rsync.ps1 C:\data D:\backup -RsyncArgs '-az --checksum'
#>

[CmdletBinding()]
param (
    [Parameter(Mandatory, Position = 0)]
    [string]$Src,

    [Parameter(Mandatory, Position = 1)]
    [string]$Dest,

    [ValidateRange(1, [int]::MaxValue)]
    [int]$Workers = [Math]::Min(8, [Environment]::ProcessorCount),

    [ValidateSet('recursive-files', 'top-level')]
    [string]$Mode = 'recursive-files',

    [string]$RsyncArgs = '-a --partial --info=stats2',

    [switch]$AutoTune,

    [switch]$ShowRsyncOutput,

    [switch]$ShowRsyncOutputOnFailure,

    [switch]$DryRun
)

Set-StrictMode -Version Latest
$ErrorActionPreference = 'Stop'

#region Helper functions

function Format-Bytes {
    param([double]$Bytes)
    $units = @('B', 'KB', 'MB', 'GB', 'TB')
    $value = $Bytes
    foreach ($unit in $units) {
        if ($value -lt 1024 -or $unit -eq 'TB') {
            return '{0:F1}{1}' -f $value, $unit
        }
        $value /= 1024
    }
    return '{0:F1}B' -f $Bytes
}

function Test-RemoteDestination {
    param([string]$Destination)
    if ($Destination -match '^rsync://') { return $true }
    if ($Destination -match '::')        { return $true }
    # Windows drive letters and UNC paths are local
    if ($Destination -match '^[a-zA-Z]:[/\\]') { return $false }
    if ($Destination -match '^\\\\')            { return $false }
    # user@host:/path or host:/path patterns are remote
    return $Destination -match '^[^:/\\]+@?[^:/\\]*:[^/]'
}

function Get-AutoTunedWorkers {
    param(
        [int]   $TotalUnits,
        [double]$AvgSizeBytes,
        [bool]  $IsRemote,
        [int]   $CpuCount,
        [int]   $MaxWorkers
    )
    $cpu   = [Math]::Max(1, $CpuCount)
    $MB1   = 1   * 1024 * 1024
    $MB32  = 32  * 1024 * 1024
    $MB256 = 256 * 1024 * 1024

    if ($AvgSizeBytes -lt $MB1) {
        $base = [Math]::Min(32, $cpu * 4)
    } elseif ($AvgSizeBytes -lt $MB32) {
        $base = [Math]::Min(16, $cpu * 2)
    } elseif ($AvgSizeBytes -lt $MB256) {
        $base = [Math]::Min(8, $cpu)
    } else {
        $base = [Math]::Min(4, [Math]::Max(2, [int]($cpu / 2)))
    }

    if ($IsRemote) {
        $base = [Math]::Max(2, [Math]::Min($base, 8))
    } else {
        $base = [Math]::Min($base + 2, 24)
    }

    $tuned = [Math]::Max(1, [Math]::Min($base, $TotalUnits))
    return [Math]::Min($tuned, $MaxWorkers)
}

function Get-RecursiveFileList {
    param([string]$SrcDir)
    $root = $SrcDir.TrimEnd('\', '/')
    Get-ChildItem -LiteralPath $root -Recurse -File |
        ForEach-Object { $_.FullName.Substring($root.Length).TrimStart('\', '/') }
}

function Get-SizedFiles {
    param([string]$SrcDir, [string[]]$RelFiles)
    foreach ($rel in $RelFiles) {
        $abs = Join-Path $SrcDir $rel
        try   { $size = (Get-Item -LiteralPath $abs -ErrorAction Stop).Length }
        catch { $size = 0 }
        [PSCustomObject]@{ Size = [long]$size; RelPath = $rel }
    }
}

function Get-BalancedChunks {
    param($SizedFiles, [int]$Workers)
    $buckets     = @(1..$Workers | ForEach-Object { [System.Collections.Generic.List[string]]::new() })
    $bucketSizes = [long[]]::new($Workers)

    foreach ($entry in ($SizedFiles | Sort-Object -Property Size -Descending)) {
        $minIdx = 0
        for ($i = 1; $i -lt $Workers; $i++) {
            if ($bucketSizes[$i] -lt $bucketSizes[$minIdx]) { $minIdx = $i }
        }
        $buckets[$minIdx].Add($entry.RelPath)
        $bucketSizes[$minIdx] += $entry.Size
    }

    $buckets | Where-Object { $_.Count -gt 0 }
}

#endregion

#region Runspace script blocks
# These blocks run inside isolated runspaces — all needed logic must be self-contained.

$rsyncChunkScriptBlock = {
    param(
        [string]  $SrcDir,
        [string[]]$RelFiles,
        [string]  $Dest,
        [string[]]$RsyncArgsList,
        [bool]    $IsDryRun
    )

    $filesFromTmp = [System.IO.Path]::GetTempFileName()
    try {
        # Build a NUL-delimited file list using forward slashes for rsync
        $lines    = $RelFiles | ForEach-Object { $_ -replace '\\', '/' }
        $allBytes = [System.Collections.Generic.List[byte]]::new()
        foreach ($line in $lines) {
            $allBytes.AddRange([System.Text.Encoding]::UTF8.GetBytes($line))
            $allBytes.Add([byte]0)
        }
        [System.IO.File]::WriteAllBytes($filesFromTmp, $allBytes.ToArray())

        $srcRoot = ($SrcDir.TrimEnd('\', '/') -replace '\\', '/') + '/'

        $cmdArgs = [System.Collections.Generic.List[string]]::new($RsyncArgsList)
        if ($IsDryRun) { $cmdArgs.Insert(0, '--dry-run') }
        $cmdArgs.AddRange([string[]]@('--from0', "--files-from=$filesFromTmp", $srcRoot, $Dest))

        $psi = [System.Diagnostics.ProcessStartInfo]::new('rsync')
        $psi.RedirectStandardOutput = $true
        $psi.RedirectStandardError  = $true
        $psi.UseShellExecute        = $false
        $psi.CreateNoWindow         = $true
        foreach ($arg in $cmdArgs) { $psi.ArgumentList.Add($arg) }

        try {
            $proc   = [System.Diagnostics.Process]::Start($psi)
            $stdout = $proc.StandardOutput.ReadToEnd()
            $stderr = $proc.StandardError.ReadToEnd()
            $proc.WaitForExit()
            $exitCode = $proc.ExitCode
        } catch {
            $stdout   = ''
            $stderr   = "rsync not found on PATH: $($_.Exception.Message)"
            $exitCode = 127
        }

        return [PSCustomObject]@{
            Label    = "chunk($($RelFiles.Count) files)"
            ExitCode = $exitCode
            Stdout   = $stdout
            Stderr   = $stderr
        }
    } finally {
        Remove-Item -LiteralPath $filesFromTmp -Force -ErrorAction SilentlyContinue
    }
}

$rsyncItemScriptBlock = {
    param(
        [string]  $Item,
        [string]  $Dest,
        [string[]]$RsyncArgsList,
        [bool]    $IsDryRun
    )

    $srcPath = $Item -replace '\\', '/'

    $cmdArgs = [System.Collections.Generic.List[string]]::new($RsyncArgsList)
    if ($IsDryRun) { $cmdArgs.Insert(0, '--dry-run') }
    $cmdArgs.Add($srcPath)
    $cmdArgs.Add($Dest)

    $psi = [System.Diagnostics.ProcessStartInfo]::new('rsync')
    $psi.RedirectStandardOutput = $true
    $psi.RedirectStandardError  = $true
    $psi.UseShellExecute        = $false
    $psi.CreateNoWindow         = $true
    foreach ($arg in $cmdArgs) { $psi.ArgumentList.Add($arg) }

    try {
        $proc   = [System.Diagnostics.Process]::Start($psi)
        $stdout = $proc.StandardOutput.ReadToEnd()
        $stderr = $proc.StandardError.ReadToEnd()
        $proc.WaitForExit()
        $exitCode = $proc.ExitCode
    } catch {
        $stdout   = ''
        $stderr   = "rsync not found on PATH: $($_.Exception.Message)"
        $exitCode = 127
    }

    return [PSCustomObject]@{
        Label    = $Item
        ExitCode = $exitCode
        Stdout   = $stdout
        Stderr   = $stderr
    }
}

#endregion

#region Parallel job runner (as-completed using WaitHandle.WaitAny)

function Invoke-ParallelJobs {
    param(
        [object[]]    $Jobs,       # PSCustomObjects with .PS and .Handle
        [bool]        $ShowOutput,
        [bool]        $ShowOnFailOnly
    )

    $results = [System.Collections.Generic.List[object]]::new()
    $pending = [System.Collections.Generic.List[object]]::new()
    foreach ($j in $Jobs) { $pending.Add($j) }

    while ($pending.Count -gt 0) {
        $handles = @($pending | ForEach-Object { $_.Handle.AsyncWaitHandle })
        $idx     = [System.Threading.WaitHandle]::WaitAny($handles)
        $job     = $pending[$idx]
        $pending.RemoveAt($idx)

        $r = $job.PS.EndInvoke($job.Handle)[0]
        $job.PS.Dispose()
        $results.Add($r)

        $status = if ($r.ExitCode -eq 0) { 'OK' } else { "ERR($($r.ExitCode))" }
        Write-Host "[$status] $($r.Label)"

        $printOutput = $ShowOutput -and (-not $ShowOnFailOnly -or $r.ExitCode -ne 0)
        if ($printOutput) {
            if ($r.Stdout.Trim()) {
                Write-Host "--- stdout: $($r.Label) ---"
                Write-Host $r.Stdout.TrimEnd()
            }
            if ($r.Stderr.Trim()) {
                Write-Warning "--- stderr: $($r.Label) ---`n$($r.Stderr.TrimEnd())"
            }
        }
    }

    return $results
}

#endregion

#region Main

$timer = [System.Diagnostics.Stopwatch]::StartNew()

$rsyncArgsList = $RsyncArgs -split '\s+' | Where-Object { $_ -ne '' }
$srcPath       = $Src.TrimEnd('\', '/')
$isRemote      = Test-RemoteDestination -Destination $Dest
$cpuCount      = [Environment]::ProcessorCount
$showOutput    = $ShowRsyncOutput -or $ShowRsyncOutputOnFailure
$showOnFail    = $ShowRsyncOutputOnFailure -and -not $ShowRsyncOutput

if ($ShowRsyncOutput -and $ShowRsyncOutputOnFailure) {
    Write-Warning 'Both -ShowRsyncOutput and -ShowRsyncOutputOnFailure set; showing output for all tasks.'
}

$allResults = $null

if ($Mode -eq 'recursive-files' -and (Test-Path -LiteralPath $srcPath -PathType Container)) {

    $relFiles = @(Get-RecursiveFileList -SrcDir $srcPath)
    if (-not $relFiles -or $relFiles.Count -eq 0) {
        Write-Error "No files found for: $Src"
        exit 2
    }

    $sizedFiles = @(Get-SizedFiles -SrcDir $srcPath -RelFiles $relFiles)
    $totalBytes = ($sizedFiles | Measure-Object -Property Size -Sum).Sum
    $avgSize    = if ($sizedFiles.Count -gt 0) { $totalBytes / $sizedFiles.Count } else { 0 }

    if ($AutoTune) {
        $effectiveWorkers = Get-AutoTunedWorkers -TotalUnits $sizedFiles.Count -AvgSizeBytes $avgSize `
            -IsRemote $isRemote -CpuCount $cpuCount -MaxWorkers $Workers
        $destKind = if ($isRemote) { 'remote' } else { 'local' }
        Write-Host "Auto-tuned workers=$effectiveWorkers ($destKind dest, avg file $(Format-Bytes $avgSize))"
    } else {
        $effectiveWorkers = [Math]::Min($Workers, $sizedFiles.Count)
    }

    $chunks = @(Get-BalancedChunks -SizedFiles $sizedFiles -Workers $effectiveWorkers)
    Write-Host "Starting $($chunks.Count) rsync chunk task(s) with $effectiveWorkers worker(s)..."

    if ($VerbosePreference -ne 'SilentlyContinue') {
        $srcRoot     = ($srcPath -replace '\\', '/') + '/'
        $previewArgs = (@('rsync') + $rsyncArgsList + @('--from0', '--files-from=<tmpfile>', $srcRoot, $Dest))
        if ($DryRun) { $previewArgs = @('rsync', '--dry-run') + $previewArgs[1..($previewArgs.Count-1)] }
        Write-Verbose "Command template: $($previewArgs -join ' ')"
    }

    $pool = [RunspaceFactory]::CreateRunspacePool(1, $effectiveWorkers)
    $pool.Open()

    $jobs = foreach ($chunk in $chunks) {
        $ps = [PowerShell]::Create()
        $ps.RunspacePool = $pool
        [void]$ps.AddScript($rsyncChunkScriptBlock)
        [void]$ps.AddParameter('SrcDir',       $srcPath)
        [void]$ps.AddParameter('RelFiles',      [string[]]$chunk)
        [void]$ps.AddParameter('Dest',          $Dest)
        [void]$ps.AddParameter('RsyncArgsList', [string[]]$rsyncArgsList)
        [void]$ps.AddParameter('IsDryRun',      $DryRun.IsPresent)
        [PSCustomObject]@{ PS = $ps; Handle = $ps.BeginInvoke() }
    }

    $allResults = Invoke-ParallelJobs -Jobs $jobs -ShowOutput $showOutput -ShowOnFailOnly $showOnFail
    $pool.Close()
    $pool.Dispose()

} else {

    if (Test-Path -LiteralPath $srcPath -PathType Container) {
        $items = @(Get-ChildItem -LiteralPath $srcPath | ForEach-Object { $_.FullName })
    } elseif ($srcPath -match '[*?]') {
        $items = @(Resolve-Path -Path $srcPath | ForEach-Object { $_.ProviderPath })
    } else {
        $items = @($srcPath)
    }

    if (-not $items -or $items.Count -eq 0) {
        Write-Error "No files found for: $Src"
        exit 2
    }

    $sizeSamples = foreach ($item in $items) {
        try {
            $fi = Get-Item -LiteralPath $item -ErrorAction Stop
            if ($fi.PSIsContainer) { [long](64 * 1024 * 1024) } else { $fi.Length }
        } catch { [long]0 }
    }
    $totalBytes = ($sizeSamples | Measure-Object -Sum).Sum
    $avgSize    = if ($items.Count -gt 0) { $totalBytes / $items.Count } else { 0 }

    if ($AutoTune) {
        $effectiveWorkers = Get-AutoTunedWorkers -TotalUnits $items.Count -AvgSizeBytes $avgSize `
            -IsRemote $isRemote -CpuCount $cpuCount -MaxWorkers $Workers
        $destKind = if ($isRemote) { 'remote' } else { 'local' }
        Write-Host "Auto-tuned workers=$effectiveWorkers ($destKind dest, avg unit $(Format-Bytes $avgSize))"
    } else {
        $effectiveWorkers = [Math]::Min($Workers, $items.Count)
    }

    Write-Host "Starting $($items.Count) rsync task(s) with $effectiveWorkers worker(s)..."

    if ($VerbosePreference -ne 'SilentlyContinue') {
        Write-Verbose 'Task list preview:'
        $previewItems = $items | Select-Object -First 10
        for ($i = 0; $i -lt $previewItems.Count; $i++) {
            Write-Verbose ('  {0}. {1}' -f ($i + 1), $previewItems[$i])
        }
        if ($items.Count -gt 10) { Write-Verbose "  ... and $($items.Count - 10) more" }
    }

    $pool = [RunspaceFactory]::CreateRunspacePool(1, $effectiveWorkers)
    $pool.Open()

    $jobs = foreach ($item in $items) {
        if ($VerbosePreference -ne 'SilentlyContinue') {
            $previewArgs = (@('rsync') + $rsyncArgsList + @(($item -replace '\\', '/'), $Dest))
            if ($DryRun) { $previewArgs = @('rsync', '--dry-run') + $previewArgs[1..($previewArgs.Count-1)] }
            Write-Verbose "Scheduling: $($previewArgs -join ' ')"
        }
        $ps = [PowerShell]::Create()
        $ps.RunspacePool = $pool
        [void]$ps.AddScript($rsyncItemScriptBlock)
        [void]$ps.AddParameter('Item',          $item)
        [void]$ps.AddParameter('Dest',          $Dest)
        [void]$ps.AddParameter('RsyncArgsList', [string[]]$rsyncArgsList)
        [void]$ps.AddParameter('IsDryRun',      $DryRun.IsPresent)
        [PSCustomObject]@{ PS = $ps; Handle = $ps.BeginInvoke() }
    }

    $allResults = Invoke-ParallelJobs -Jobs $jobs -ShowOutput $showOutput -ShowOnFailOnly $showOnFail
    $pool.Close()
    $pool.Dispose()
}

#endregion

#region Summary

$failed  = @($allResults | Where-Object { $_.ExitCode -ne 0 })
$elapsed = [Math]::Round($timer.Elapsed.TotalSeconds, 2)
Write-Host "`nCompleted: $($allResults.Count) tasks, $($failed.Count) failed in ${elapsed}s."

if ($failed.Count -gt 0) {
    Write-Host 'Failures detail:'
    foreach ($r in $failed) {
        Write-Host "--- $($r.Label) (exit $($r.ExitCode)) ---"
        if ($r.Stdout.Trim()) { Write-Host $r.Stdout.TrimEnd() }
        if ($r.Stderr.Trim()) { Write-Host $r.Stderr.TrimEnd() }
    }
    exit 1
}

exit 0

#endregion