$ErrorActionPreference = "Stop"

Write-Host "=== Whisper Dictation Builder ===" -ForegroundColor Cyan
Write-Host ""

$ScriptDir = Split-Path -Parent $MyInvocation.MyCommand.Definition
$RepoRoot = Resolve-Path "$ScriptDir/.."
$DepsDir = "$RepoRoot/deps"
$BuildDir = "$RepoRoot/build"
$SdlVer = "2.28.5"
$SdlDirName = "SDL2-$SdlVer"
$SdlPath = "$DepsDir/$SdlDirName"

# Detect available GPU backends
$UseGPU = $false
$GPUBackend = "none"

Write-Host "[1/6] Detecting GPU capabilities..." -ForegroundColor Yellow

# Check for NVIDIA GPU
try {
    $NvidiaGpu = Get-WmiObject Win32_VideoController | Where-Object { $_.Name -like "*NVIDIA*" }
    if ($NvidiaGpu) {
        Write-Host "  [OK] NVIDIA GPU detected: $($NvidiaGpu.Name)" -ForegroundColor Green
        
        # Check for CUDA
        $CudaPath = $env:CUDA_PATH
        if (!$CudaPath) {
            # Try to find latest CUDA version if CUDA_PATH not set
            $cudaDir = "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA"
            if (Test-Path $cudaDir) {
                $versions = Get-ChildItem $cudaDir -Directory | Sort-Object Name -Descending
                if ($versions.Count -gt 0) {
                    $CudaPath = $versions[0].FullName
                    Write-Host "  [INFO] Found CUDA at: $CudaPath" -ForegroundColor Cyan
                }
            }
        }
        
        if ($CudaPath -and (Test-Path "$CudaPath\bin\nvcc.exe")) {
            # Check CUDA version
            $nvccOutput = (& "$CudaPath\bin\nvcc.exe" --version 2>&1) -join "`n"
            if ($nvccOutput -match "release\s+(\d+)\.(\d+)") {
                $cudaMajor = [int]$matches[1]
                $cudaMinor = [int]$matches[2]
                
                if ($cudaMajor -ge 12) {
                    Write-Host "  [OK] CUDA $cudaMajor.$cudaMinor toolkit - EXCELLENT!" -ForegroundColor Green
                    Write-Host "  [OK] Full GPU acceleration enabled" -ForegroundColor Green
                    $UseGPU = $true
                    $GPUBackend = "cuda"
                } elseif ($cudaMajor -eq 11 -and $cudaMinor -ge 7) {
                    Write-Host "  [OK] CUDA $cudaMajor.$cudaMinor toolkit found" -ForegroundColor Yellow
                    Write-Host "  [WARN] CUDA 12.4+ recommended for MSVC 2022, will try anyway" -ForegroundColor Yellow
                    $UseGPU = $true
                    $GPUBackend = "cuda"
                } else {
                    Write-Host "  [WARN] CUDA $cudaMajor.$cudaMinor is too old (need 11.7+)" -ForegroundColor Yellow
                }
            } else {
                Write-Host "  [WARN] Could not parse CUDA version from nvcc" -ForegroundColor Yellow
            }
        } else {
            Write-Host "  [WARN] CUDA toolkit not found or nvcc.exe missing" -ForegroundColor Yellow
            Write-Host "  [INFO] After installing CUDA 13.0, restart your terminal" -ForegroundColor Cyan
        }
    }
} catch {
    Write-Host "  [WARN] GPU detection error: $_" -ForegroundColor Yellow
    Write-Host "  [INFO] Falling back to CPU-only" -ForegroundColor Gray
}

# Check for AMD GPU
try {
    $AmdGpu = Get-WmiObject Win32_VideoController | Where-Object { $_.Name -like "*AMD*" -or $_.Name -like "*Radeon*" }
    if ($AmdGpu -and !$UseGPU) {
        Write-Host "  [OK] AMD GPU detected: $($AmdGpu.Name)" -ForegroundColor Green
        Write-Host "  [INFO] ROCm support available but requires manual setup" -ForegroundColor Cyan
    }
} catch {
    # Ignore
}

