param(
    [switch]$DownloadTagger
)

$ErrorActionPreference = "Stop"
$Root = Split-Path -Parent $MyInvocation.MyCommand.Path
$Venv = Join-Path $Root ".venv"
$Python = Join-Path $Venv "Scripts\python.exe"
$SdScripts = Join-Path $Root "sd-scripts"
$SdScriptsCommit = "c162e9039cc3228057b3fc6ae64ce770b87ab9bf"
$TorchCu124Version = "2.6.0"
$TorchvisionCu124Version = "0.21.0"
$TorchCu128Version = "2.7.1"
$TorchvisionCu128Version = "0.22.1"
$SupportedPythonVersions = @("3.11", "3.10")
$BootstrapPythonVersion = "3.11.9"
$BootstrapPythonUrl = "https://www.python.org/ftp/python/3.11.9/python-3.11.9-amd64.exe"
$BootstrapPythonInstallerName = "python-3.11.9-amd64.exe"
$env:PYTHONUTF8 = "1"
$env:PYTHONIOENCODING = "utf-8"
$env:XFORMERS_FORCE_DISABLE_TRITON = "1"

function Invoke-NativeStep {
    param(
        [string]$Name,
        [scriptblock]$Command
    )

    $PreviousErrorActionPreference = $ErrorActionPreference
    $HasNativePreference = Get-Variable -Name PSNativeCommandUseErrorActionPreference -ErrorAction SilentlyContinue
    if ($HasNativePreference) {
        $PreviousNativePreference = $PSNativeCommandUseErrorActionPreference
    }

    try {
        $ErrorActionPreference = "Continue"
        if ($HasNativePreference) {
            $PSNativeCommandUseErrorActionPreference = $false
        }

        & $Command
        $ExitCode = $LASTEXITCODE
    } finally {
        $ErrorActionPreference = $PreviousErrorActionPreference
        if ($HasNativePreference) {
            $PSNativeCommandUseErrorActionPreference = $PreviousNativePreference
        }
    }

    if ($ExitCode -ne 0) {
        throw "$Name failed. Exit code: $ExitCode"
    }
}

function Test-SdScriptsReady {
    if (!(Test-Path $SdScripts)) {
        return $false
    }

    $RequiredPaths = @(
        "library",
        "anima_train_network.py",
        "sdxl_train_network.py",
        "train_network.py",
        "requirements.txt"
    )

    foreach ($RelativePath in $RequiredPaths) {
        if (!(Test-Path (Join-Path $SdScripts $RelativePath))) {
            return $false
        }
    }

    return $true
}

function Assert-CommandAvailable {
    param(
        [string]$CommandName,
        [string]$InstallHint
    )

    if (!(Get-Command $CommandName -ErrorAction SilentlyContinue)) {
        throw "$CommandName was not found. $InstallHint"
    }
}

function New-PythonCandidate {
    param(
        [string]$Label,
        [string]$Exe,
        [string[]]$Args = @()
    )

    return @{
        Label = $Label
        Exe = $Exe
        Args = @($Args)
    }
}

function Get-PythonCandidates {
    $Candidates = @()

    if ($env:COPITA_PYTHON) {
        $Candidates += New-PythonCandidate "COPITA_PYTHON" $env:COPITA_PYTHON
    }

    $PackageRoot = Split-Path -Parent $Root
    $LocalPythonPaths = @()
    foreach ($BaseDir in @($Root, $PackageRoot)) {
        foreach ($RelativePath in @(
            "python\python.exe",
            "Python\python.exe",
            "python310\python.exe",
            "python311\python.exe",
            "Python310\python.exe",
            "Python311\python.exe"
        )) {
            $LocalPythonPaths += Join-Path $BaseDir $RelativePath
        }
    }

    foreach ($LocalPython in $LocalPythonPaths) {
        $Candidates += New-PythonCandidate "local Python ($LocalPython)" $LocalPython
    }

    $Candidates += New-PythonCandidate "Python Launcher 3.11" "py" @("-3.11")
    $Candidates += New-PythonCandidate "Python Launcher 3.10" "py" @("-3.10")
    $Candidates += New-PythonCandidate "default python" "python"

    return $Candidates
}

