# ============================================================
# Windows-Side Setup - Local LLM Stack (Razer Blade 18)
#
# Run this FIRST from PowerShell (Admin) on a fresh Windows machine.
# After this completes, open WSL2 and run setup-wsl.sh.
#
# What this does:
#   1. Verifies NVIDIA drivers + CUDA
#   2. Installs Ollama via winget
#   3. Pulls all 5 LLM models
#   4. Installs WSL2 with Ubuntu 24.04
#
# Usage (PowerShell as Admin):
#   Set-ExecutionPolicy -Scope Process Bypass
#   .\setup-windows.ps1
#
# After reboot, open Ubuntu (WSL2) and run:
#   bash ~/setup-wsl.sh
# ============================================================

$ErrorActionPreference = "Stop"

function Write-OK    { param($msg) Write-Host "  [OK] $msg" -ForegroundColor Green }
function Write-Warn  { param($msg) Write-Host "  [!!] $msg" -ForegroundColor Yellow }
function Write-Fail  { param($msg) Write-Host "  [FAIL] $msg" -ForegroundColor Red }
function Write-Step  { param($msg) Write-Host "`n=== $msg ===" -ForegroundColor Cyan }

Write-Host ""
Write-Host "  =====================================" -ForegroundColor Cyan
Write-Host "  Local LLM Stack - Windows Setup" -ForegroundColor Cyan
Write-Host "  Razer Blade 18 / RTX 5090" -ForegroundColor Cyan
Write-Host "  =====================================" -ForegroundColor Cyan
Write-Host ""

# -- 1. Check NVIDIA Drivers ----------------------------------
Write-Step "Step 1/4: NVIDIA Drivers + CUDA"

$nvidiaSmi = Get-Command nvidia-smi -ErrorAction SilentlyContinue
if ($nvidiaSmi) {
    $gpuInfo = & nvidia-smi --query-gpu=name,driver_version,memory.total --format=csv,noheader 2>$null
    if ($gpuInfo) {
        Write-OK "NVIDIA driver installed"
        Write-Host "  GPU: $($gpuInfo.Trim())"
    } else {
        Write-Warn "nvidia-smi found but returned no data"
    }
} else {
    Write-Fail "NVIDIA drivers not found!"
    Write-Host "  Download from: https://www.nvidia.com/Download/index.aspx"
    Write-Host "  Install Game Ready or Studio drivers, then re-run this script."
    Write-Host ""
    $continue = Read-Host "  Continue anyway? (y/N)"
    if ($continue -ne "y") { exit 1 }
}

# -- 2. Install Ollama -----------------------------------------
Write-Step "Step 2/4: Ollama"

$ollamaCmd = Get-Command ollama -ErrorAction SilentlyContinue
if ($ollamaCmd) {
    $ollamaVer = & ollama --version 2>&1
    Write-OK "Ollama already installed ($ollamaVer)"
} else {
    Write-Host "  Installing Ollama via winget..."
    try {
        & winget install --id Ollama.Ollama --accept-source-agreements --accept-package-agreements
        Write-OK "Ollama installed"
        Write-Warn "You may need to restart your terminal for 'ollama' to be on PATH"
    } catch {
        Write-Fail "winget install failed. Download manually: https://ollama.com/download"
        $continue = Read-Host "  Continue anyway? (y/N)"
        if ($continue -ne "y") { exit 1 }
    }
}

# Verify Ollama is running
Write-Host "  Checking if Ollama is running..."
try {
    $response = Invoke-WebRequest -Uri "http://localhost:11434/api/tags" -TimeoutSec 3 -UseBasicParsing -ErrorAction Stop
    Write-OK "Ollama is running on port 11434"
} catch {
    Write-Warn "Ollama not running. Starting..."
    Start-Process ollama -ArgumentList "serve" -WindowStyle Hidden
    Start-Sleep -Seconds 3
    try {
        $response = Invoke-WebRequest -Uri "http://localhost:11434/api/tags" -TimeoutSec 3 -UseBasicParsing -ErrorAction Stop
        Write-OK "Ollama started"
    } catch {
        Write-Fail "Could not start Ollama. Start manually: ollama serve"
        Write-Host "  Then re-run this script."
        $continue = Read-Host "  Continue anyway? (y/N)"
        if ($continue -ne "y") { exit 1 }
    }
}