# Check for Vulkan
if (!$UseGPU) {
    $VulkanSDK = $env:VULKAN_SDK
    if ($VulkanSDK -and (Test-Path "$VulkanSDK")) {
        $UseGPU = $true
        $GPUBackend = "vulkan"
        Write-Host "  [OK] Vulkan SDK found at: $VulkanSDK" -ForegroundColor Green
    }
}

if (!$UseGPU) {
    Write-Host "  -> Building with optimized multi-core CPU" -ForegroundColor Cyan
    Write-Host "     Using all $env:NUMBER_OF_PROCESSORS CPU threads for maximum performance" -ForegroundColor Gray
    Write-Host "     (To enable GPU: Install CUDA 12.4+ or Vulkan SDK)" -ForegroundColor DarkGray
}

Write-Host ""

# Setup SDL2
Write-Host "[2/6] Setting up SDL2..." -ForegroundColor Yellow
if (-not (Test-Path "$SdlPath/cmake/sdl2-config.cmake")) {
    if (-not (Test-Path $DepsDir)) { New-Item -ItemType Directory -Path $DepsDir | Out-Null }
    $ZipFile = "$DepsDir/SDL2-devel-$SdlVer-VC.zip"
    $Url = "https://github.com/libsdl-org/SDL/releases/download/release-$SdlVer/SDL2-devel-$SdlVer-VC.zip"
    
    Write-Host "  Downloading SDL2..." -NoNewline
    if (-not (Test-Path $ZipFile)) { 
        Invoke-WebRequest -Uri $Url -OutFile $ZipFile 
    }
    Write-Host " Done" -ForegroundColor Green
    
    Write-Host "  Extracting..." -NoNewline
    Expand-Archive -Path $ZipFile -DestinationPath $DepsDir -Force
    Write-Host " Done" -ForegroundColor Green
} else {
    Write-Host "  [OK] SDL2 already configured" -ForegroundColor Green
}

Write-Host ""

# Configure CMake
Write-Host "[3/6] Configuring CMake..." -ForegroundColor Yellow
Set-Location $RepoRoot

$SdlPathResolved = Resolve-Path $SdlPath -ErrorAction SilentlyContinue
if (-not $SdlPathResolved) {
    $SdlPathResolved = $SdlPath
}
# Convert to forward slashes for CMake
$SdlDir = ($SdlPathResolved -replace '\\', '/') + '/cmake'

$CMakeArgs = @(
    "-S", $RepoRoot,
    "-B", "build",
    "-DWHISPER_SDL2=ON",
    "-DSDL2_DIR=$SdlDir"
)

# Add GPU backend flags
if ($UseGPU) {
    if ($GPUBackend -eq "cuda") {
        Write-Host "  Enabling CUDA backend..." -ForegroundColor Cyan
        $CMakeArgs += "-DGGML_CUDA=ON"
        # Explicitly set CUDA_PATH to avoid finding old versions
        if ($CudaPath) {
            $env:CUDA_PATH = $CudaPath
            $env:CUDA_HOME = $CudaPath
            # Force CMake to use this specific toolkit root
            $CMakeArgs += "-DCUDAToolkit_ROOT=$CudaPath"
            Write-Host "  Setting CUDA_PATH to: $CudaPath" -ForegroundColor Cyan
        }
    } elseif ($GPUBackend -eq "vulkan") {
        Write-Host "  Enabling Vulkan backend..." -ForegroundColor Cyan
        $CMakeArgs += "-DGGML_VULKAN=ON"
    }
} else {
    Write-Host "  Building CPU-only with all cores..." -ForegroundColor Cyan
    # Explicitly disable all GPU backends
    $CMakeArgs += "-DGGML_CUDA=OFF"
    $CMakeArgs += "-DGGML_VULKAN=OFF"
    $CMakeArgs += "-DGGML_METAL=OFF"
    $CMakeArgs += "-DGGML_HIPBLAS=OFF"
}

# Run CMake
& cmake @CMakeArgs

