hammerspoon/WindowManager.lua

-- WindowManager.lua

local obj = {}
local json = require("hs.json")
local styledtext = require("hs.styledtext")

-- ==========================================
-- CONFIGURATION & CENTRAL CONFIG LOADING
-- ==========================================
obj.layoutFile = os.getenv("HOME") .. "/.hammerspoon/saved_layout.json"
local configFile = os.getenv("HOME") .. "/.hammerspoon/config.json"

local stubbornAppsList = {
    ["Gemini"] = true, ["AFFiNE"] = true, ["Terminal"] = true,
    ["System Settings"] = true, ["Hammerspoon"] = true
}

local ignoreListItems = {
    ["TheBoringNotch"] = true, ["Control Center"] = true, ["Notification Center"] = true,
    ["Dock"] = true, ["BetterDisplay"] = true, ["Stats"] = true, ["DockDoor"] = true,
    ["Bartender 6"] = true, ["Bartender"] = true, ["Arq"] = true, ["Arq Agent"] = true, ["TypeWhisper"] = true,
    ["com.typewhisper.mac"] = true
}

-- Load Central Config if available
local cfgData = hs.json.read(configFile)
if cfgData then
    if cfgData.stubbornAppsList then stubbornAppsList = cfgData.stubbornAppsList end
    if cfgData.ignoreListItems then ignoreListItems = cfgData.ignoreListItems end
end

obj.saveInterval = 300 
obj.saveCountdown = obj.saveInterval 
obj.isRescued = false 
obj.isTransitioning = false 
obj.isRestoring = false 
obj.isMenuUpdating = false -- Safety guard to prevent concurrent IPC menu rendering loops
obj.wakeTimer = nil     
obj.lastScreenCount = #hs.screen.allScreens()
obj.lastSavedTime = "Never"

-- ==========================================
-- INTERNAL UTILITIES
-- ==========================================

local function log(msg)
    print(string.format("WindowManager [%s]: %s", os.date("%H:%M:%S"), msg))
end

local function wrapText(text, limit)
    local lines = {}
    local currentLine = ""
    for word in text:gmatch("%S+") do
        if #currentLine + #word >= limit then
            table.insert(lines, currentLine)
            currentLine = "            " .. word 
        else
            currentLine = currentLine == "" and "      ↳ " .. word or currentLine .. " " .. word
        end
    end
    table.insert(lines, currentLine)
    return lines
end

-- ==========================================
-- CORE LOGIC
-- ==========================================

function obj.saveLayout(silent)
    local currentScreens = #hs.screen.allScreens()

    if currentScreens ~= obj.lastScreenCount then
        log(string.format("Screen mismatch (%d vs %d). Syncing count.", currentScreens, obj.lastScreenCount))
        obj.lastScreenCount = currentScreens
    end
    
    local layout = { 
        saveTime = os.date("%H:%M:%S"), 
        screenCount = currentScreens,
        windows = {} 
    }
    
    for _, win in ipairs(hs.window.allWindows()) do
        pcall(function()
            local app = win:application()
            if app and win:isVisible() and win:frame().w > 0 then
                local appName = app:name() or ""
                local bundleID = app:bundleID() or ""
                
                if not ignoreListItems[appName] and not ignoreListItems[bundleID] then
                    table.insert(layout.windows, {
                        appName = appName,
                        bundleID = bundleID,
                        winTitle = win:title(),
                        x = math.floor(win:frame().x),
                        y = math.floor(win:frame().y),
                        w = math.floor(win:frame().w),
                        h = math.floor(win:frame().h)
                    })
                end
            end
        end)
    end
    
    hs.json.write(layout, obj.layoutFile, true, true)
    obj.lastSavedTime = layout.saveTime
    
    if silent then
        log("Autosave successful.")
    else
        log("Manual save successful.")
        hs.alert.show("Layout Saved", 1.5) 
    end
end

