#Requires -Version 5.1
# mistral.rs Installation Script for Windows
# Automatic hardware detection and feature configuration

$ErrorActionPreference = "Stop"

# Color output functions
function Write-Info { Write-Host "info: $args" -ForegroundColor Blue }
function Write-Success { Write-Host "success: $args" -ForegroundColor Green }
function Write-Warn { Write-Host "warning: $args" -ForegroundColor Yellow }
function Write-Err { Write-Host "error: $args" -ForegroundColor Red; exit 1 }

# MISTRALRS_INSTALL_YES=1 auto-confirms every prompt (non-interactive installs, `mistralrs update`).
function Read-Confirm($prompt) {
    if ($env:MISTRALRS_INSTALL_YES -eq "1") { return "y" }
    return Read-Host $prompt
}

# Banner
function Show-Banner {
    Write-Host ""
    Write-Host "  __  __ _     _             _              " -ForegroundColor Cyan
    Write-Host " |  \/  (_)___| |_ _ __ __ _| |  _ __ ___   " -ForegroundColor Cyan
    Write-Host " | |\/| | / __| __| '__/ `` | | | '__/ __|  " -ForegroundColor Cyan
    Write-Host " | |  | | \__ \ |_| | | (_| | |_| |  \__ \  " -ForegroundColor Cyan
    Write-Host " |_|  |_|_|___/\__|_|  \__,_|_(_)_|  |___/  " -ForegroundColor Cyan
    Write-Host ""
    Write-Host "Fast, flexible LLM inference." -ForegroundColor Blue
    Write-Host ""
}

# Minimum required Rust version (from Cargo.toml rust-version)
$RequiredRustVersion = "1.94"
$MistralRsRepoUrl = "https://github.com/EricLBuehler/mistral.rs"
$MistralRsBranch = "master"
$MistralRsCliPackage = "mistralrs-cli"

# Check if Rust is installed
function Test-Rust {
    try {
        $null = Get-Command cargo -ErrorAction Stop
        return $true
    } catch {
        return $false
    }
}

# Get installed Rust version (major.minor)
function Get-RustVersion {
    try {
        $output = & rustc --version 2>$null
        if ($output -match 'rustc (\d+\.\d+)') {
            return $matches[1]
        }
    } catch {}
    return $null
}

# Compare two version strings (returns $true if $v1 >= $v2)
function Test-VersionGte {
    param([string]$v1, [string]$v2)

    $v1Parts = $v1 -split '\.'
    $v2Parts = $v2 -split '\.'

    $v1Major = [int]$v1Parts[0]
    $v1Minor = [int]$v1Parts[1]
    $v2Major = [int]$v2Parts[0]
    $v2Minor = [int]$v2Parts[1]

    if ($v1Major -gt $v2Major) {
        return $true
    } elseif ($v1Major -eq $v2Major -and $v1Minor -ge $v2Minor) {
        return $true
    }
    return $false
}

# Install Rust via rustup
function Install-Rust {
    Write-Info "Installing Rust via rustup..."
    $rustupInit = "$env:TEMP\rustup-init.exe"

    try {
        Invoke-WebRequest -Uri "https://win.rustup.rs/x86_64" -OutFile $rustupInit -UseBasicParsing
        & $rustupInit -y

        # Add cargo to PATH for current session
        $env:PATH = "$env:USERPROFILE\.cargo\bin;$env:PATH"

        Write-Success "Rust installed successfully"
    } catch {
        Write-Err "Failed to install Rust: $_"
    }
}

# Update Rust to latest version
function Update-Rust {
    Write-Info "Updating Rust to latest version..."
    try {
        & rustup update stable
        Write-Success "Rust updated successfully"
    } catch {
        Write-Err "Failed to update Rust: $_"
    }
}

# Get CUDA compute capability
function Get-CudaComputeCap {
    try {
        $nvidiaSmi = Get-Command nvidia-smi -ErrorAction Stop
        $output = & nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>$null
        if ($output) {
            $cc = ($output -split "`n")[0].Trim() -replace '\.',''
            return $cc
        }
    } catch {}
    return $null
}

# CUDA toolkit version as major*100+minor (e.g. 13.1 -> 1301). $null if nvcc is unavailable.
function Get-CudaVersionCode {
    try {
        $null = Get-Command nvcc -ErrorAction Stop
        $output = & nvcc --version 2>$null | Out-String
        if ($output -match "release (\d+)\.(\d+)") {
            return [int]$Matches[1] * 100 + [int]$Matches[2]
        }
    } catch {}
    return $null
}