function Test-PythonCandidate {
    param([hashtable]$Candidate)

    $Exe = [string]$Candidate["Exe"]
    $Args = @($Candidate["Args"])
    $Label = [string]$Candidate["Label"]

    if ([string]::IsNullOrWhiteSpace($Exe)) {
        return $null
    }

    $LooksLikePath = $Exe.Contains("\") -or $Exe.Contains("/") -or [System.IO.Path]::IsPathRooted($Exe)
    if ($LooksLikePath) {
        if (!(Test-Path -LiteralPath $Exe)) {
            return $null
        }
    } elseif (!(Get-Command $Exe -ErrorAction SilentlyContinue)) {
        return $null
    }

    $VersionOutput = & $Exe @Args -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')"
    if ($LASTEXITCODE -ne 0 -or !$VersionOutput) {
        Write-Host "Skipping ${Label}: could not run Python."
        return $null
    }

    $Version = ($VersionOutput | Select-Object -Last 1).Trim()
    if ($SupportedPythonVersions -notcontains $Version) {
        Write-Host "Skipping ${Label}: Python $Version is not supported."
        return $null
    }

    $null = & $Exe @Args -m venv --help 2>$null
    if ($LASTEXITCODE -ne 0) {
        Write-Host "Skipping ${Label}: Python $Version does not provide the venv module."
        return $null
    }

    $Candidate["Version"] = $Version
    return $Candidate
}

function Install-LocalBootstrapPython {
    $PackageRoot = Split-Path -Parent $Root
    $LocalPythonDir = Join-Path $PackageRoot "python"
    $InstallerDir = Join-Path $Root "downloads"
    $InstallerPath = Join-Path $InstallerDir $BootstrapPythonInstallerName
    $InstallLog = Join-Path $InstallerDir "python-3.11-install.log"

    Write-Host ""
    Write-Host "Python 3.10 or 3.11 was not found."
    Write-Host "Downloading Python $BootstrapPythonVersion for Copita only..."
    Write-Host "This does not change the default Python on this PC."

    New-Item -ItemType Directory -Force -Path $InstallerDir | Out-Null
    if (!(Test-Path -LiteralPath $InstallerPath)) {
        Invoke-WebRequest -Uri $BootstrapPythonUrl -OutFile $InstallerPath
    } else {
        Write-Host "Using cached Python installer: $InstallerPath"
    }

    Write-Host "Installing Python $BootstrapPythonVersion into:"
    Write-Host $LocalPythonDir
    Invoke-NativeStep "Installing local Python $BootstrapPythonVersion" {
        & $InstallerPath "/quiet" "InstallAllUsers=0" "TargetDir=$LocalPythonDir" `
            "Include_launcher=0" "InstallLauncherAllUsers=0" "AssociateFiles=0" "Shortcuts=0" `
            "PrependPath=0" "AppendPath=0" "Include_pip=1" "Include_exe=1" "Include_lib=1" "Include_dev=1" "Include_tools=1" `
            "Include_doc=0" "Include_tcltk=0" "Include_test=0" "CompileAll=0" "/log" "$InstallLog"
    }

    $LocalPython = Join-Path $LocalPythonDir "python.exe"
    if (!(Test-Path -LiteralPath $LocalPython)) {
        throw "Python installer finished, but python.exe was not found at $LocalPython"
    }
}

function Resolve-BootstrapPython {
    foreach ($Candidate in (Get-PythonCandidates)) {
        $Resolved = Test-PythonCandidate $Candidate
        if ($Resolved) {
            return $Resolved
        }
    }

    Install-LocalBootstrapPython

    foreach ($Candidate in (Get-PythonCandidates)) {
        $Resolved = Test-PythonCandidate $Candidate
        if ($Resolved) {
            return $Resolved
        }
    }

    throw "Could not prepare Python 3.10 or 3.11 for Copita. Please send setup_log.txt and Copita\downloads\python-3.11-install.log."
}

function Test-SupportedPythonVersion {
    param([string]$Version)
    return $SupportedPythonVersions -contains $Version
}

function Get-CopitaPythonVersion {
    if (!(Test-Path $Python)) {
        return ""
    }

    $Version = & $Python -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')"
    if ($LASTEXITCODE -ne 0 -or !$Version) {
        return ""
    }

    return ($Version | Select-Object -Last 1).Trim()
}

function Assert-SupportedPythonVersion {
    if (!(Test-Path $Python)) {
        throw "Python environment was not created."
    }

    $Version = Get-CopitaPythonVersion
    if (!$Version) {
        throw "Could not check Python version."
    }

    if (!(Test-SupportedPythonVersion $Version)) {
        throw "Copita supports Python 3.10 or 3.11. Current environment is Python $Version. Please install Python 3.10 or 3.11, delete Copita\.venv, and run setup_first_time.bat again."
    }

    Write-Host "Python version: $Version"
}

function Test-CudaTorchBuild {
    if (!(Test-Path $Python)) {
        return $false
    }

    $DesiredTorch = Get-DesiredTorchSpec

    & $Python -c "import importlib.util, sys; spec=importlib.util.find_spec('torch'); sys.exit(1 if spec is None else 0)"
    if ($LASTEXITCODE -ne 0) {
        return $false
    }

    $CheckScript = @"
import sys
import torch

required_cuda = sys.argv[1]
required_torch = sys.argv[2]
cuda = torch.version.cuda or ''
version = torch.__version__.split('+', 1)[0]
ok = cuda.startswith(required_cuda) and version == required_torch
print(f'torch={torch.__version__} cuda={cuda} required=torch {required_torch} / CUDA {required_cuda}')
sys.exit(0 if ok else 1)
"@
    & $Python -c $CheckScript $DesiredTorch.CudaVersion $DesiredTorch.TorchVersion
    return $LASTEXITCODE -eq 0
}

function Get-PrimaryNvidiaGpuInfo {
    $Info = @{
        Name = ""
        ComputeCapability = ""
    }

    if (!(Get-Command "nvidia-smi" -ErrorAction SilentlyContinue)) {
        return $Info
    }

    $Output = & nvidia-smi --query-gpu=name,compute_cap --format=csv,noheader 2>$null
    if ($LASTEXITCODE -ne 0 -or !$Output) {
        return $Info
    }

    $Line = ($Output | Select-Object -First 1)
    if (!$Line) {
        return $Info
    }

    $Parts = $Line -split ",", 2
    $Info.Name = $Parts[0].Trim()
    if ($Parts.Count -gt 1) {
        $Info.ComputeCapability = $Parts[1].Trim()
    }
    return $Info
}

function Test-NeedsCuda128Torch {
    $GpuInfo = Get-PrimaryNvidiaGpuInfo
    $Compute = $GpuInfo.ComputeCapability
    if ($Compute -match '^(\d+)(?:\.(\d+))?$') {
        $Major = [int]$Matches[1]
        if ($Major -ge 12) {
            return $true
        }
    }

    $Name = $GpuInfo.Name
    if ($Name -match '(?i)\bRTX\s+50\d{2}\b' -or $Name -match '(?i)\bRTX\s+5\d{3}\b') {
        return $true
    }

    return $false
}

function Get-DesiredTorchSpec {
    if (Test-NeedsCuda128Torch) {
        return @{
            TorchVersion = $TorchCu128Version
            TorchvisionVersion = $TorchvisionCu128Version
            CudaVersion = "12.8"
            IndexUrl = "https://download.pytorch.org/whl/cu128"
            Label = "CUDA 12.8 PyTorch for RTX 50 series"
        }
    }

    return @{
        TorchVersion = $TorchCu124Version
        TorchvisionVersion = $TorchvisionCu124Version
        CudaVersion = "12.4"
        IndexUrl = "https://download.pytorch.org/whl/cu124"
        Label = "CUDA 12.4 PyTorch"
    }
}

function Install-CudaTorch {
    $DesiredTorch = Get-DesiredTorchSpec
    $GpuInfo = Get-PrimaryNvidiaGpuInfo
    if ($GpuInfo.Name) {
        Write-Host ("Detected GPU: {0} (compute capability {1})" -f $GpuInfo.Name, $GpuInfo.ComputeCapability)
    }
    Write-Host ("Installing {0}..." -f $DesiredTorch.Label)
    Write-Host ("torch {0} / torchvision {1}" -f $DesiredTorch.TorchVersion, $DesiredTorch.TorchvisionVersion)
    Invoke-NativeStep "Installing CUDA PyTorch" {
        & $Python -m pip install --upgrade --force-reinstall `
            "torch==$($DesiredTorch.TorchVersion)" `
            "torchvision==$($DesiredTorch.TorchvisionVersion)" `
            --index-url $DesiredTorch.IndexUrl
    }
}

function Test-PythonPackageInstalled {
    param([string]$PackageName)

    $CheckScript = @"
import importlib.metadata as metadata
import sys

try:
    metadata.version(sys.argv[1])
    print('1')
except metadata.PackageNotFoundError:
    print('0')
"@

    $Result = & $Python -c $CheckScript $PackageName
    return (($Result | Select-Object -Last 1) -eq "1")
}

function Remove-OnnxRuntimeGpuIfInstalled {
    if (Test-PythonPackageInstalled "onnxruntime-gpu") {
        Invoke-NativeStep "Removing onnxruntime-gpu" {
            & $Python -m pip uninstall -y onnxruntime-gpu
        }
    } else {
        Write-Host "onnxruntime-gpu is not installed. Skipping removal."
    }
}

function Update-SdScripts {
    if ((Test-SdScriptsReady) -and !(Test-Path (Join-Path $SdScripts ".git"))) {
        Write-Host "Using bundled sd-scripts."
        return
    }

    if (!(Test-Path $SdScripts)) {
        Assert-CommandAvailable "git" "Bundled sd-scripts is missing, so Git is required to download it."
        Write-Host "Downloading sd-scripts..."
        Invoke-NativeStep "Cloning sd-scripts" { git clone https://github.com/kohya-ss/sd-scripts.git $SdScripts }
    }

    if (!(Test-SdScriptsReady)) {
        Assert-CommandAvailable "git" "Bundled sd-scripts is incomplete, so Git is required to download it again."
        Write-Host "Removing incomplete sd-scripts folder..."
        Remove-Item -LiteralPath $SdScripts -Recurse -Force
        Write-Host "Downloading sd-scripts..."
        Invoke-NativeStep "Cloning sd-scripts" { git clone https://github.com/kohya-ss/sd-scripts.git $SdScripts }
    }

    if (!(Test-Path (Join-Path $SdScripts ".git"))) {
        Write-Host "Using bundled sd-scripts."
        return
    }

    Assert-CommandAvailable "git" "Git is required to verify an sd-scripts git checkout."
    Write-Host "Using verified sd-scripts revision: $SdScriptsCommit"
    Push-Location $SdScripts
    try {
        Invoke-NativeStep "Fetching sd-scripts" { git fetch origin }
        Invoke-NativeStep "Checking out verified sd-scripts" { git checkout --force $SdScriptsCommit }
        Invoke-NativeStep "Resetting sd-scripts" { git reset --hard $SdScriptsCommit }
    } finally {
        Pop-Location
    }
}

function Apply-SdScriptsCompatibilityPatch {
    $AnimaUtils = Join-Path $SdScripts "library\anima_utils.py"
    if (!(Test-Path $AnimaUtils)) {
        return
    }

    $Text = Get-Content -Raw -LiteralPath $AnimaUtils
    $Old = 'rename_hooks = WeightTransformHooks(rename_hook=lambda k: k[len("net.") :] if k.startswith("net.") else k)'
    $New = 'rename_hooks = WeightTransformHooks(rename_hook=normalize_anima_checkpoint_key)'
    $Helper = @'
def normalize_anima_checkpoint_key(key: str) -> str:
    for prefix in ("model.diffusion_model.", "diffusion_model.", "net."):
        if key.startswith(prefix):
            return key[len(prefix) :]
    return key


'@
    $FilterHelper = @'
def filter_anima_model_state_dict(sd, model):
    model_keys = set(model.state_dict().keys())
    ignored_keys = [key for key in sd.keys() if key not in model_keys]
    if not ignored_keys:
        return sd

    logger.info(
        "Ignoring %d non-DiT checkpoint keys: %s",
        len(ignored_keys),
        ignored_keys[:5],
    )
    return {key: value for key, value in sd.items() if key in model_keys}


'@

    $BrokenInlinePattern = '(?ms)^ {8}def normalize_anima_checkpoint_key\(k\):\r?\n {8}for prefix in \("model\.diffusion_model\.", "diffusion_model\.", "net\."\):\r?\n {12}if k\.startswith\(prefix\):\r?\n {16}return k\[len\(prefix\) :\]\r?\n {8}return k\r?\n\r?\n {4}rename_hooks = WeightTransformHooks\(rename_hook=normalize_anima_checkpoint_key\)'
    if ([regex]::IsMatch($Text, $BrokenInlinePattern)) {
        Write-Host "Repairing broken Anima checkpoint compatibility patch..."
        $Text = [regex]::Replace($Text, $BrokenInlinePattern, "    $New")
    }

    if ($Text.Contains($Old)) {
        Write-Host "Applying Anima checkpoint compatibility patch..."
        $Text = $Text.Replace($Old, $New)
    }

    if (-not [regex]::IsMatch($Text, '(?m)^def normalize_anima_checkpoint_key\(')) {
        $Marker = "def load_anima_model("
        if ($Text.Contains($Marker)) {
            $Text = $Text.Replace($Marker, "$Helper$Marker")
        } else {
            throw "Could not find load_anima_model in $AnimaUtils"
        }
    }

    if (-not [regex]::IsMatch($Text, '(?m)^def filter_anima_model_state_dict\(')) {
        $Marker = "def load_anima_model("
        if ($Text.Contains($Marker)) {
            $Text = $Text.Replace($Marker, "$FilterHelper$Marker")
        } else {
            throw "Could not find load_anima_model in $AnimaUtils"
        }
    }

    $OldFp8Block = @'
    if fp8_scaled:
        apply_fp8_monkey_patch(model, sd, use_scaled_mm=False)

        if loading_device.type != "cpu":
            # make sure all the model weights are on the loading_device
            logger.info(f"Moving weights to {loading_device}")
            for key in sd.keys():
                sd[key] = sd[key].to(loading_device)

    missing, unexpected = model.load_state_dict(sd, strict=False, assign=True)
'@
    $NewFp8Block = @'
    if fp8_scaled:
        apply_fp8_monkey_patch(model, sd, use_scaled_mm=False)
        sd = filter_anima_model_state_dict(sd, model)

        if loading_device.type != "cpu":
            # make sure all the model weights are on the loading_device
            logger.info(f"Moving weights to {loading_device}")
            for key in sd.keys():
                sd[key] = sd[key].to(loading_device)
    else:
        sd = filter_anima_model_state_dict(sd, model)

    missing, unexpected = model.load_state_dict(sd, strict=False, assign=True)
'@
    if ($Text.Contains($OldFp8Block)) {
        Write-Host "Applying Anima bundled-VAE checkpoint filter patch..."
        $Text = $Text.Replace($OldFp8Block, $NewFp8Block)
    } else {
        $Fp8Pattern = '(?ms)^    if fp8_scaled:\r?\n        apply_fp8_monkey_patch\(model, sd, use_scaled_mm=False\)\r?\n\r?\n        if loading_device.type != "cpu":\r?\n            # make sure all the model weights are on the loading_device\r?\n            logger\.info\(f"Moving weights to \{loading_device\}"\)\r?\n            for key in sd\.keys\(\):\r?\n                sd\[key\] = sd\[key\]\.to\(loading_device\)\r?\n\r?\n    missing, unexpected = model\.load_state_dict\(sd, strict=False, assign=True\)'
        if ([regex]::IsMatch($Text, $Fp8Pattern)) {
            Write-Host "Applying Anima bundled-VAE checkpoint filter patch..."
            $Text = [regex]::Replace($Text, $Fp8Pattern, $NewFp8Block.TrimEnd())
        }
    }

    if ($Text.Contains("rename_hook=normalize_anima_checkpoint_key") -and $Text.Contains("model.diffusion_model.") -and $Text.Contains("filter_anima_model_state_dict(sd, model)")) {
        Write-Host "Anima checkpoint compatibility patch: OK"
        Set-Content -LiteralPath $AnimaUtils -Value $Text -Encoding UTF8 -NoNewline
    } else {
        Write-Host "Anima checkpoint compatibility patch: skipped"
    }
}

function Test-SdScriptsIntegration {
    $Smoke = @"
import sys
sys.path.insert(0, r'''$Root''')
import trainer

sdxl_train_network, sdxl_merge, sdxl_resize = trainer.import_sdxl_scripts()
anima_train_network, generic_merge, generic_resize = trainer.import_anima_scripts()

for module in (sdxl_train_network, anima_train_network):
    parser = module.setup_parser()
    args = parser.parse_args([])
    args_util = trainer.training_args_util(module)
    args_util.verify_command_line_training_args(args)
    args_util.read_config_from_file(args, parser)

checks = [
    (sdxl_merge, 'merge'),
    (sdxl_resize, 'resize'),
    (generic_merge, 'merge'),
    (generic_resize, 'resize'),
]
for module, attr in checks:
    if not hasattr(module, attr):
        raise RuntimeError('%s has no %s' % (module.__name__, attr))

print('Copita sd-scripts integration: OK')
"@
    Invoke-NativeStep "Checking Copita sd-scripts integration" { & $Python -X utf8=0 -c $Smoke }
}

Write-Host ""
Write-Host "=========================================="
Write-Host " Copita - Setup"
Write-Host "=========================================="
Write-Host ""

if (Test-Path $Python) {
    $ExistingVersion = Get-CopitaPythonVersion
    if (!$ExistingVersion) {
        Write-Host "Existing Copita Python environment could not be checked. Recreating Copita\.venv..."
        Remove-Item -LiteralPath $Venv -Recurse -Force
    } elseif (!(Test-SupportedPythonVersion $ExistingVersion)) {
        Write-Host "Existing Copita Python environment is Python $ExistingVersion, which is not supported."
        Write-Host "Recreating Copita\.venv with Python 3.10 or 3.11..."
        Remove-Item -LiteralPath $Venv -Recurse -Force
    }
}

if (!(Test-Path $Python)) {
    $BootstrapPython = Resolve-BootstrapPython
    $BootstrapPythonExe = [string]$BootstrapPython["Exe"]
    $BootstrapPythonArgs = @($BootstrapPython["Args"])
    Write-Host ("Creating Python environment with {0} (Python {1})..." -f $BootstrapPython["Label"], $BootstrapPython["Version"])
    Invoke-NativeStep "Creating Python environment" { & $BootstrapPythonExe @BootstrapPythonArgs -m venv $Venv }
}

Assert-SupportedPythonVersion

Write-Host "Installing Python packages..."
Invoke-NativeStep "Upgrading pip" { & $Python -m pip install --upgrade pip wheel "setuptools<82" }
Remove-OnnxRuntimeGpuIfInstalled
if (!(Test-CudaTorchBuild)) {
    Install-CudaTorch
}
Invoke-NativeStep "Installing tool packages" { & $Python -m pip install -r (Join-Path $Root "requirements.txt") }

Update-SdScripts
Apply-SdScriptsCompatibilityPatch
Invoke-NativeStep "Checking Anima compatibility patch" { & $Python -m py_compile (Join-Path $SdScripts "library\anima_utils.py") }

if (Test-Path (Join-Path $SdScripts "requirements.txt")) {
    Write-Host "Installing sd-scripts packages..."
    Push-Location $SdScripts
    try {
        Invoke-NativeStep "Installing sd-scripts packages" { & $Python -m pip install -r "requirements.txt" }
    } finally {
        Pop-Location
    }
}

if (!(Test-CudaTorchBuild)) {
    Install-CudaTorch
}

Test-SdScriptsIntegration

if ($DownloadTagger) {
    Write-Host "Downloading caption tagger..."
    $TaggerDir = Join-Path $Root "models\tagger"
    New-Item -ItemType Directory -Force -Path $TaggerDir | Out-Null
    $DownloadScript = @"
import pathlib
import urllib.request

tagger_dir = pathlib.Path(r'''$TaggerDir''')
base_url = 'https://huggingface.co/SmilingWolf/wd-swinv2-tagger-v3/resolve/main'
files = ('model.onnx', 'selected_tags.csv', 'config.json', 'sw_jax_cv_config.json')

for name in files:
    out = tagger_dir / name
    if out.exists() and out.stat().st_size > 0:
        print(f'{name}: already exists')
        continue
    tmp = out.with_suffix(out.suffix + '.part')
    url = f'{base_url}/{name}'
    print(f'Downloading {name}...')
    request = urllib.request.Request(url, headers={'User-Agent': 'Copita setup'})
    with urllib.request.urlopen(request, timeout=60) as response, tmp.open('wb') as handle:
        while True:
            chunk = response.read(1024 * 1024)
            if not chunk:
                break
            handle.write(chunk)
    if tmp.stat().st_size <= 0:
        raise RuntimeError(f'Downloaded file is empty: {name}')
    tmp.replace(out)
    print(f'{name}: OK')
"@
    Invoke-NativeStep "Downloading caption tagger" { & $Python -c $DownloadScript }
}

Write-Host "Checking package consistency..."
Invoke-NativeStep "Checking package consistency" { & $Python -m pip check }

Write-Host ""
Write-Host "Setup finished."
Write-Host "Next: run start_copita.bat"