if ($LASTEXITCODE -ne 0) {
    Write-Host ""
    Write-Host "[WARN] CMake configuration had warnings, but continuing..." -ForegroundColor Yellow
    if ($UseGPU) {
        Write-Host "  GPU backend may not be available. Will use CPU fallback." -ForegroundColor Yellow
    }
}

Write-Host ""

# Build
Write-Host "[4/6] Building win-dictation..." -ForegroundColor Yellow
cmake --build build --config Release --target win-dictation -j $env:NUMBER_OF_PROCESSORS

if ($LASTEXITCODE -ne 0) {
    Write-Host ""
    Write-Host "[ERROR] Build failed!" -ForegroundColor Red
    exit 1
}

Write-Host "  [OK] Build successful" -ForegroundColor Green
Write-Host ""

# Deploy DLLs
Write-Host "[5/6] Deploying dependencies..." -ForegroundColor Yellow
$BinDir = "$BuildDir/bin/Release"

# Copy SDL2
$SdlDll = "$SdlPath/lib/x64/SDL2.dll"
if (Test-Path $SdlDll) {
    Copy-Item -Path $SdlDll -Destination $BinDir -Force
    Write-Host "  [OK] Copied SDL2.dll" -ForegroundColor Green
}

# Copy CUDA DLLs if needed
if ($UseGPU -and $GPUBackend -eq "cuda") {
    Write-Host "  Deploying CUDA runtime DLLs..." -ForegroundColor Cyan
    
    # Find CUDA installation
    $cudaPath = $env:CUDA_PATH
    if (!$cudaPath) {
        $cudaDir = "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA"
        $versions = Get-ChildItem $cudaDir -Directory -ErrorAction SilentlyContinue | Sort-Object Name -Descending
        if ($versions.Count -gt 0) {
            $cudaPath = $versions[0].FullName
        }
    }
    
    if ($cudaPath) {
        # CUDA 13.0 / 12.x DLLs
        $CudaDlls = @("cudart64_*.dll", "cublas64_*.dll", "cublasLt64_*.dll", "cudnn64_*.dll", "cufft64_*.dll")
        foreach ($pattern in $CudaDlls) {
            # Check bin root
            $dll = Get-ChildItem "$cudaPath\bin\$pattern" -ErrorAction SilentlyContinue | Select-Object -First 1
            
            # Check bin/x64 (CUDA 13 layout)
            if (!$dll) {
                $dll = Get-ChildItem "$cudaPath\bin\x64\$pattern" -ErrorAction SilentlyContinue | Select-Object -First 1
            }

            if ($dll) {
                Copy-Item -Path $dll.FullName -Destination $BinDir -Force
                Write-Host "  [OK] Copied $($dll.Name)" -ForegroundColor Green
            }
        }
        
        # Also copy from system if available (for CUDA 13)
        $systemCuda = "C:\Windows\System32\cudart64_*.dll"
        $sysDll = Get-ChildItem $systemCuda -ErrorAction SilentlyContinue | Select-Object -First 1
        if ($sysDll -and !(Test-Path "$BinDir\cudart64_*.dll")) {
            Copy-Item -Path $sysDll.FullName -Destination $BinDir -Force
            Write-Host "  [OK] Copied system $($sysDll.Name)" -ForegroundColor Green
        }
    } else {
        Write-Host "  [WARN] Could not find CUDA installation for DLL deployment" -ForegroundColor Yellow
    }
}

# Copy ggml DLLs
$GgmlDlls = Get-ChildItem "$BuildDir" -Recurse -Filter "ggml*.dll" -ErrorAction SilentlyContinue
foreach ($dll in $GgmlDlls) {
    if ($dll.FullName -like "*Release*" -or $dll.FullName -like "*bin*") {
        Copy-Item -Path $dll.FullName -Destination $BinDir -Force -ErrorAction SilentlyContinue
    }
}

Write-Host ""

# Download models
Write-Host "[6/6] Checking Whisper models..." -ForegroundColor Yellow

