param(
  [string]$VenvDir = ".venv",
  [string]$PythonExe = "python",
  [string]$PipExe = ""
)

$ErrorActionPreference = "Stop"

if (-not (Test-Path $VenvDir)) {
  & $PythonExe -m venv $VenvDir
}

if ([string]::IsNullOrWhiteSpace($PipExe)) {
  $pipCandidate = Join-Path $VenvDir "Scripts\\pip.exe"
  if (Test-Path $pipCandidate) {
    $PipExe = $pipCandidate
  } else {
    $PipExe = "pip"
  }
}

$cudaVariant = $env:RAGGO_TORCH_CUDA
if ([string]::IsNullOrWhiteSpace($cudaVariant)) {
  $nvidiaSmi = Get-Command "nvidia-smi" -ErrorAction SilentlyContinue
  if ($null -ne $nvidiaSmi) {
    $cudaVersion = (& nvidia-smi --query-gpu=cuda_version --format=csv,noheader 2>$null | Select-Object -First 1).Trim()
    if (-not [string]::IsNullOrWhiteSpace($cudaVersion)) {
      $parts = $cudaVersion.Split(".")
      $major = [int]$parts[0]
      $minor = if ($parts.Length -gt 1) { [int]$parts[1] } else { 0 }
      if ($major -ge 12) {
        $cudaVariant = "cu121"
      } elseif ($major -eq 11 -and $minor -ge 8) {
        $cudaVariant = "cu118"
      }
    }
  }
}

if (-not [string]::IsNullOrWhiteSpace($cudaVariant)) {
  Write-Host "Detected CUDA; installing torch for $cudaVariant"
  & $PipExe install --index-url "https://download.pytorch.org/whl/$cudaVariant" torch
} else {
  Write-Host "CUDA not detected; installing CPU torch via default index"
  & $PipExe install torch
}

& $PipExe install -r python_service/requirements.txt
