Komorebi script to save/load window state

kmrb.ps1

# Powershell script to save and load komorebi window state
#
# Usage:
#   .\kmrb.ps1           = load window state for current monitor config if previously saved
#     -list              = list all existing configs
#     -config name       = load window state from specified config name
#     -save              = save current window state for current monitor config
#     -save -config name = save current window state to specified config name
#     -dbg               = show komorebic commands invoked

param (
    [Parameter(Mandatory = $false)][switch]$list = $false,
    [Parameter(Mandatory = $false)][switch]$save = $false,
    [Parameter(Mandatory = $false)][string]$config = "",
    [Parameter(Mandatory = $false)][switch]$dbg = $false
)
class Kmrb {
    [string]$config_dir
    [string]$config_id
    [string]$config_file
    [PSCustomObject]$current_state
    [string]$config
    [bool]$dbg
    [array]$monitors
    [string]$monitors_str
    [PSCustomObject]$target_state

    [void]main($list, $save, $config, $dbg) {
        # Store parameters
        $this.config = $config
        $this.dbg = $dbg

        # Get the config directory
        $this.config_dir = Join-Path -Path (Join-Path -Path ($Env:HOMEDRIVE + $Env:HOMEPATH) -ChildPath ".config") -ChildPath "kmrb"
        if (-not (Test-Path -Path $this.config_dir)) {
            New-Item -ItemType Directory -Path $this.config_dir | Out-Null
        }

        if ($list) {
            $this.list_configs()
            exit 0
        }

        # Check komorebi
        $this.check_komorebi($save)

        # Get current state of komorebi
        $this.get_state()

        # Do what user requests
        if ($save) {
            $this.save()
        }
        else {
            $this.load()
        }
    }

    [string]popen([string]$cmd) {
        if ($this.dbg) {
            # Show command if dbg flag is set
            Write-Host $cmd
        }

        # Run shell command and get output and exit code
        $process = Start-Process -FilePath "powershell.exe" -ArgumentList "-Command", $cmd -NoNewWindow -RedirectStandardOutput "output.txt" -RedirectStandardError "error.txt" -PassThru -Wait
        $output = Get-Content -Path "output.txt" -Raw
        $errorOutput = Get-Content -Path "error.txt" -Raw

        Remove-Item -Path "output.txt", "error.txt"

        if ($process.ExitCode -ne 0) {
            Write-Host "Failed to run: $cmd"
            Write-Host "Returned: $($process.ExitCode)`n$errorOutput"
            exit 1
        }

        return $output
    }

    [int]run([string]$cmd) {
        if ($this.dbg) {
            # Show command if dbg flag is set
            Write-Host $cmd
        }

        Invoke-Expression -Command $cmd *> $null
        return $LASTEXITCODE
    }

    check_komorebi($save) {
        # Stop komorebi since not $save
        # $save means state is good and we don't want to stop komorebi
        if (!$save) {
            $this.run("komorebic stop")
        }

        # Start komorebi if not running
        $process = Get-Process -Name komorebi -ErrorAction SilentlyContinue
        if ($null -eq $process) {
            $this.run("komorebic start")
        }
    }

    [void]list_configs() {
        foreach ($config_file in Get-ChildItem -Path $this.config_dir -Filter "*.json") {
            $state = Get-Content -Path $config_file.FullName | ConvertFrom-Json
            $result = $this.get_monitors($state)
            $mons_str = $result[1]
            Write-Host "$($config_file.Name) = $mons_str"
        }
    }

    [array]get_monitors([PSCustomObject]$state) {
        # Extract the elements list
        $elements = $state.monitors.elements

        # Generate a list of monitors in order of appearance
        $mons = $elements | ForEach-Object { $_.device_id }
        $mons_str = $mons -join "+"

        return @($mons, $mons_str)
    }