# -- 3. Pull Models --------------------------------------------
Write-Step "Step 3/4: Pull Ollama Models"
# Curated for coding + reasoning + TTS; update from https://ollama.com/models as needed
$models = @(
    @{ name = "qwen2.5-coder:32b";  desc = "19 GB - primary coding model" },
    @{ name = "qwen2.5-coder:7b";   desc = "4.7 GB - fast coding" },
    @{ name = "deepseek-r1:32b";     desc = "19 GB - chain-of-thought" },
    @{ name = "llama3.1:8b";         desc = "4.9 GB - fast general tasks" },
    @{ name = "sematre/orpheus:en";  desc = "4 GB - text-to-speech (8 voices)" }
)

# Check which models are already pulled
$existingModels = @()
try {
    $tagsJson = Invoke-WebRequest -Uri "http://localhost:11434/api/tags" -TimeoutSec 5 -UseBasicParsing -ErrorAction Stop
    $tagsObj = $tagsJson.Content | ConvertFrom-Json
    $existingModels = $tagsObj.models | ForEach-Object { $_.name }
} catch {
    Write-Warn "Could not check existing models (Ollama may not be running)"
}

foreach ($model in $models) {
    $alreadyPulled = $existingModels | Where-Object { $_ -like "*$($model.name)*" }
    if ($alreadyPulled) {
        Write-OK "$($model.name) - already pulled"
    } else {
        Write-Host "  Pulling $($model.name) ($($model.desc))..."
        & ollama pull $model.name
        if ($LASTEXITCODE -eq 0) {
            Write-OK "$($model.name) - pulled"
        } else {
            Write-Warn "$($model.name) - pull failed (you can retry later: ollama pull $($model.name))"
        }
    }
}

# Verify
Write-Host ""
Write-Host "  Installed models:"
& ollama list

# -- 4. Install WSL2 -------------------------------------------
Write-Step "Step 4/4: WSL2 (Ubuntu 24.04)"

$wslInstalled = $false
try {
    $wslList = & wsl --list --quiet 2>&1
    if ($wslList -match "Ubuntu") {
        Write-OK "WSL2 Ubuntu already installed"
        $wslInstalled = $true
    }
} catch {
    # wsl not available
}

if (-not $wslInstalled) {
    Write-Host "  Installing WSL2 with Ubuntu 24.04..."
    Write-Host "  This may require a reboot."
    Write-Host ""
    & wsl --install -d Ubuntu-24.04
    Write-Host ""
    Write-Warn "If prompted, REBOOT your machine now."
    Write-Warn "After reboot, Ubuntu will ask you to set up a username/password."
}

# -- Summary ---------------------------------------------------
Write-Host ''
Write-Host '  =====================================' -ForegroundColor Green
Write-Host '  Windows-Side Setup Complete!' -ForegroundColor Green
Write-Host '  =====================================' -ForegroundColor Green
Write-Host ''
Write-Host '  Next steps:' -ForegroundColor Yellow
Write-Host '  1. If WSL2 was just installed, reboot and set up Ubuntu username/password'
Write-Host '  2. Open Ubuntu (WSL2) terminal'
Write-Host '  3. Run the WSL2 setup script:'
Write-Host ''
Write-Host '     curl -fsSL https://raw.githubusercontent.com/saravanakumardb1/learning_ai_common_plat/main/__LOCAL_LLMs/windows_specific/setup-wsl.sh | bash' -ForegroundColor Cyan
Write-Host ''
Write-Host '     Or if you already cloned the repo:'
Write-Host '     bash ~/code/mygh/learning_ai_common_plat/__LOCAL_LLMs/windows_specific/setup-wsl.sh' -ForegroundColor Cyan
Write-Host ''