# cleanup_ab.ps1
#
# Kill ONLY the Chrome processes that were spawned by agent-browser.
# Leaves your normal Chrome browser untouched.
#
# Usage (from PowerShell):
#   .\cleanup_ab.ps1

$ErrorActionPreference = 'SilentlyContinue'

# 1. Find agent-browser root processes.
$abProcesses = @(Get-Process -Name 'agent-browser-win32-x64' -ErrorAction SilentlyContinue)
if ($abProcesses.Count -eq 0) {
    Write-Output "No agent-browser process running. Nothing to clean up."
    exit 0
}

# 2. One fast WMI query, then build a parent->children map in memory.
$allProcs = Get-CimInstance -Query "SELECT ProcessId, ParentProcessId FROM Win32_Process"
$childMap = @{}
foreach ($proc in $allProcs) {
    $ppid = [int]$proc.ParentProcessId
    if (-not $childMap.ContainsKey($ppid)) {
        $childMap[$ppid] = New-Object System.Collections.Generic.List[int]
    }
    [void]$childMap[$ppid].Add([int]$proc.ProcessId)
}

# 3. BFS from each agent-browser root to collect the full descendant tree.
$toKill = @{}
$queue = New-Object System.Collections.Generic.Queue[int]
foreach ($ab in $abProcesses) {
    if (-not $toKill.ContainsKey($ab.Id)) {
        $toKill[$ab.Id] = $true
        $queue.Enqueue($ab.Id)
    }
}
while ($queue.Count -gt 0) {
    $pid = $queue.Dequeue()
    if ($childMap.ContainsKey($pid)) {
        foreach ($child in $childMap[$pid]) {
            if (-not $toKill.ContainsKey($child)) {
                $toKill[$child] = $true
                $queue.Enqueue($child)
            }
        }
    }
}

# 4. Kill the agent-browser roots first (Chrome children usually exit on their own).
foreach ($ab in $abProcesses) {
    try { Stop-Process -Id $ab.Id -Force } catch { }
}

# Give Windows a moment to reap children.
Start-Sleep -Milliseconds 500

# 5. Final sweep of any descendants still alive.
foreach ($pid in $toKill.Keys) {
    try { Stop-Process -Id $pid -Force } catch { }
}

$killed = $toKill.Count
Write-Output "Cleaned up $killed agent-browser process(es) and descendants."
Write-Output "Your normal Chrome (if any) was left alone."