    [void]get_state() {
        # Get komorebi state as a json object
        $this.current_state = $this.popen("komorebic state") | ConvertFrom-Json

        # Extract the monitors list
        $result = $this.get_monitors($this.current_state)
        $this.monitors = $result[0]
        $this.monitors_str = $result[1]

        if ($this.config.Length -ne 0) {
            # Get save name from CLI
            $this.config_id = $this.config
        }
        else {
            # Generate a hash from monitors in order to identify the state
            $md5 = [System.Security.Cryptography.MD5]::Create()
            $hashBytes = $md5.ComputeHash([System.Text.Encoding]::UTF8.GetBytes($this.monitors_str))
            $this.config_id = [BitConverter]::ToString($hashBytes) -replace '-', ''
        }

        $this.config_file = Join-Path -Path $this.config_dir -ChildPath "$($this.config_id).json"
    }

    [string]get_version() {
        $version = $this.popen("komorebic -V")

        # Split the string by newlines
        $lines = $version -split "`n"

        # Take the first line
        $firstLine = $lines[0]

        # Split the first line by space
        $parts = $firstLine -split " "

        # Take the last piece
        $lastPiece = $parts[-1].Trim()

        # Return the last piece
        return $lastPiece
    }

    [void]save() {
        # Remove empty workspaces
        foreach ($monitor in $this.current_state.monitors.elements) {
            $newWorkspaces = @()
            for ($idx = 0; $idx -lt $monitor.workspaces.elements.Count; $idx++) {
                $workspace = $monitor.workspaces.elements[$idx]
                if ($workspace.containers.elements.Count -ne 0) {
                    $newWorkspaces += $workspace
                }
            }

            $monitor.workspaces.elements = $newWorkspaces
        }

        # Save the current state to a json file
        Write-Host "Saving config file: $($this.config_file)"
        $json = $this.current_state | ConvertTo-Json -Depth 100 -Compress

        $json | Out-File -FilePath $this.config_file -Encoding utf8
    }

    [void]load() {
        # Check if the config file exists
        if (-not (Test-Path -Path $this.config_file)) {
            if ($this.config.Length -ne 0) {
                Write-Host "No config file found matching name: $($this.config)"
            }
            else {
                Write-Host "No config file $($this.config_id).json found matching state: $($this.monitors_str)"
                $this.list_configs()
            }
            exit 1
        }

        # Load the config file
        Write-Host "Loading config file: $($this.config_file)"
        $this.target_state = Get-Content -Path $this.config_file | ConvertFrom-Json

        $regex = [regex]::new('((?<=[a-z0-9])[A-Z]|(?!^)A-Z)')

        $monitor_id = 0
        $exes = @()
        foreach ($monitor in $this.target_state.monitors.elements) {
            # Ensure there are enough workspaces
            $this.run("komorebic ensure-workspaces $monitor_id $($monitor.workspaces.elements.Count)")

            $workspace_id = 0
            foreach ($workspace in $monitor.workspaces.elements) {
                # Set workspace layout
                $layout = $regex.Replace($workspace.layout.Default, '-$1').ToLower()
                $this.run("komorebic workspace-layout $monitor_id $workspace_id $layout")

                foreach ($container in $workspace.containers.elements) {
                    foreach ($window in $container.windows.elements) {
                        # Check if window executable already handled
                        if ($exes.Contains($window.exe)) {
                            continue
                        }
                        $exes += $window.exe

                        # Set workspace rule for each window by executable name
                        $this.run("komorebic initial-workspace-rule exe $($window.exe) $monitor_id $workspace_id")
                        $this.run("komorebic workspace-rule exe $($window.exe) $monitor_id $workspace_id")
                    }
                }

                $workspace_id++
            }

            $monitor_id++
        }

        # Apply changes
        $this.run("komorebic complete-configuration")
        $this.run("komorebic retile")
    }
}

$kmrb = [Kmrb]::new()
$kmrb.main($list, $save, $config, $dbg)