param(
  [switch]$Cuda,
  [string]$CudaToolkitRoot = "",
  [switch]$Clean,
  [switch]$Package
)

$ErrorActionPreference = "Stop"
$pluginRoot = Split-Path -Parent $PSScriptRoot
$sourceRoot = Join-Path $pluginRoot "backend\transcription\worker-native"
$backend = if ($Cuda) { "cuda" } else { "cpu" }
$buildRoot = Join-Path $env:TEMP "lumi-whisper-worker-$backend"
$targetRoot = Join-Path $pluginRoot "data\runtime\worker\windows-x64-$backend\bin"

function Find-Tool([string]$name) {
  $command = Get-Command $name -ErrorAction SilentlyContinue
  if ($command) { return $command.Source }
  $candidate = Get-ChildItem (Join-Path $env:APPDATA "Python\Python*\Scripts\$name") -ErrorAction SilentlyContinue |
    Sort-Object FullName -Descending |
    Select-Object -First 1
  if ($candidate) { return $candidate.FullName }
  throw "$name was not found. Install CMake and Ninja before building the transcription worker."
}

$cmake = Find-Tool "cmake.exe"
$ninja = Find-Tool "ninja.exe"
$cudaRoot = if ($Cuda -and $CudaToolkitRoot) {
  Get-Item -LiteralPath $CudaToolkitRoot -ErrorAction Stop
} elseif ($Cuda) {
  Get-ChildItem "$env:ProgramFiles\NVIDIA GPU Computing Toolkit\CUDA\v*" -Directory -ErrorAction SilentlyContinue |
    Where-Object { Test-Path (Join-Path $_.FullName "bin\nvcc.exe") } |
    Sort-Object FullName -Descending |
    Select-Object -First 1
} else { $null }
if ($Cuda -and !$cudaRoot) { throw "The CUDA Toolkit was not found. Install it before building the CUDA worker." }
if ($Cuda -and !(Test-Path (Join-Path $cudaRoot.FullName "bin\nvcc.exe"))) {
  throw "The selected CUDA Toolkit does not contain bin\\nvcc.exe."
}
$cudaToolRoot = $null
if ($cudaRoot) {
  $cudaToolRoot = Join-Path $env:TEMP "lumi-cuda-sdk"
  if (Test-Path $cudaToolRoot) { Remove-Item $cudaToolRoot -Force -Recurse }
  New-Item -ItemType Junction -Path $cudaToolRoot -Target $cudaRoot.FullName | Out-Null
}
$vcvars = Get-ChildItem "${env:ProgramFiles(x86)}\Microsoft Visual Studio\2022\*\VC\Auxiliary\Build\vcvars64.bat" -ErrorAction SilentlyContinue |
  Select-Object -First 1
if (!$vcvars) { throw "Visual Studio 2022 C++ Build Tools were not found." }
$windowsSdkBin = Get-ChildItem "${env:ProgramFiles(x86)}\Windows Kits\10\bin\*\x64" -Directory -ErrorAction SilentlyContinue |
  Where-Object { (Test-Path (Join-Path $_.FullName "rc.exe")) -and (Test-Path (Join-Path $_.FullName "mt.exe")) } |
  Sort-Object FullName -Descending |
  Select-Object -First 1
