#Requires -Version 5.1
<#
.SYNOPSIS
  Mirror a public image (default: Docker Hub via optional China mirror) to your registry,
  keeping only selected platforms (default linux/amd64 + linux/arm64), then push one multi-arch tag.

.DESCRIPTION
  Uses skopeo copy per digest + docker buildx imagetools create.
  Docker Desktop users: credentials are resolved from credential helper into a UTF-8 no-BOM auth file.

.PARAMETER Source
  Docker Hub style: redis:latest, nginx:alpine, bitnami/redis:7, library/redis:7, docker.io/prom/prometheus:latest
  Or any registry: ghcr.io/org/image:tag (no Docker-Hub mirror rewrite)

.PARAMETER DestRepo
  Target repository without tag, e.g. registry.cn-shanghai.aliyuncs.com/myns/redis

.PARAMETER DestTag
  Tag on destination (default: tag parsed from -Source)

.PARAMETER MirrorRegistry
  Pull-through mirror host for Docker Hub only, e.g. docker.m.daocloud.io. Empty = use docker.io directly.

.PARAMETER Platforms
  OS/arch list to keep, default linux/amd64 and linux/arm64.

.PARAMETER SkopeoParallelCopies
  skopeo --image-parallel-copies (layers in flight). 0 = auto (higher default for better upload utilization).
  If upload still looks underused, try 32 or 40; if CPU/memory spikes, lower to 12.

  Other tuning (outside this script): use ACR region closest to you; Docker Desktop -> Resources give more CPUs/RAM;
  avoid VPN throttling; wired Ethernet; temporarily pause heavy uploads on same link.

.EXAMPLE
  .\mirror-image-to-registry.ps1 -Source nginx:alpine -DestRepo registry.cn-shanghai.aliyuncs.com/myns/nginx -DestTag alpine

.EXAMPLE
  .\mirror-image-to-registry.ps1 -Source docker.io/library/redis:latest -DestRepo registry.example.com/prod/redis -MirrorRegistry ""
#>

param(
    [Parameter(Mandatory = $true)][string] $Source,
    [Parameter(Mandatory = $true)][string] $DestRepo,
    [string] $DestTag = "",
    [string] $MirrorRegistry = "docker.m.daocloud.io",
    [string[]] $Platforms = @("linux/amd64", "linux/arm64"),
    [string] $SkopeoImage = "quay.io/skopeo/stable:latest",
    [string] $DestUsername = "",
    [string] $DestPassword = "",
    [int] $SkopeoParallelCopies = 0
)

$ErrorActionPreference = "Stop"

# Parallel layer copies: 0 = auto (scale with CPU cores, capped)
$nCpu = [Environment]::ProcessorCount
if ($SkopeoParallelCopies -le 0) {
    $script:EffectiveParallelCopies = [Math]::Min(40, [Math]::Max(16, $nCpu * 4))
    Write-Host "SkopeoParallelCopies=auto -> $script:EffectiveParallelCopies (logical CPUs: $nCpu)"
} else {
    $script:EffectiveParallelCopies = $SkopeoParallelCopies
    Write-Host "SkopeoParallelCopies=$script:EffectiveParallelCopies"
}

function Test-CmdInPath {
    param([string]$Name)
    return [bool](Get-Command $Name -ErrorAction SilentlyContinue)
}

function Get-RegistryHostFromRepo {
    param([string]$Repo)
    return ($Repo -split '/')[0]
}

function Get-DockerConfigPath {
    return (Join-Path $env:USERPROFILE ".docker\config.json")
}

function Write-JsonFileUtf8NoBom {
    param([string]$Path, [string]$JsonText)
    $utf8NoBom = New-Object System.Text.UTF8Encoding $false
    [System.IO.File]::WriteAllText($Path, $JsonText, $utf8NoBom)
}

function New-SkopeoAuthJsonFile {
    param(
        [Parameter(Mandatory = $true)][string]$RegistryHost,
        [Parameter(Mandatory = $true)][string]$Username,
        [Parameter(Mandatory = $true)][string]$Password
    )
    $pair = "${Username}:${Password}"
    $authB64 = [Convert]::ToBase64String([System.Text.Encoding]::UTF8.GetBytes($pair))
    $payload = @{
        auths = @{
            $RegistryHost = @{
                auth = $authB64
            }
        }
    }
    $path = [System.IO.Path]::GetTempFileName() + ".json"
    $jsonText = $payload | ConvertTo-Json -Compress -Depth 6
    Write-JsonFileUtf8NoBom -Path $path -JsonText $jsonText
    return $path
}