function obj.restoreLayout()
    obj.isRestoring = true
    obj.isTransitioning = false
    obj.isRescued = false 

    local data = hs.json.read(obj.layoutFile)
    if not data or not data.windows then 
        log("Restore FAILED: No data found in JSON.")
        obj.isRestoring = false
        return 
    end
    
    hs.screen.restoreGamma() 
    
    local savedByApp = {}
    for _, winData in ipairs(data.windows) do
        if not ignoreListItems[winData.appName] and not ignoreListItems[winData.bundleID] then
            savedByApp[winData.appName] = savedByApp[winData.appName] or {}
            table.insert(savedByApp[winData.appName], winData)
        end
    end

    for appName, entries in pairs(savedByApp) do
        table.sort(entries, function(a, b)
            if a.y == b.y then return a.x < b.x end
            return a.y < b.y
        end)
    end

    log("RESTORE STARTING...")

    for appName, savedEntries in pairs(savedByApp) do
        local firstEntry = savedEntries[1]
        local app = hs.application.get(firstEntry.bundleID) or hs.application.get(appName)
        
        if app then
            local physicalWins = {}
            for _, w in ipairs(app:allWindows()) do 
                if w:isVisible() and w:frame().w > 0 then table.insert(physicalWins, w) end 
            end
            
            table.sort(physicalWins, function(a, b)
                local af, bf = a:frame(), b:frame()
                if af.y == bf.y then return af.x < bf.x end
                return af.y < bf.y
            end)
            
            for i, winData in ipairs(savedEntries) do
                local win = physicalWins[i]
                if win then
                    log(string.format("Moving %s (%d/%d)", appName, i, #savedEntries))
                    if win:isMinimized() then win:unminimize() end
                    
                    pcall(function()
                        win:setFrame({x=winData.x, y=winData.y, w=winData.w, h=winData.h}, 0)
                        
                        if stubbornAppsList[appName] then
                            hs.timer.doAfter(0.5, function()
                                pcall(function() win:setFrame({x=winData.x, y=winData.y, w=winData.w, h=winData.h}, 0) end)
                            end)
                        end
                    end)
                end
            end
        else
            log(string.format("Skip: %s is not running.", appName))
        end
    end
    
    hs.alert.show("Layout Restored", 1.5)
    hs.timer.doAfter(5, function() 
        obj.isRestoring = false 
        obj.isTransitioning = false 
        log("RESTORE CYCLE COMPLETE.")
    end)
end

function obj.rescueWindowsToLaptop()
    local primary = hs.screen.primaryScreen()
    if not primary then return end
    local maxFrame = primary:frame()
    log("RESCUE: Cascading windows on laptop.")
    
    local allWindows = hs.window.allWindows()
    local staggerOffset = 0
    
    for _, win in ipairs(allWindows) do
        pcall(function()
            local app = win:application()
            if app and win:isVisible() and win:frame().w > 0 then
                local appName = app:name() or ""
                local bundleID = app:bundleID() or ""
                
                if not ignoreListItems[appName] and not ignoreListItems[bundleID] then
                    local f = win:frame()
                    if f.w > maxFrame.w then f.w = maxFrame.w - 100 end
                    if f.h > maxFrame.h then f.h = maxFrame.h - 100 end
                    
                    f.x = maxFrame.x + 50 + staggerOffset
                    f.y = maxFrame.y + 50 + staggerOffset
                    
                    win:setFrame(f, 0)
                    staggerOffset = staggerOffset + 30
                    if staggerOffset > 150 then staggerOffset = 0 end
                end
            end
        end)
    end
    obj.isRescued = true
    obj.isTransitioning = false
    hs.alert.show("Windows Cascaded", 1.5)
end

-- ==========================================
-- MENUBAR & WATCHERS
-- ==========================================
local timerMenu = hs.menubar.new()

function updateMenu()
    -- Safety Guard: Bail if update is already running from a previous thread block
    if obj.isMenuUpdating then return end
    obj.isMenuUpdating = true

    if timerMenu then
        local screens = hs.screen.allScreens()
        timerMenu:setTitle(string.format("💠 %d:%02d", math.floor(obj.saveCountdown / 60), obj.saveCountdown % 60))
        local menuTable = {
            { title = "📅 Last Saved: " .. (obj.lastSavedTime or "Never"), disabled = true },
            { title = ((#screens > 1) and "🖥️ Docked" or "💻 Laptop") .. " (" .. #screens .. " Screens)", disabled = true },
            { title = "-" },
            { title = "📸 Save Layout (⇧⌘W)", fn = function() obj.saveLayout(false); obj.saveCountdown = obj.saveInterval end },
            { title = "🔄 Restore Layout (⇧⌘R)", fn = function() obj.restoreLayout() end },
            { title = "🚀 Rescue Windows (Cascade on Laptop) (⇧⌘⌃L)", fn = obj.rescueWindowsToLaptop },
            { title = "-" },
            { title = "📦 Saved Apps:", disabled = true }
        }
        
        local data = hs.json.read(obj.layoutFile)
        if data and data.windows then
            local seen, names = {}, {}
            for _, w in ipairs(data.windows) do 
                if not seen[w.appName] then table.insert(names, w.appName .. ", "); seen[w.appName] = true end 
            end
            for _, line in ipairs(wrapText(table.concat(names):gsub(", $", ""), 45)) do 
                table.insert(menuTable, { title = line, disabled = true }) 
            end
        end
        timerMenu:setMenu(menuTable)
    end

    obj.isMenuUpdating = false
end

obj.powerWatcher = hs.caffeinate.watcher.new(function(event)
    if event == hs.caffeinate.watcher.systemWillSleep or event == hs.caffeinate.watcher.screensDidSleep or event == hs.caffeinate.watcher.screensDidLock then
        log("POWER: Sleep event.")
        obj.isTransitioning = true
        -- Stop both timers entirely so they don't fire or stack requests while suspended
        if obj.autoSaveTimer then obj.autoSaveTimer:stop() end
        if obj.clockTimer then obj.clockTimer:stop() end
    elseif event == hs.caffeinate.watcher.systemDidWake or event == hs.caffeinate.watcher.screensDidWake or event == hs.caffeinate.watcher.screensDidUnlock then
        log("POWER: Wake event.")
        obj.saveCountdown = obj.saveInterval
        obj.isMenuUpdating = false -- Explicitly clean guard flag upon waking up
        
        -- Safely start timers up only after wakeup initialization
        if obj.autoSaveTimer then obj.autoSaveTimer:start() end
        if obj.clockTimer then obj.clockTimer:start() end
        if obj.wakeTimer then obj.wakeTimer:stop() end
        
        local currentScreens = #hs.screen.allScreens()
        
        if currentScreens > 1 then
            obj.wakeTimer = hs.timer.doAfter(12, function() 
                obj.isTransitioning = false
                obj.isRestoring = false 
                obj.lastScreenCount = currentScreens
                obj.restoreLayout()
                obj.wakeTimer = nil
            end)
        else
            log("WAKE SKIP: Single screen detected. Syncing count only.")
            obj.isTransitioning = false
            obj.isRestoring = false
            obj.lastScreenCount = currentScreens
            updateMenu()
        end
    end
end):start()

hs.hotkey.bind({"shift", "cmd"}, "W", function() obj.saveLayout(false); obj.saveCountdown = obj.saveInterval end)
hs.hotkey.bind({"shift", "cmd"}, "R", function() obj.restoreLayout() end)
hs.hotkey.bind({"shift", "cmd", "ctrl"}, "L", obj.rescueWindowsToLaptop)

obj.screenWatcher = hs.screen.watcher.new(function()
    if obj.isTransitioning or obj.isRestoring or obj.wakeTimer then 
        log("DOCK EVENT: Ignored.")
        return 
    end
    
    local currentScreens = #hs.screen.allScreens()
    
    if currentScreens == obj.lastScreenCount then
        log("DOCK EVENT: Ignored (Screen count unchanged).")
        return
    end

    log("DOCK EVENT: Detected.")
    obj.isTransitioning = true
    
    hs.timer.doAfter(15, function()
        if obj.isTransitioning and not obj.isRestoring then
            obj.isTransitioning = false
            log("STATE GUARD: Emergency clear of busy flag.")
        end
    end)

    hs.timer.doAfter(7, function() 
        obj.lastScreenCount = #hs.screen.allScreens()
        obj.isTransitioning = false
        if obj.lastScreenCount > 1 then obj.restoreLayout() else obj.rescueWindowsToLaptop() end
        updateMenu()
    end)
end):start()

obj.autoSaveTimer = hs.timer.doEvery(obj.saveInterval, function() 
    obj.saveLayout(true) 
    obj.saveCountdown = obj.saveInterval 
end)

obj.clockTimer = hs.timer.doEvery(1, function()
    obj.saveCountdown = obj.saveCountdown - 1
    if obj.saveCountdown < 0 then obj.saveCountdown = obj.saveInterval end
    updateMenu()
end)

updateMenu()

-- ==========================================
-- RAYCAST INTERACTIVE CHOOSER MENU
-- ==========================================
function obj.showMenu()
    local min = math.floor(obj.saveCountdown / 60)
    local sec = obj.saveCountdown % 60

    -- 1. Base Core Operational Rows
    local choices = {
        {
            text = string.format("⏱️ Next Autosave: %d:%02d", min, sec),
            subText = "📅 Last Saved Profile State: " .. obj.lastSavedTime,
            action = "info"
        },
        {
            text = "📸 Save State Layout",
            subText = "Snapshot active positions to saved_layout.json (Resets auto-timer)",
            action = "save"
        },
        {
            text = "🔄 Restore State Layout",
            subText = "Force apps and window sizes back to your saved profile state",
            action = "restore"
        },
        {
            text = "🚀 Rescue Windows",
            subText = "Cascade active window threads onto primary laptop screen space",
            action = "rescue"
        },
        {
            text = "───────────────────────────────────────────────────────",
            subText = "📦 Captured Windows Inside Profile State File:",
            action = "info"
        }
    }

    -- 2. Dynamically read layout file and append saved windows
    local data = hs.json.read(obj.layoutFile)
    if data and data.windows and #data.windows > 0 then
        for _, win in ipairs(data.windows) do
            -- Verify if the application is currently running on the system server
            local isRunning = hs.application.get(win.bundleID) or hs.application.get(win.appName)
            local statusIndicator = isRunning and "🟢" or "🔴"
            
            local cleanTitle = (win.winTitle and win.winTitle ~= "") and win.winTitle or "Untitled Window"
            if #cleanTitle > 60 then cleanTitle = string.sub(cleanTitle, 1, 57) .. "..." end

            table.insert(choices, {
                text = string.format("%s %s", statusIndicator, win.appName),
                subText = string.format("↳ Title: \"%s\" | Bounds: %dx%d at (%d,%d)", cleanTitle, win.w, win.h, win.x, win.y),
                action = "info"
            })
        end
    else
        table.insert(choices, {
            text = "⚠️ No Saved Window Metrics Found",
            subText = "Run a Save operation to record your desktop layout profile context.",
            action = "info"
        })
    end

    if obj.instanceChooser then obj.instanceChooser:hide() end

    obj.instanceChooser = hs.chooser.new(function(choice)
        if choice then
            if choice.action == "save" then
                obj.saveLayout(false)
                obj.saveCountdown = obj.saveInterval
            elseif choice.action == "restore" then
                obj.restoreLayout()
            elseif choice.action == "rescue" then
                obj.rescueWindowsToLaptop()
            end
        end
    end)

    obj.instanceChooser:placeholderText("Select a background workspace operation...")
    obj.instanceChooser:choices(choices)
    obj.instanceChooser:show()
end

-- Export module instance globally for direct IPC command routing
WindowManager = obj

return obj