# Check if MKL is installed
function Test-MKL {
    if ($env:MKLROOT -and (Test-Path $env:MKLROOT)) {
        return $true
    }

    $mklPaths = @(
        "C:\Program Files (x86)\Intel\oneAPI\mkl\latest",
        "C:\Program Files\Intel\oneAPI\mkl\latest",
        "$env:USERPROFILE\intel\oneapi\mkl\latest"
    )

    foreach ($path in $mklPaths) {
        if (Test-Path $path) {
            return $true
        }
    }
    return $false
}

# Check if CPU is Intel
function Test-IntelCpu {
    try {
        $cpu = Get-CimInstance -ClassName Win32_Processor | Select-Object -First 1
        return $cpu.Manufacturer -match "Intel" -or $cpu.Name -match "Intel"
    } catch {
        return $false
    }
}

# Check if cuDNN is installed
function Test-CuDNN {
    # Check common cuDNN library paths on Windows
    $cudnnPaths = @(
        "$env:CUDA_PATH\bin\cudnn*.dll",
        "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\*\bin\cudnn*.dll",
        "C:\Program Files\NVIDIA\CUDNN\*\bin\cudnn*.dll"
        "C:\Program Files\NVIDIA\CUDNN\*\bin\*\x64\cudnn*.dll"
    )

    foreach ($pattern in $cudnnPaths) {
        if (Get-Item $pattern -ErrorAction SilentlyContinue) {
            return $true
        }
    }
    return $false
}

# Check if NCCL is installed
function Test-NCCL {
    $ncclPaths = @(
        "$env:NCCL_ROOT\bin\nccl*.dll",
        "$env:NCCL_ROOT\lib\nccl*.lib",
        "$env:NCCL_HOME\bin\nccl*.dll",
        "$env:NCCL_HOME\lib\nccl*.lib",
        "$env:CUDA_PATH\bin\nccl*.dll",
        "$env:CUDA_PATH\lib\x64\nccl*.lib",
        "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\*\bin\nccl*.dll",
        "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\*\lib\x64\nccl*.lib"
    )

    foreach ($pattern in $ncclPaths) {
        if (Get-Item $pattern -ErrorAction SilentlyContinue) {
            return $true
        }
    }
    return $false
}

# Build feature string based on detected hardware
function Get-Features {
    $features = @()

    # Check for CUDA
    $cudaCC = Get-CudaComputeCap
    if ($cudaCC) {
        $features += "cuda"

        $ccMajor = $cudaCC.Substring(0, 1)
        $ccMinor = if ($cudaCC.Length -gt 1) { $cudaCC.Substring(1) } else { "0" }
        Write-Info "CUDA detected (compute capability: $ccMajor.$ccMinor)"

        if ($env:MISTRALRS_INSTALL_NO_NCCL -eq "1") {
            Write-Info "MISTRALRS_INSTALL_NO_NCCL=1 set - skipping nccl"
        } elseif (Test-NCCL) {
            $features += "nccl"
            Write-Info "NCCL detected - enabling nccl for CUDA multi-GPU tensor parallelism"
        } elseif ($env:MISTRALRS_INSTALL_NCCL -eq "1") {
            $features += "nccl"
            Write-Warn "MISTRALRS_INSTALL_NCCL=1 set but NCCL was not detected; the build may fail unless NCCL is on the linker path"
        } else {
            Write-Warn "NCCL not found - skipping nccl. Install NCCL or set MISTRALRS_INSTALL_NCCL=1 to force it; NCCL is the preferred CUDA multi-GPU path."
        }

        # Check for cuDNN
        if (Test-CuDNN) {
            $features += "cudnn"
            Write-Info "cuDNN detected - enabling cudnn"
        } else {
            Write-Info "cuDNN not found - skipping cudnn feature"
        }

        # Add flash attention based on compute capability
        if ($cudaCC -eq "90") {
            $features += "flash-attn-v3"
            Write-Info "Hopper GPU detected - enabling flash-attn-v3"
        } elseif ([int]$cudaCC -ge 80) {
            $features += "flash-attn"
            Write-Info "Ampere+ GPU detected - enabling flash-attn"
        }

        # cuTile: optimized CUDA kernels. Needs CUDA >= 13.1 (its JIT tool tileiras ships with 13.1+); runs on Ampere (80-89) or Blackwell+ (>=100), not Hopper (90-99).
        $cudaVer = Get-CudaVersionCode
        $ccNum = [int]$cudaCC
        if ($cudaVer -and $cudaVer -ge 1301 -and ((($ccNum -ge 80) -and ($ccNum -lt 90)) -or ($ccNum -ge 100))) {
            $features += "cutile"
            Write-Info "CUDA >= 13.1 and supported arch - enabling cutile (optimized kernels)"
        }
    } else {
        Write-Info "No NVIDIA GPU detected"
    }

    # Check for MKL on Intel
    if ((Test-IntelCpu) -and (Test-MKL)) {
        $features += "mkl"
        Write-Info "Intel MKL detected - enabling mkl"
    }

    return $features -join " "
}