function Get-InlineAuthFromDockerConfig {
    param([string]$DockerConfigPath, [string]$RegistryHost)
    if (-not (Test-Path -LiteralPath $DockerConfigPath)) { return $null }
    $txt = Get-Content -LiteralPath $DockerConfigPath -Raw -ErrorAction Stop
    $cfg = $txt | ConvertFrom-Json
    if (-not $cfg.auths) { return $null }
    foreach ($p in $cfg.auths.PSObject.Properties) {
        if ($p.Name -eq $RegistryHost -and $p.Value.auth) {
            return [string]$p.Value.auth
        }
    }
    return $null
}

function Get-CredentialFromDockerHelper {
    param([string]$RegistryHost, [string]$HelperName)
    if (-not $HelperName) { return $null }
    $exeName = "docker-credential-$HelperName"
    $exe = Get-Command $exeName -ErrorAction SilentlyContinue
    if (-not $exe) {
        $candidates = @(
            (Join-Path ${env:ProgramFiles} "Docker\Docker\resources\bin\$exeName.exe"),
            (Join-Path ${env:ProgramFiles} "Docker\Docker\resources\bin\$HelperName.exe")
        )
        foreach ($c in $candidates) {
            if (Test-Path -LiteralPath $c) { $exe = Get-Command $c -ErrorAction SilentlyContinue; break }
        }
    }
    if (-not $exe) {
        Write-Warning "Credential helper '$exeName' not found on PATH or Docker resources."
        return $null
    }
    $psi = New-Object System.Diagnostics.ProcessStartInfo
    $psi.FileName = $exe.Source
    $psi.Arguments = "get"
    $psi.RedirectStandardInput = $true
    $psi.RedirectStandardOutput = $true
    $psi.RedirectStandardError = $true
    $psi.UseShellExecute = $false
    $psi.CreateNoWindow = $true
    $p = New-Object System.Diagnostics.Process
    $p.StartInfo = $psi
    [void]$p.Start()
    $p.StandardInput.WriteLine($RegistryHost)
    $p.StandardInput.Close()
    $out = $p.StandardOutput.ReadToEnd()
    $err = $p.StandardError.ReadToEnd()
    $p.WaitForExit()
    if ($p.ExitCode -ne 0) {
        Write-Warning "docker-credential-$HelperName get failed (exit $($p.ExitCode)): $err"
        return $null
    }
    try {
        $j = $out | ConvertFrom-Json
        return @{ Username = $j.Username; Secret = $j.Secret }
    } catch {
        return $null
    }
}

function Resolve-DestinationAuthFile {
    param([string]$RegistryHost)
    if ($DestUsername -and $DestPassword) {
        Write-Host "Using -DestUsername / -DestPassword for registry push auth."
        return (New-SkopeoAuthJsonFile -RegistryHost $RegistryHost -Username $DestUsername -Password $DestPassword)
    }
    if ($env:ALIYUN_CR_USER -and $env:ALIYUN_CR_PASS) {
        Write-Host "Using ALIYUN_CR_USER / ALIYUN_CR_PASS for registry push auth."
        return (New-SkopeoAuthJsonFile -RegistryHost $RegistryHost -Username $env:ALIYUN_CR_USER -Password $env:ALIYUN_CR_PASS)
    }
    if ($env:DEST_REGISTRY_USER -and $env:DEST_REGISTRY_PASS) {
        Write-Host "Using DEST_REGISTRY_USER / DEST_REGISTRY_PASS for registry push auth."
        return (New-SkopeoAuthJsonFile -RegistryHost $RegistryHost -Username $env:DEST_REGISTRY_USER -Password $env:DEST_REGISTRY_PASS)
    }
    $dockerCfg = Get-DockerConfigPath
    $inline = Get-InlineAuthFromDockerConfig -DockerConfigPath $dockerCfg -RegistryHost $RegistryHost
    if ($inline) {
        Write-Host "Using inline auth from Docker config.json for $RegistryHost"
        $path = [System.IO.Path]::GetTempFileName() + ".json"
        $payload = @{ auths = @{ $RegistryHost = @{ auth = $inline } } }
        $jsonText = $payload | ConvertTo-Json -Compress -Depth 6
        Write-JsonFileUtf8NoBom -Path $path -JsonText $jsonText
        return $path
    }
    if (Test-Path -LiteralPath $dockerCfg) {
        $cfg = (Get-Content -LiteralPath $dockerCfg -Raw) | ConvertFrom-Json
        $helperName = $null
        if ($cfg.credHelpers) {
            foreach ($hp in $cfg.credHelpers.PSObject.Properties) {
                if ($hp.Name -eq $RegistryHost) {
                    $helperName = [string]$hp.Value
                    break
                }
            }
        }
        if (-not $helperName -and $cfg.credsStore) {
            $helperName = $cfg.credsStore
        }
        if ($helperName) {
            Write-Host "Resolving credentials via docker-credential-$helperName (Docker Desktop)..."
            $c = Get-CredentialFromDockerHelper -RegistryHost $RegistryHost -HelperName $helperName
            if ($c -and $c.Username -and $c.Secret) {
                return (New-SkopeoAuthJsonFile -RegistryHost $RegistryHost -Username $c.Username -Password $c.Secret)
            }
        }
    }
    return $null
}