if (-not (Test-Path "$BinDir/models")) { 
    New-Item -ItemType Directory -Path "$BinDir/models" | Out-Null 
}

# Download tiny.en for CPU-only systems (faster, smaller)
$TinyModelName = "ggml-tiny.en.bin"
$TinyModelSrc = "$RepoRoot/models/$TinyModelName"
$TinyModelDest = "$BinDir/models/$TinyModelName"

if (-not (Test-Path $TinyModelSrc)) {
    Write-Host "  Downloading tiny.en model (for CPU-only systems)..." -ForegroundColor Cyan
    
    $DownloadScript = "$RepoRoot/models/download-ggml-model.sh"
    if (Test-Path $DownloadScript) {
        # Try bash if available
        $bash = Get-Command bash -ErrorAction SilentlyContinue
        if ($bash) {
            & bash $DownloadScript tiny.en
        } else {
            # Direct download
            $TinyModelUrl = "https://huggingface.co/ggerganov/whisper.cpp/resolve/main/ggml-tiny.en.bin"
            Invoke-WebRequest -Uri $TinyModelUrl -OutFile $TinyModelSrc
        }
    }
}

if (Test-Path $TinyModelSrc) { 
    Copy-Item -Path $TinyModelSrc -Destination $TinyModelDest -Force
    Write-Host "  [OK] tiny.en model ready" -ForegroundColor Green
} else {
    Write-Host "  [WARN] Could not download tiny.en model" -ForegroundColor Yellow
}

# Download base.en for GPU systems (better accuracy)
$BaseModelName = "ggml-base.en.bin"
$BaseModelSrc = "$RepoRoot/models/$BaseModelName"
$BaseModelDest = "$BinDir/models/$BaseModelName"

if (-not (Test-Path $BaseModelSrc)) {
    Write-Host "  Downloading base.en model (for GPU systems)..." -ForegroundColor Cyan
    
    $DownloadScript = "$RepoRoot/models/download-ggml-model.sh"
    if (Test-Path $DownloadScript) {
        # Try bash if available
        $bash = Get-Command bash -ErrorAction SilentlyContinue
        if ($bash) {
            & bash $DownloadScript base.en
        } else {
            # Direct download
            $BaseModelUrl = "https://huggingface.co/ggerganov/whisper.cpp/resolve/main/ggml-base.en.bin"
            Invoke-WebRequest -Uri $BaseModelUrl -OutFile $BaseModelSrc
        }
    }
}

if (Test-Path $BaseModelSrc) { 
    Copy-Item -Path $BaseModelSrc -Destination $BaseModelDest -Force
    Write-Host "  [OK] base.en model ready" -ForegroundColor Green
} else {
    Write-Host "  [WARN] Could not download base.en model" -ForegroundColor Yellow
    Write-Host "  Please download manually from:" -ForegroundColor Yellow
    Write-Host "  https://huggingface.co/ggerganov/whisper.cpp/resolve/main/ggml-base.en.bin" -ForegroundColor Gray
    Write-Host "  And place it in: $BaseModelDest" -ForegroundColor Gray
}

Write-Host "  [INFO] App will auto-select optimal model based on GPU availability" -ForegroundColor Cyan

Write-Host ""
Write-Host "=== Build Complete ===" -ForegroundColor Green
Write-Host ""
Write-Host "Configuration:" -ForegroundColor Cyan
Write-Host "  - GPU Backend: " -NoNewline
if ($UseGPU) {
    Write-Host "$GPUBackend" -ForegroundColor Green
} else {
    Write-Host "CPU-only" -ForegroundColor Yellow
}
Write-Host "  - CPU Threads: $env:NUMBER_OF_PROCESSORS" -ForegroundColor Cyan
Write-Host "  - Model: base.en" -ForegroundColor Cyan
Write-Host ""
Write-Host "Run the application:" -ForegroundColor White
Write-Host "  $BinDir\win-dictation.exe" -ForegroundColor Yellow
Write-Host ""
Write-Host "Hotkey: Ctrl+Shift+R to toggle recording" -ForegroundColor Gray
Write-Host ""