script.go

package main

import (
	"fmt"
	"sort"
	"strings"
	"time"
)

// ---------------------------------------------------------------- script

func psQuote(s string) string { return "'" + strings.ReplaceAll(s, "'", "''") + "'" }

const psBroadcast = `# --- Broadcast the change so new shells pick it up ---
if (-not ('Win32.NativeMethods' -as [type])) {
  Add-Type -Namespace Win32 -Name NativeMethods -MemberDefinition @'
[DllImport("user32.dll", SetLastError = true, CharSet = CharSet.Auto)]
public static extern IntPtr SendMessageTimeout(IntPtr hWnd, uint Msg, UIntPtr wParam,
    string lParam, uint fuFlags, uint uTimeout, out UIntPtr lpdwResult);
'@
}
$r = [UIntPtr]::Zero
[Win32.NativeMethods]::SendMessageTimeout([IntPtr]0xffff, 0x1A, [UIntPtr]::Zero,
    'Environment', 2, 5000, [ref]$r) | Out-Null
`

func psPreamble(scope Scope) string {
	var b strings.Builder
	b.WriteString("#requires -Version 5.1\n")
	if scope == ScopeMachine {
		b.WriteString("#requires -RunAsAdministrator\n")
	}
	b.WriteString("Set-StrictMode -Version Latest\n$ErrorActionPreference = 'Stop'\n\n")
	return b.String()
}

func buildScript(rep *Report) string {
	if !rep.Scope.writable() {
		return "# Process scope is read-only. Re-run with --scope user or --scope machine.\n"
	}
	ts := time.Now().Format("20060102-150405")
	var b strings.Builder

	fmt.Fprintf(&b, "# Generated by %s v%s on %s\n", appName, appVersion, rep.GeneratedAt)
	fmt.Fprintf(&b, "# Scope: %s   %d B -> %d B (saved %d B)\n",
		rep.Scope, rep.OriginalLen, rep.OptimizedLen, rep.Saved())
	b.WriteString(psPreamble(rep.Scope))

	b.WriteString("# --- 1. Backup (restore by double-clicking the .reg file) ---\n")
	fmt.Fprintf(&b, "$backup = Join-Path $env:USERPROFILE 'env-backup-%s-%s.reg'\n", rep.Scope, ts)
	fmt.Fprintf(&b, "reg export \"%s\" $backup /y | Out-Null\n", rep.Scope.hive())
	b.WriteString("Write-Host \"Backup written to $backup\" -ForegroundColor Green\n\n")

	if len(rep.GroupOrder) > 0 {
		b.WriteString("# --- 2. Group variables ---\n")
		for _, g := range rep.GroupOrder {
			fmt.Fprintf(&b,
				"Set-ItemProperty -Path '%s' -Name '%s' -Value %s -Type ExpandString\n",
				rep.Scope.psPath(), g, psQuote(rep.Groups[g]))
		}
		b.WriteString("\n")
	}

	b.WriteString("# --- 3. PATH (ExpandString keeps %VARS% working; setx would truncate) ---\n")
	fmt.Fprintf(&b, "Set-ItemProperty -Path '%s' -Name 'Path' -Value %s -Type ExpandString\n\n",
		rep.Scope.psPath(), psQuote(rep.Optimized))

	b.WriteString(psBroadcast)
	b.WriteString("Write-Host 'PATH updated. Open a new terminal to use it.' -ForegroundColor Green\n")
	return b.String()
}

// buildRestoreScript rewinds PATH and managed vars, deleting any that were
// created after the snapshot was taken.
func buildRestoreScript(sn *Snapshot, removeVars []string) string {
	var b strings.Builder
	fmt.Fprintf(&b, "# %s v%s — restore snapshot #%d (%s, taken %s)\n",
		appName, appVersion, sn.ID, sn.Scope, sn.TakenAt.Format("2006-01-02 15:04:05"))
	b.WriteString(psPreamble(sn.Scope))

	fmt.Fprintf(&b, "Set-ItemProperty -Path '%s' -Name 'Path' -Value %s -Type ExpandString\n",
		sn.Scope.psPath(), psQuote(sn.Path))

	if len(sn.Vars) > 0 {
		names := make([]string, 0, len(sn.Vars))
		for k := range sn.Vars {
			names = append(names, k)
		}
		sort.Strings(names)
		b.WriteString("\n# Managed variables as they were\n")
		for _, k := range names {
			fmt.Fprintf(&b,
				"Set-ItemProperty -Path '%s' -Name '%s' -Value %s -Type ExpandString\n",
				sn.Scope.psPath(), k, psQuote(sn.Vars[k]))
		}
	}
	if len(removeVars) > 0 {
		b.WriteString("\n# Variables created after this snapshot\n")
		for _, k := range removeVars {
			fmt.Fprintf(&b,
				"Remove-ItemProperty -Path '%s' -Name '%s' -ErrorAction SilentlyContinue\n",
				sn.Scope.psPath(), k)
		}
	}
	b.WriteString("\n")
	b.WriteString(psBroadcast)
	b.WriteString("Write-Host 'Restore complete. Open a new terminal.' -ForegroundColor Green\n")
	return b.String()
}