function Split-ImageRefPathAndTag {
    param([string]$Ref)
    if ($Ref -match '@sha256:') {
        throw "Digest-pinned sources (@sha256:...) are not supported. Use a tag."
    }
    $tag = "latest"
    $path = $Ref
    if ($Ref -match '^(.+):([^:/]+)$') {
        $path = $matches[1]
        $tag = $matches[2]
    }
    return @{ Path = $path; Tag = $tag }
}

function Test-IsDockerHubOnlyPath {
    param([string]$PathPart)
    if ($PathPart -match '^[a-zA-Z0-9][a-zA-Z0-9.-]*\.[a-zA-Z]+/') { return $false }
    return $true
}

function Resolve-DockerHubRepoPath {
    param([string]$PathPart)
    if ($PathPart -notmatch '/') {
        return "library/$PathPart"
    }
    return $PathPart
}

# Returns: IsDockerHub, InspectRef (docker://...), SourceRepoPath for copy by digest (host/repo without tag), Tag
function Resolve-SourceInspectReference {
    param(
        [string]$Source,
        [string]$MirrorRegistry
    )
    $s = $Source.Trim()
    if ($s.StartsWith("docker://")) {
        $s = $s.Substring(9)
    }
    $st = Split-ImageRefPathAndTag $s
    $path = $st.Path
    $tag = $st.Tag

    if (Test-IsDockerHubOnlyPath $path) {
        $repo = Resolve-DockerHubRepoPath $path
        $hubRef = "docker.io/${repo}:${tag}"
        if ($MirrorRegistry) {
            $inspect = "docker://${MirrorRegistry}/${repo}:${tag}"
        } else {
            $inspect = "docker://${hubRef}"
        }
        return @{
            IsDockerHub   = $true
            InspectRef    = $inspect
            SourceRepoPath = if ($MirrorRegistry) { "${MirrorRegistry}/${repo}" } else { "docker.io/${repo}" }
            Tag           = $tag
        }
    }

    # Full registry reference e.g. ghcr.io/org/img:tag
    $firstSlash = $path.IndexOf('/')
    if ($firstSlash -lt 1) { throw "Invalid source: $Source" }
    $hostPart = $path.Substring(0, $firstSlash)
    $repoPath = $path.Substring($firstSlash + 1)
    if (-not $repoPath) { throw "Invalid source: $Source" }
    $inspect = "docker://${hostPart}/${repoPath}:${tag}"
    return @{
        IsDockerHub    = $false
        InspectRef     = $inspect
        SourceRepoPath = "${hostPart}/${repoPath}"
        Tag            = $tag
    }
}

function Get-DockerAuthVolumeArgs {
    param([string]$AuthFilePath)
    if ($AuthFilePath -and (Test-Path -LiteralPath $AuthFilePath)) {
        return @("-v", "${AuthFilePath}:/auth.json:ro")
    }
    return @()
}

