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

-- ==========================================
-- CONFIGURATION
-- ==========================================
obj.layoutFile = os.getenv("HOME") .. "/.hammerspoon/saved_layout.json"
obj.saveInterval = 300 
obj.isRescued = false 
obj.isTransitioning = false 
obj.isRestoring = false 
obj.wakeTimer = nil     
obj.lastScreenCount = #hs.screen.allScreens()
obj.lastSavedTime = "Never"

local ignoreListItems = {
    ["TheBoringNotch"] = true, ["Control Center"] = true, ["Notification Center"] = true,
    ["Dock"] = true, ["BetterDisplay"] = true, ["Stats"] = true, ["DockDoor"] = true
}

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

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 (PURE NATIVE)
-- ==========================================

function obj.saveLayout(silent)
    if obj.isTransitioning or obj.isRestoring or #hs.screen.allScreens() ~= obj.lastScreenCount then 
        return 
    end
    
    local layout = { 
        saveTime = os.date("%H:%M:%S"), 
        screenCount = #hs.screen.allScreens(),
        windows = {} 
    }
    
    for _, win in ipairs(hs.window.allWindows()) do
        local app = win:application()
        if app and win:isVisible() and win:frame().w > 0 then
            local appName = app:name() or ""
            if not ignoreListItems[appName] then
                table.insert(layout.windows, {
                    appName = appName,
                    bundleID = app: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
    hs.json.write(layout, obj.layoutFile, true, true)
    obj.lastSavedTime = layout.saveTime
    if not silent then hs.alert.show("Layout Saved", 1.5) end
end

function obj.restoreLayout()
    if obj.isRestoring then return end

    local data = hs.json.read(obj.layoutFile)
    if not data or not data.windows then return end
    
    obj.isRestoring = true
    obj.isRescued = false 
    
    -- 1. Group saved windows by App Name
    local savedByApp = {}
    for _, winData in ipairs(data.windows) do
        savedByApp[winData.appName] = savedByApp[winData.appName] or {}
        table.insert(savedByApp[winData.appName], winData)
    end

    -- 2. Sort saved windows by position
    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

    print("WindowManager: Restore Starting (Native Engine)")

    -- 3. Move windows using Pure Hammerspoon logic
    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
            
            -- Sort physical windows to match
            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
                    if win:isMinimized() then win:unminimize() end
                    -- Move window to the exact saved frame
                    win:setFrame({
                        x = winData.x,
                        y = winData.y,
                        w = winData.w,
                        h = winData.h
                    }, 0) -- 0 duration for instant move
                end
            end
        end
    end
    
    hs.alert.show("Layout Restored", 1.5)
    
    hs.timer.doAfter(2, function() 
        obj.isRestoring = false 
        print("WindowManager: Restore Lock Cleared.")
    end)
end

function obj.rescueWindowsToLaptop()
    local primary = hs.screen.primaryScreen()
    if not primary then return end
    print("WindowManager: Rescue Triggered.")
    for _, win in ipairs(hs.window.allWindows()) do
        if win and win:isVisible() and win:frame().w > 0 then
            win:setFrame(primary:fullFrame(), 0)
        end
    end
    obj.isRescued = true
    hs.alert.show("Rescued to Laptop", 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 = "📸 Save Layout (⇧⌘W)", fn = function() obj.saveLayout(false); saveCountdown = obj.saveInterval end },
            { title = "🔄 Restore Layout (⇧⌘R)", fn = function() obj.isRestoring = false; obj.restoreLayout() end },
            { title = "🚀 Rescue Windows (Bring to Laptop) (⇧⌘⌃L)", fn = obj.rescueWindowsToLaptop }
        }
        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
        print("WindowManager: SLEEP")
        obj.isTransitioning = true
    elseif event == hs.caffeinate.watcher.systemDidWake or event == hs.caffeinate.watcher.screensDidWake or event == hs.caffeinate.watcher.screensDidUnlock then
        print("WindowManager: WAKE")
        if obj.wakeTimer then obj.wakeTimer:stop() end
        obj.wakeTimer = hs.timer.doAfter(10, function() 
            obj.isTransitioning = false
            obj.isRestoring = false 
            obj.restoreLayout()
            obj.wakeTimer = nil
        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.isRestoring = false; 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 return end
    print("WindowManager: DOCK")
    obj.isTransitioning = true
    hs.timer.doAfter(5, 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()
return obj