# Install mistralrs-cli
function Install-MistralRS {
    param([string]$Features)

    # MISTRALRS_INSTALL_TAG pins a git tag; otherwise build the latest master.
    if ($env:MISTRALRS_INSTALL_TAG) {
        $gitRef = @("--tag", $env:MISTRALRS_INSTALL_TAG)
        $refDesc = "tag $($env:MISTRALRS_INSTALL_TAG)"
    } else {
        $gitRef = @("--branch", $MistralRsBranch)
        $refDesc = "branch $MistralRsBranch"
    }

    if ($Features) {
        Write-Info "Installing mistralrs-cli from GitHub $refDesc with features: $Features"
        & cargo install --force --locked --git $MistralRsRepoUrl @gitRef $MistralRsCliPackage --features "$Features"
    } else {
        Write-Info "Installing mistralrs-cli from GitHub $refDesc with default features"
        & cargo install --force --locked --git $MistralRsRepoUrl @gitRef $MistralRsCliPackage
    }

    if ($LASTEXITCODE -ne 0) {
        Write-Err "Installation failed"
    }
}

# MISTRALRS_INSTALL_TAG pins a specific release (e.g. v0.8.4); default is the latest stable release.
$ReleaseBase = if ($env:MISTRALRS_INSTALL_TAG) {
    "https://github.com/EricLBuehler/mistral.rs/releases/download/$($env:MISTRALRS_INSTALL_TAG)"
} else {
    "https://github.com/EricLBuehler/mistral.rs/releases/latest/download"
}
$PrebuiltDir = "$env:USERPROFILE\.mistralrs"
$BinDir = "$env:USERPROFILE\.local\bin"

# Download and install the Windows CPU prebuilt. Returns $true on success.
function Install-Prebuilt {
    $asset = "mistralrs-cpu-x86_64-pc-windows-msvc.zip"
    $tmp = Join-Path $env:TEMP $asset
    Write-Info "Downloading $asset"
    try {
        # Start-BitsTransfer shows a native progress bar and is fast; fall back to IWR with its bar.
        try {
            Start-BitsTransfer -Source "$ReleaseBase/$asset" -Destination $tmp -ErrorAction Stop
        } catch {
            $ProgressPreference = 'Continue'
            Invoke-WebRequest -Uri "$ReleaseBase/$asset" -OutFile $tmp -UseBasicParsing
        }
    } catch {
        return $false
    }
    try {
        if (Test-Path $PrebuiltDir) { Remove-Item -Recurse -Force $PrebuiltDir }
        New-Item -ItemType Directory -Force -Path $PrebuiltDir | Out-Null
        Expand-Archive -Path $tmp -DestinationPath $PrebuiltDir -Force
        Remove-Item -Force $tmp
        New-Item -ItemType Directory -Force -Path $BinDir | Out-Null
        Copy-Item -Force "$PrebuiltDir\mistralrs.exe" "$BinDir\mistralrs.exe"
        & "$BinDir\mistralrs.exe" --version *> $null
        if ($LASTEXITCODE -ne 0) { return $false }
    } catch {
        return $false
    }
    return $true
}