function Invoke-SkopeoCopyWithAuth {
    param(
        [string]$SkopeoImage,
        [string]$SourceRef,
        [string]$DestRef,
        [string]$AuthFilePath,
        [int]$ParallelCopies = 16,
        [int]$GoMaxProcs = 0
    )
    $parallelArgs = @()
    if ($ParallelCopies -gt 0) {
        $parallelArgs = @("--image-parallel-copies", "$ParallelCopies")
    }
    if ($GoMaxProcs -le 0) { $GoMaxProcs = [Environment]::ProcessorCount }
    $vol = Get-DockerAuthVolumeArgs -AuthFilePath $AuthFilePath
    if (-not (Test-CmdInPath "skopeo")) {
        if (-not $AuthFilePath) {
            throw "No auth file for skopeo copy. Use -DestUsername/-DestPassword, DEST_REGISTRY_*, ALIYUN_CR_*, or docker login with resolvable credentials."
        }
        # GOMAXPROCS lets the Go runtime use more threads for TLS/registry I/O alongside parallel blob copies
        $dockerLine = @("run", "--rm", "-e", "GOMAXPROCS=$GoMaxProcs") + $vol + @($SkopeoImage) + @("copy", "--authfile", "/auth.json") + $parallelArgs + @($SourceRef, $DestRef)
        & docker @dockerLine
        $script:SkopeoLastExitCode = $LASTEXITCODE
        return
    }
    if ($AuthFilePath) {
        $env:GOMAXPROCS = "$GoMaxProcs"
        try {
            & skopeo copy --authfile $AuthFilePath @parallelArgs $SourceRef $DestRef
        } finally {
            Remove-Item Env:GOMAXPROCS -ErrorAction SilentlyContinue
        }
    } else {
        $env:GOMAXPROCS = "$GoMaxProcs"
        try {
            & skopeo copy @parallelArgs $SourceRef $DestRef
        } finally {
            Remove-Item Env:GOMAXPROCS -ErrorAction SilentlyContinue
        }
    }
    $script:SkopeoLastExitCode = $LASTEXITCODE
}

function Invoke-SkopeoInspectRaw {
    param([string]$SkopeoImage, [string]$InspectRef)
    if (Test-CmdInPath "skopeo") {
        $rawLines = & skopeo inspect --raw $InspectRef 2>&1
        $code = $LASTEXITCODE
    } else {
        $rawLines = & docker run --rm $SkopeoImage inspect --raw $InspectRef 2>&1
        $code = $LASTEXITCODE
    }
    return [PSCustomObject]@{ RawLines = $rawLines; ExitCode = $code }
}

function Test-PlatformMatch {
    param($ManifestEntry, [string]$OsWant, [string]$ArchWant)
    if (-not $ManifestEntry.platform) { return $false }
    $os = $ManifestEntry.platform.os
    $arch = $ManifestEntry.platform.architecture
    if ($os -ne $OsWant) { return $false }
    if ($arch -eq $ArchWant) { return $true }
    return $false
}

if (-not (Test-CmdInPath "docker")) {
    Write-Error "docker not found."
}

$resolved = Resolve-SourceInspectReference -Source $Source -MirrorRegistry $MirrorRegistry
if (-not $DestTag) {
    $DestTag = $resolved.Tag
}

$destHost = Get-RegistryHostFromRepo -Repo $DestRepo
$destImage = "${DestRepo}:${DestTag}"

if (-not (Test-CmdInPath "skopeo")) {
    Write-Host "Using skopeo from Docker image: $SkopeoImage"
}

Write-Host "Inspect: $($resolved.InspectRef)"
Write-Host "Destination: $destImage"
Write-Host "Platforms: $($Platforms -join ', ')"

$inspectResult = Invoke-SkopeoInspectRaw -SkopeoImage $SkopeoImage -InspectRef $resolved.InspectRef
if ($inspectResult.ExitCode -ne 0) {
    throw "skopeo inspect failed (exit $($inspectResult.ExitCode)): $($inspectResult.RawLines)"
}

$rawLines = $inspectResult.RawLines
$rawJson = if ($rawLines -is [array]) { $rawLines -join "`n" } else { [string]$rawLines }
$rawJson = $rawJson.Trim()
if ($rawJson -notmatch '^\s*\{') {
    $idx = $rawJson.LastIndexOf('{')
    if ($idx -ge 0) { $rawJson = $rawJson.Substring($idx) }
}

