# RuvLLM ESP32 - Windows PowerShell Installer
# Run: .\install.ps1 [command]

param(
    [Parameter(Position=0)]
    [string]$Command = "install",

    [Parameter(Position=1)]
    [string]$Arg1 = ""
)

$ErrorActionPreference = "Stop"

# Colors
function Write-Color($Text, $Color) {
    Write-Host $Text -ForegroundColor $Color
}

function Write-Banner {
    Write-Color @"

╔══════════════════════════════════════════════════════════╗
║          RuvLLM ESP32 - Windows Installer                ║
║     Tiny LLM + RAG + Federation for Microcontrollers     ║
╚══════════════════════════════════════════════════════════╝

"@ Cyan
}

# Check if command exists
function Test-Command($cmdname) {
    return [bool](Get-Command -Name $cmdname -ErrorAction SilentlyContinue)
}

# Install Rust
function Install-Rust {
    if (Test-Command rustc) {
        $version = rustc --version
        Write-Color "✓ Rust: $version" Green
        return
    }

    Write-Color "Installing Rust..." Yellow

    # Download and run rustup
    $rustupUrl = "https://win.rustup.rs/x86_64"
    $rustupPath = "$env:TEMP\rustup-init.exe"

    Invoke-WebRequest -Uri $rustupUrl -OutFile $rustupPath
    Start-Process -FilePath $rustupPath -ArgumentList "-y" -Wait

    # Refresh PATH
    $env:Path = [System.Environment]::GetEnvironmentVariable("Path","Machine") + ";" + [System.Environment]::GetEnvironmentVariable("Path","User")

    Write-Color "✓ Rust installed" Green
}

# Install ESP32 toolchain
function Install-ESPToolchain {
    Write-Color "`nInstalling ESP32 toolchain..." Yellow

    # Install espup
    if (-not (Test-Command espup)) {
        Write-Host "Installing espup..."
        cargo install espup
    } else {
        Write-Color "✓ espup already installed" Green
    }

    # Run espup install
    Write-Host "Running espup install (this may take 5-10 minutes)..."
    espup install

    # Install espflash
    if (-not (Test-Command espflash)) {
        Write-Host "Installing espflash..."
        cargo install espflash
    } else {
        Write-Color "✓ espflash already installed" Green
    }

    # Install ldproxy
    if (-not (Test-Command ldproxy)) {
        Write-Host "Installing ldproxy..."
        cargo install ldproxy
    } else {
        Write-Color "✓ ldproxy already installed" Green
    }

    Write-Color "✓ ESP32 toolchain ready" Green
    Write-Color "`n⚠ Please restart PowerShell before building!" Yellow
}

# Build project
function Build-Project {
    Write-Color "`nBuilding RuvLLM ESP32..." Yellow

    # Source ESP environment if exists
    $exportScript = "$env:USERPROFILE\.espressif\esp-idf-export.ps1"
    if (Test-Path $exportScript) {
        . $exportScript
    }

    cargo build --release

    if ($LASTEXITCODE -eq 0) {
        Write-Color "✓ Build successful!" Green
    } else {
        Write-Color "✗ Build failed" Red
        exit 1
    }
}

# Flash to device
function Flash-Device {
    param([string]$Port = "COM6")

    Write-Color "`nFlashing to $Port..." Yellow

    # Detect port if not specified
    if ($Port -eq "COM6") {
        $ports = [System.IO.Ports.SerialPort]::GetPortNames()
        if ($ports.Count -gt 0) {
            $Port = $ports[0]
            Write-Color "Auto-detected port: $Port" Cyan
        }
    }

    espflash flash --port $Port --monitor target\xtensa-esp32-espidf\release\ruvllm-esp32-flash
}

# Setup cluster
function Setup-Cluster {
    param([int]$NumChips = 2)

    Write-Color "`nSetting up $NumChips-chip cluster..." Yellow

    $config = @"
# RuvLLM ESP32 Cluster Configuration
# Generated by install.ps1

[cluster]
name = "ruvllm-cluster"
chips = $NumChips
topology = "pipeline"  # pipeline, tensor, hybrid

[chips]
"@

    for ($i = 1; $i -le $NumChips; $i++) {
        $role = if ($i -eq 1) { "master" } else { "worker" }
        $port = "COM$($i + 5)"

        $config += @"

[[chips.nodes]]
id = $i
role = "$role"
port = "$port"
layers = [$([math]::Floor(($i-1) * 2 / $NumChips)), $([math]::Floor($i * 2 / $NumChips - 1))]
"@
    }

    $config | Out-File -FilePath "cluster.toml" -Encoding utf8

    Write-Color "✓ Created cluster.toml" Green
    Write-Host "`nEdit cluster.toml to set correct COM ports, then run:"
    Write-Host "  .\cluster-flash.ps1"
}

# Show help
function Show-Help {
    Write-Host @"
Usage: .\install.ps1 [command] [options]

Commands:
  install     Install all dependencies and build (default)
  build       Build the project only
  flash       Flash to ESP32 (optionally specify port)
  deps        Install dependencies only
  cluster     Setup cluster configuration
  help        Show this help

Examples:
  .\install.ps1                    # Full install and build
  .\install.ps1 flash COM6         # Flash to COM6
  .\install.ps1 cluster 5          # Setup 5-chip cluster
"@
}

# Main
Write-Banner

switch ($Command.ToLower()) {
    "install" {
        Install-Rust
        Install-ESPToolchain
        Write-Color "`n⚠ Restart PowerShell, then run: .\install.ps1 build" Yellow
    }
    "build" {
        Build-Project
        Write-Color "`nTo flash: .\install.ps1 flash COM6" Cyan
    }
    "flash" {
        $port = if ($Arg1) { $Arg1 } else { "COM6" }
        Flash-Device -Port $port
    }
    "deps" {
        Install-Rust
        Install-ESPToolchain
    }
    "cluster" {
        $chips = if ($Arg1) { [int]$Arg1 } else { 2 }
        Setup-Cluster -NumChips $chips
    }
    "help" {
        Show-Help
    }
    default {
        Write-Color "Unknown command: $Command" Red
        Show-Help
        exit 1
    }
}