param(
  [ValidateSet("default", "cpu", "cu121", "cu124", "cu126", "cu128")]
  [string]$TorchFlavor = "default",
  [switch]$Force
)

$ErrorActionPreference = "Stop"
$Root = Split-Path -Parent $MyInvocation.MyCommand.Path
$Venv = Join-Path $Root ".venv"
$Python = $null
$PythonArgs = @()

function Assert-LastExitCode {
  param([string]$Step)
  if ($LASTEXITCODE -ne 0) {
    throw "$Step failed with exit code $LASTEXITCODE"
  }
}

function Test-PythonCandidate {
  param(
    [string]$Exe,
    [string[]]$Args
  )
  if (!(Get-Command $Exe -ErrorAction SilentlyContinue)) {
    return $false
  }
  & $Exe @Args -c "import sys; raise SystemExit(0 if sys.version_info >= (3, 10) else 1)" *> $null
  return $LASTEXITCODE -eq 0
}

$Candidates = @(
  @{ Exe = "py"; Args = @("-3.12") },
  @{ Exe = "py"; Args = @("-3.11") },
  @{ Exe = "py"; Args = @("-3.10") },
  @{ Exe = "py"; Args = @("-3.13") },
  @{ Exe = "python"; Args = @() },
  @{ Exe = "python3"; Args = @() }
)

foreach ($Candidate in $Candidates) {
  if (Test-PythonCandidate -Exe $Candidate.Exe -Args $Candidate.Args) {
    $Python = $Candidate.Exe
    $PythonArgs = $Candidate.Args
    break
  }
}

if ($null -eq $Python) {
  throw "No working Python 3.10+ runtime was found. Check 'python --version' or install Python 3.11 from python.org with 'Add python.exe to PATH' enabled."
}

if ((Test-Path $Venv) -and $Force) {
  Remove-Item -Recurse -Force $Venv
}

if (!(Test-Path $Venv)) {
  & $Python @PythonArgs -m venv $Venv
  Assert-LastExitCode "Create virtual environment"
}

$VenvPython = Join-Path $Venv "Scripts\python.exe"
if (!(Test-Path $VenvPython)) {
  throw "Virtual environment creation failed. Tried: $Python $($PythonArgs -join ' ') -m venv $Venv"
}
& $VenvPython -m pip install --upgrade pip setuptools wheel
Assert-LastExitCode "Upgrade pip/setuptools/wheel"

switch ($TorchFlavor) {
  "cpu" {
    & $VenvPython -m pip install torch --index-url https://download.pytorch.org/whl/cpu
    Assert-LastExitCode "Install torch CPU"
  }
  "cu121" {
    & $VenvPython -m pip install torch --index-url https://download.pytorch.org/whl/cu121
    Assert-LastExitCode "Install torch cu121"
  }
  "cu124" {
    & $VenvPython -m pip install torch --index-url https://download.pytorch.org/whl/cu124
    Assert-LastExitCode "Install torch cu124"
  }
  "cu126" {
    & $VenvPython -m pip install torch --index-url https://download.pytorch.org/whl/cu126
    Assert-LastExitCode "Install torch cu126"
  }
  "cu128" {
    & $VenvPython -m pip install torch --index-url https://download.pytorch.org/whl/cu128
    Assert-LastExitCode "Install torch cu128"
  }
  default {
    & $VenvPython -m pip install torch
    Assert-LastExitCode "Install torch"
  }
}

& $VenvPython -m pip install `
  "numpy>=1.24.0" `
  "pydantic>=2.0.0" `
  "pydantic-settings>=2.0.0" `
  "transformers>=4.30.0" `
  "click>=8.1.0" `
  "rich>=13.0.0" `
  "sentencepiece>=0.1.99" `
  "tqdm>=4.65.0"
Assert-LastExitCode "Install runtime dependencies"

& $VenvPython -m pip install -e (Join-Path $Root "code\Taotern_SSM") --no-deps
Assert-LastExitCode "Install Taotern_SSM"
& $VenvPython -m pip install -e (Join-Path $Root "code\TaoTrain") --no-deps
Assert-LastExitCode "Install TaoTrain"

Write-Host ""
Write-Host "Setup complete."
Write-Host "Python: $VenvPython"
& $VenvPython -c "import torch; print('torch:', torch.__version__); print('cuda available:', torch.cuda.is_available()); print('cuda device:', torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'none')"
Write-Host ""
Write-Host "Run fixed chat with:"
Write-Host "  .\run_chat_fixed.ps1"