$index = $rawJson | ConvertFrom-Json

# Single manifest (not a list): copy tag-to-tag
if (-not $index.manifests) {
    Write-Host "Source is a single manifest (not a multi-arch index). Copying tag as-is..."
    $srcTagRef = $resolved.InspectRef
    $authFileForPush = Resolve-DestinationAuthFile -RegistryHost $destHost
    if (-not $authFileForPush) { throw "Could not resolve push credentials for $destHost." }
    try {
        Invoke-SkopeoCopyWithAuth -SkopeoImage $SkopeoImage -SourceRef $srcTagRef -DestRef "docker://${destImage}" -AuthFilePath $authFileForPush -ParallelCopies $script:EffectiveParallelCopies
        if ($script:SkopeoLastExitCode -ne 0) { throw "skopeo copy failed (exit $($script:SkopeoLastExitCode))" }
        Write-Host "Done: $destImage"
    } finally {
        if ($authFileForPush -and (Test-Path -LiteralPath $authFileForPush)) {
            Remove-Item -LiteralPath $authFileForPush -Force -ErrorAction SilentlyContinue
        }
    }
    exit 0
}

$digestByPlatform = @{}
foreach ($plat in $Platforms) {
    $parts = $plat -split '/', 2
    if ($parts.Length -ne 2) { throw "Invalid platform (use os/arch): $plat" }
    $osW = $parts[0]
    $archW = $parts[1]
    $found = $null
    foreach ($m in $index.manifests) {
        if (Test-PlatformMatch -ManifestEntry $m -OsWant $osW -ArchWant $archW) {
            $found = $m.digest
            break
        }
    }
    if (-not $found) {
        throw "Platform $plat not found in source manifest list. Inspect with: docker buildx imagetools inspect $($resolved.InspectRef -replace '^docker://','')"
    }
    $digestByPlatform[$plat] = $found
}

foreach ($p in $Platforms) {
    Write-Host "$p : $($digestByPlatform[$p])"
}

$srcBase = $resolved.SourceRepoPath
$dstBase = $DestRepo

$authFileForPush = Resolve-DestinationAuthFile -RegistryHost $destHost
if (-not $authFileForPush) {
    throw @"
Could not resolve push credentials for $destHost.

Use: -DestUsername / -DestPassword, or env DEST_REGISTRY_USER + DEST_REGISTRY_PASS, or ALIYUN_CR_USER + ALIYUN_CR_PASS, or docker login.
"@
}

try {
    Write-Host ""
    Write-Host "Tip: After the 'Copying blob ...' lines, the console may stay quiet for many minutes while large layers upload to the registry (not frozen). Watch Task Manager -> Network (send) if unsure."
    $step = 1
    $total = $Platforms.Count
    foreach ($plat in $Platforms) {
        $d = $digestByPlatform[$plat]
        Write-Host ""
        Write-Host "[$step/$total] skopeo copy $plat ..."
        $step++
        Invoke-SkopeoCopyWithAuth -SkopeoImage $SkopeoImage -SourceRef "docker://${srcBase}@${d}" -DestRef "docker://${dstBase}@${d}" -AuthFilePath $authFileForPush -ParallelCopies $script:EffectiveParallelCopies
        if ($script:SkopeoLastExitCode -ne 0) { throw "skopeo copy failed for $plat (exit $($script:SkopeoLastExitCode))" }
    }

    $refsForBuildx = @()
    foreach ($plat in $Platforms) {
        $refsForBuildx += "${dstBase}@$($digestByPlatform[$plat])"
    }

    Write-Host ""
    Write-Host "[$($total + 1)/$($total + 1)] docker buildx imagetools create..."
    & docker buildx imagetools create -t $destImage @refsForBuildx
    if ($LASTEXITCODE -ne 0) { throw "docker buildx imagetools create failed" }

    Write-Host ""
    Write-Host "Done: $destImage"
    Write-Host "Verify: docker buildx imagetools inspect $destImage"
} finally {
    if ($authFileForPush -and (Test-Path -LiteralPath $authFileForPush)) {
        Remove-Item -LiteralPath $authFileForPush -Force -ErrorAction SilentlyContinue
    }
}