# WSL2 Environment Upgrade Script
# This script upgrades Debian system packages and configures WSL2 settings

# Set error action
$ErrorActionPreference = 'Stop'

# Color definitions
$Red = [System.ConsoleColor]::Red
$Green = [System.ConsoleColor]::Green
$Yellow = [System.ConsoleColor]::Yellow

# Logging functions
function Write-Info {
    param([string]$Message)
    Write-Host '[INFO] ' -NoNewline -ForegroundColor $Green
    Write-Host $Message
}

function Write-Warn {
    param([string]$Message)
    Write-Host '[WARN] ' -NoNewline -ForegroundColor $Yellow
    Write-Host $Message
}

function Write-Error {
    param([string]$Message)
    Write-Host '[ERROR] ' -NoNewline -ForegroundColor $Red
    Write-Host $Message
}

function Wait-KeyPress {
    Write-Host "`nPress any key to exit..." -NoNewline
    $null = $Host.UI.RawUI.ReadKey('NoEcho,IncludeKeyDown')
    Write-Host ''
}

function Get-UserConfirmation {
    param([string]$Message)
    $response = Read-Host -Prompt "$Message [y/N]"
    return $response -eq 'y' -or $response -eq 'Y'
}

function Get-CustomPath {
    param([string]$DefaultPath)
    Write-Host "`nCurrent WSL installation path: $DefaultPath"
    $response = Read-Host -Prompt "Enter new installation path or press Enter to keep current"
    if ([string]::IsNullOrWhiteSpace($response)) {
        return $DefaultPath
    }
    return $response
}

# Check administrator privileges
if (-NOT ([Security.Principal.WindowsPrincipal][Security.Principal.WindowsIdentity]::GetCurrent()).IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)) {
    Write-Warn 'Administrator privileges required'
    Write-Info 'Please run this script as Administrator'
    Wait-KeyPress
    exit 1
}

Write-Info "`nStarting WSL2 environment configuration...`n"

Write-Info 'Note: Please make sure you have manually updated WSL2 kernel before proceeding.'
Write-Info 'You can update WSL2 kernel by running "wsl --update" in an administrator PowerShell.'
Write-Host ''

# Check if Debian is installed
$debianInstalled = $false
try {
    wsl -d Debian echo "Check Debian" 2>$null
    $debianInstalled = $true
    Write-Info "Found existing Debian installation"
} catch {
    Write-Info "No existing Debian installation found"
}

# Handle existing Debian installation
if ($debianInstalled) {
    if (Get-UserConfirmation "Do you want to remove the existing Debian installation?") {
        Write-Info "Removing existing Debian installation..."
        try {
            wsl --unregister Debian
            Write-Info "Existing Debian installation removed"
            $debianInstalled = $false
        } catch {
            Write-Error "Failed to remove Debian: $($_.Exception.Message)"
            Wait-KeyPress
            exit 1
        }
    }
}

# Get custom installation path if Debian needs to be installed
if (-not $debianInstalled) {
    $defaultPath = "$env:USERPROFILE\WSL\Debian"
    $customPath = Get-CustomPath $defaultPath
    
    if ($customPath -ne $defaultPath) {
        Write-Info "Setting up custom installation path: $customPath"
        
        # Create directory if it doesn't exist
        if (-not (Test-Path $customPath)) {
            try {
                New-Item -ItemType Directory -Path $customPath -Force | Out-Null
                Write-Info "Created directory: $customPath"
            } catch {
                Write-Error "Failed to create directory: $($_.Exception.Message)"
                Wait-KeyPress
                exit 1
            }
        }
        
        # Set WSL default installation path
        Write-Info "Setting WSL installation path..."
        try {
            [System.Environment]::SetEnvironmentVariable('WSLENV', "WSLPATH/up:$customPath", [System.EnvironmentVariableTarget]::User)
            Write-Info "WSL installation path set to: $customPath"
        } catch {
            Write-Error "Failed to set WSL installation path: $($_.Exception.Message)"
            Wait-KeyPress
            exit 1
        }
    }
    
    # Install Debian
    Write-Info "Installing Debian..."
    try {
        wsl --install -d Debian
        Write-Info "Debian installation completed"
        
        # Wait for Debian to be ready
        Write-Info "Waiting for Debian to be ready..."
        Start-Sleep -Seconds 10
    } catch {
        Write-Error "Failed to install Debian: $($_.Exception.Message)"
        Wait-KeyPress
        exit 1
    }
}

# Shutdown WSL to ensure clean state
Write-Info 'Shutting down WSL to ensure clean state...'
try {
    wsl --shutdown
    Start-Sleep -Seconds 2
    Write-Info 'WSL shutdown completed'
} catch {
    Write-Error "Failed to shutdown WSL: $($_.Exception.Message)"
    Wait-KeyPress
    exit 1
}

# Update Debian system
Write-Info 'Updating Debian system packages...'
try {
    # Update package lists
    Write-Info 'Updating package lists...'
    wsl -d Debian -u root apt update

    # Upgrade packages
    Write-Info 'Upgrading packages...'
    wsl -d Debian -u root apt upgrade -y

    # Clean up
    Write-Info 'Cleaning up...'
    wsl -d Debian -u root apt autoremove -y
    wsl -d Debian -u root apt clean

    Write-Info 'Debian system update completed'
} catch {
    Write-Error "Failed to update Debian system: $($_.Exception.Message)"
    Wait-KeyPress
    exit 1
}

# Configure WSL memory limit
Write-Info 'Configuring WSL memory limit...'
try {
    $wslConfigPath = "$env:USERPROFILE\.wslconfig"
    $wslConfig = @"
[wsl2]
memory=8GB
processors=4
swap=2GB
"@

    if (Test-Path $wslConfigPath) {
        Write-Info 'Backing up existing .wslconfig...'
        Copy-Item $wslConfigPath "$wslConfigPath.bak"
    }

    $wslConfig | Out-File $wslConfigPath -Encoding UTF8
    Write-Info 'WSL memory limit configured'
} catch {
    Write-Error "Failed to configure WSL memory limit: $($_.Exception.Message)"
    Wait-KeyPress
    exit 1
}

Write-Host "`n============================================"
Write-Host "         Configuration Complete!            "
Write-Host "============================================"
Write-Host "`nYour WSL2 environment has been configured:`n"
Write-Info '1. Debian system packages updated'
Write-Info '2. WSL memory limit configured (8GB RAM, 4 cores, 2GB swap)'
if (-not [string]::IsNullOrWhiteSpace($customPath)) {
    Write-Info "3. WSL installation path set to: $customPath"
}
Write-Info 'Please restart your computer to apply all changes.'
Write-Host "============================================"

# Wait for user input before exit
Wait-KeyPress