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.isRescued = false 
obj.isTransitioning = false 
obj.isRestoring = false 
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 saveCountdown = obj.saveInterval
local timerMenu = hs.menubar.new()

function updateMenu()
    if timerMenu then
        local screens = hs.screen.allScreens()
        timerMenu:setTitle(string.format("💠 %d:%02d", math.floor(saveCountdown / 60), 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); 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
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
        if obj.autoSaveTimer then obj.autoSaveTimer: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.")
        saveCountdown = obj.saveInterval
        if obj.autoSaveTimer then obj.autoSaveTimer: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
        end
    end
end):start()

hs.hotkey.bind({"shift", "cmd"}, "W", function() obj.saveLayout(false); 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) 
    saveCountdown = obj.saveInterval 
end)

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

updateMenu()

-- ==========================================
-- RAYCAST INTERACTIVE CHOOSER MENU
-- ==========================================
function obj.showMenu()
    local choices = {
        {
            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"
        }
    }

    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)
                saveCountdown = obj.saveInterval
            elseif choice.action == "restore" then
                obj.restoreLayout()
            elseif choice.action == "rescue" then
                obj.rescueWindowsToLaptop()
            end
        end
    end)

    obj.instanceChooser:placeholderText("Workspace Manager Actions...")
    obj.instanceChooser:choices(choices)
    obj.instanceChooser:show()
end

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

return obj