if (!$windowsSdkBin) { throw "The Windows SDK resource and manifest tools were not found." }
$resourceCompiler = (Join-Path $windowsSdkBin.FullName "rc.exe").Replace("\", "/")
$manifestTool = (Join-Path $windowsSdkBin.FullName "mt.exe").Replace("\", "/")
$windowsSdkVersion = $windowsSdkBin.Parent.Name
$windowsSdkRoot = Join-Path "${env:ProgramFiles(x86)}" "Windows Kits\10"
$windowsSdkInclude = Join-Path $windowsSdkRoot "Include\$windowsSdkVersion"
$windowsSdkLib = Join-Path $windowsSdkRoot "Lib\$windowsSdkVersion"

if ($Clean -and (Test-Path $buildRoot)) { Remove-Item $buildRoot -Force -Recurse }
New-Item -ItemType Directory -Force -Path $buildRoot, $targetRoot | Out-Null
$cudaValue = if ($Cuda) { "ON" } else { "OFF" }
$cudaArgument = if ($cudaRoot) {
  '-DCUDAToolkit_ROOT="{0}" -DCMAKE_CUDA_COMPILER="{1}"' -f $cudaToolRoot.Replace("\", "/"), (Join-Path $cudaToolRoot "bin\nvcc.exe").Replace("\", "/")
} else { "" }
$buildScript = Join-Path $buildRoot "build-worker.cmd"
@(
  "@echo off",
  ('call "{0}" >nul' -f $vcvars.FullName),
  'set "PATHEXT=.COM;.EXE;.BAT;.CMD"',
  ('set "PATH={0};%PATH%"' -f $windowsSdkBin.FullName),
  $(if ($cudaRoot) { 'set "PATH={0};%PATH%"' -f (Join-Path $cudaToolRoot "bin") } else { "rem CUDA is disabled" }),
  ('set "INCLUDE=%INCLUDE%;{0};{1};{2};{3}"' -f (Join-Path $windowsSdkInclude "ucrt"), (Join-Path $windowsSdkInclude "shared"), (Join-Path $windowsSdkInclude "um"), (Join-Path $windowsSdkInclude "winrt")),
  ('set "LIB=%LIB%;{0};{1}"' -f (Join-Path $windowsSdkLib "ucrt\x64"), (Join-Path $windowsSdkLib "um\x64")),
  ('"{0}" -S "{1}" -B "{2}" -G Ninja -DLUMI_WHISPER_CUDA={3} -DCMAKE_BUILD_TYPE=Release -DCMAKE_MAKE_PROGRAM="{4}" -DCMAKE_RC_COMPILER="{5}" -DCMAKE_MT="{6}" {7}' -f $cmake, $sourceRoot, $buildRoot, $cudaValue, $ninja, $resourceCompiler, $manifestTool, $cudaArgument),
  "if errorlevel 1 exit /b %errorlevel%",
  ('"{0}" --build "{1}" --config Release' -f $cmake, $buildRoot),
  "exit /b %errorlevel%"
) | Set-Content -Encoding Ascii $buildScript
$commandProcessor = Join-Path $env:SystemRoot "System32\cmd.exe"
$process = Start-Process -FilePath $commandProcessor -ArgumentList "/d /c $buildScript" -Wait -PassThru -NoNewWindow
if ($process.ExitCode -ne 0) { throw "Worker build failed with exit code $($process.ExitCode)." }

$worker = Join-Path $buildRoot "lumi-whisper-worker.exe"
if (!(Test-Path $worker)) { throw "The build completed without producing lumi-whisper-worker.exe." }
Copy-Item -Force $worker (Join-Path $targetRoot "lumi-whisper-worker.exe")
if ($Cuda) {
  Get-ChildItem $targetRoot -Filter "*.dll" -ErrorAction SilentlyContinue | Remove-Item -Force
  foreach ($pattern in @("cudart64_*.dll", "cublas64_*.dll", "cublasLt64_*.dll")) {
    $dependency = Get-ChildItem (Join-Path $cudaRoot.FullName "bin") -Filter $pattern -File |
      Sort-Object Name |
      Select-Object -First 1
    if (!$dependency) { throw "The selected CUDA Toolkit is missing $pattern." }
    Copy-Item -Force $dependency.FullName $targetRoot
  }
}
Write-Host "Installed Lumi whisper worker ($backend) in $targetRoot"
if ($Package) {
  $archive = Join-Path $pluginRoot "data\tmp\lumi-whisper-worker-windows-x64-$backend.zip"
  New-Item -ItemType Directory -Force -Path (Split-Path -Parent $archive) | Out-Null
  Remove-Item $archive -Force -ErrorAction SilentlyContinue
  Compress-Archive -Path $targetRoot -DestinationPath $archive -CompressionLevel Optimal
  $file = Get-Item $archive
  $sha = (Get-FileHash $archive -Algorithm SHA256).Hash.ToLowerInvariant()
  Write-Host "Runtime package: $archive"
  Write-Host "Runtime bytes: $($file.Length)"
  Write-Host "Runtime SHA256: $sha"
}