# Main installation flow
function Main {
    Show-Banner

    Write-Info "Detected OS: Windows"

    # Prefer the prebuilt CPU binary: no Rust toolchain, no compile. Set
    # MISTRALRS_INSTALL_FROM_SOURCE=1 to force a source build instead.
    if (-not $env:MISTRALRS_INSTALL_FROM_SOURCE) {
        Write-Info "Checking for a prebuilt binary..."
        if (Install-Prebuilt) {
            $ver = (& "$BinDir\mistralrs.exe" --version 2>$null | Select-Object -First 1)
            if (-not $ver) { $ver = "mistral.rs" }
            Write-Success "$ver installed successfully (prebuilt binary)!"
            Write-Host ""
            Write-Host "Installed" -ForegroundColor White
            Write-Host "========="
            Write-Host "  binary   $PrebuiltDir\mistralrs.exe"
            Write-Host "  on PATH  $BinDir\mistralrs.exe (copy)"
            Write-Host ""
            if ($env:PATH -notmatch [regex]::Escape($BinDir)) {
                Write-Warn "Add $BinDir to your PATH to run 'mistralrs' from any terminal."
            }
            Write-Host ""
            Write-Host "  mistralrs run -m Qwen/Qwen3-4B"
            Write-Host ""
            return
        }
        Write-Warn "Prebuilt install failed; building from source instead."
        Write-Host ""
    }

    if ($env:MISTRALRS_INSTALL_TAG) {
        Write-Info "Building from source: tag $($env:MISTRALRS_INSTALL_TAG)."
    } else {
        Write-Info "Building from source: latest $MistralRsBranch (bleeding edge)."
    }

    # Check for Rust
    if (Test-Rust) {
        $rustVersionFull = & rustc --version 2>$null
        $rustVersion = Get-RustVersion
        Write-Info "Rust is installed: $rustVersionFull"

        # Check if version meets minimum requirement
        if ($rustVersion -and -not (Test-VersionGte $rustVersion $RequiredRustVersion)) {
            Write-Warn "Rust $rustVersion is below the required version $RequiredRustVersion"
            Write-Host ""
            $response = Read-Confirm "Would you like to update Rust now? [Y/n]"
            if ($response -match "^[Nn]") {
                Write-Err "Rust $RequiredRustVersion or newer is required to install mistral.rs"
            }
            Update-Rust
            # Re-check version after update
            $rustVersion = Get-RustVersion
            if (-not (Test-VersionGte $rustVersion $RequiredRustVersion)) {
                Write-Err "Failed to update Rust to required version $RequiredRustVersion"
            }
        }
    } else {
        Write-Warn "Rust is not installed"
        Write-Host ""
        $response = Read-Confirm "Would you like to install Rust now? [Y/n]"
        if ($response -match "^[Nn]") {
            Write-Err "Rust is required to install mistral.rs"
        }
        Install-Rust
    }

    Write-Host ""
    Write-Info "Detecting hardware capabilities..."

    # Build features
    $features = Get-Features

    Write-Host ""
    Write-Host "Installation Summary" -ForegroundColor White
    Write-Host "====================" -ForegroundColor White
    if ($features) {
        Write-Host "Features: " -NoNewline
        Write-Host $features -ForegroundColor Green
    } else {
        Write-Host "Features: " -NoNewline
        Write-Host "(none - CPU only)" -ForegroundColor Yellow
    }
    Write-Host ""

    # Confirm installation
    $response = Read-Confirm "Proceed with installation? [Y/n]"
    if ($response -match "^[Nn]") {
        Write-Info "Installation cancelled"
        exit 0
    }

    Write-Host ""
    Install-MistralRS -Features $features

    Write-Host ""
    $ver = (& "$env:USERPROFILE\.cargo\bin\mistralrs.exe" --version 2>$null | Select-Object -First 1)
    if (-not $ver) { $ver = "mistral.rs" }
    Write-Host ""
    Write-Success "$ver installed successfully!"
    Write-Host ""
    Write-Host "Installed" -ForegroundColor White
    Write-Host "========="
    Write-Host "  binary   $env:USERPROFILE\.cargo\bin\mistralrs.exe"
    Write-Host ""
    Write-Host "Quick Start" -ForegroundColor White
    Write-Host "===========" -ForegroundColor White
    Write-Host ""
    Write-Host "  mistralrs run -m Qwen/Qwen3-4B"
    Write-Host ""
    Write-Host "  mistralrs serve --agent -m google/gemma-4-E4B-it"
    Write-Host ""
    Write-Host "For more information, visit: https://github.com/EricLBuehler/mistral.rs"
    Write-Host ""
    Write-Host "Note: " -ForegroundColor Yellow -NoNewline
    Write-Host "Restart your terminal to use the 'mistralrs' command."
}

# Run main
Main