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.lastScreenCount = #hs.screen.allScreens()

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

local ignoreListItems = {
    ["TheBoringNotch"] = true,
    ["theboringteam.boringnotch"] = true,
    ["Control Center"] = true,
    ["Notification Center"] = true,
    ["Dock"] = true,
    ["com.surteesstudios.Bartender"] = true,
    ["pro.betterdisplay.BetterDisplay"] = true,
    ["stats"] = true,
    ["eu.exelban.Stats"] = true,
    ["com.ethanbills.DockDoor"] = true,
    ["DockDoor"] = true
}

-- ==========================================
-- POWER & SLEEP WATCHER
-- ==========================================

obj.powerWatcher = hs.caffeinate.watcher.new(function(event)
    if event == hs.caffeinate.watcher.systemWillSleep or 
       event == hs.caffeinate.watcher.systemWillPowerOff or
       event == hs.caffeinate.watcher.screensDidSleep then
        
        -- Lock immediately to prevent saving the "scrambled" sleep layout
        obj.isTransitioning = true
        
    elseif event == hs.caffeinate.watcher.systemDidWake or 
           event == hs.caffeinate.watcher.screensDidWake then
        
        -- Wait for monitors to handshake, then auto-restore
        hs.timer.doAfter(5, function()
            obj.isTransitioning = false
            obj.restoreLayout()
        end)
    end
end):start()

-- ==========================================
-- 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
-- ==========================================

function obj.saveLayout(silent)
    if obj.isTransitioning or hs.caffeinate.get("displayIdle") or hs.caffeinate.get("systemIdle") then return end
    
    local screens = hs.screen.allScreens()
    local layout = { 
        saveTime = os.date("%H:%M:%S"), 
        screenCount = #screens,
        mode = (#screens > 1) and "Docked" or "Laptop",
        windows = {} 
    }
    
    local allWindows = hs.window.allWindows()
    local consoleWin = hs.console.hswindow()
    if consoleWin then table.insert(allWindows, consoleWin) end

    for _, win in ipairs(allWindows) do
        local app = win:application()
        local screen = win:screen()
        local frame = win:frame()
        
        if app and screen and (win:isVisible() or win:title() == "Hammerspoon Console") and frame.w > 0 and frame.h > 0 then
            local appName = app:name() or ""
            local bid = app:bundleID() or ""
            if not (ignoreListItems[appName] or ignoreListItems[bid] or appName:find("TheBoringNotch")) then
                table.insert(layout.windows, {
                    appName = appName, 
                    bundleID = bid, 
                    winTitle = win:title(),
                    winID = win:id(), 
                    x = math.floor(frame.x), 
                    y = math.floor(frame.y),
                    w = math.floor(frame.w), 
                    h = math.floor(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()
    local data = hs.json.read(obj.layoutFile)
    if not data then return end
    obj.isRescued = false 
    local windowList = data.windows or data

    for _, winData in ipairs(windowList) do
        if winData.winTitle == "Hammerspoon Console" then
            local cWin = hs.console.hswindow()
            if cWin then cWin:setFrame({x=winData.x, y=winData.y, w=winData.w, h=winData.h}, 0) end
        else
            local app = hs.application.get(winData.bundleID) or hs.application.get(winData.appName)
            if app then
                local pid = app:pid()
                local appPath = app:path() or ""
                local isElectron = appPath:find("Electron") or appPath:find("Frameworks")

                for _, win in ipairs(app:allWindows()) do
                    if win:isMinimized() then win:unminimize() end
                    
                    local isExactMatch = (win:id() == winData.winID)
                    local isTitleMatch = (win:title() == winData.winTitle)
                    local isStubborn = stubbornAppsList[app:name()] or isElectron
                    
                    local isFuzzy = (winData.bundleID == "com.apple.Terminal" or winData.appName == "Gemini")
                    local shouldMove = (isExactMatch or (not winData.winID and isTitleMatch) or (#app:allWindows() == 1 and (isFuzzy or isStubborn)))

                    if shouldMove then
                        local x, y, w, h = math.floor(winData.x), math.floor(winData.y), math.floor(winData.w), math.floor(winData.h)

                        if isStubborn then
                            local safeTitle = win:title():gsub('"', '\\"')
                            local winTarget = (app:name() == "Gemini") and "window 1" or string.format('first window whose name is "%s"', safeTitle)
                            local script = string.format([[
                                tell application "System Events" to tell (first process whose unix id is %d)
                                    try
                                        set targetWin to %s
                                        set position of targetWin to {%d, %d}
                                        set size of targetWin to {%d, %d}
                                    end try
                                end tell
                            ]], pid, winTarget, x, y, w, h)
                            
                            hs.applescript.applescript(script)
                            hs.timer.doAfter(0.5, function() hs.applescript.applescript(script) end)
                        else
                            win:setFrame({x=x, y=y, w=w, h=h}, 0)
                        end
                    end
                end
            end
        end
    end
    hs.alert.show("Layout Restored", 1.5)
end

-- ==========================================
-- RESCUE LOGIC
-- ==========================================

function obj.rescueWindowsToLaptop()
    local screens = hs.screen.allScreens()
    if #screens == 1 and not obj.isRescued then
        local primary = hs.screen.primaryScreen()
        local primaryFrame = primary:frame()
        for _, win in ipairs(hs.window.allWindows()) do
            local app = win:application()
            if app then
                local name = app:name() or ""
                local bid = app:bundleID() or ""
                if not (ignoreListItems[name] or ignoreListItems[bid] or name:find("TheBoringNotch")) then
                    win:unminimize()
                    local appPath = app:path() or ""
                    local isElectron = appPath:find("Electron") or appPath:find("Frameworks")
                    if stubbornAppsList[name] or isElectron then
                        local pid = app:pid()
                        local safeTitle = win:title():gsub('"', '\\"')
                        local winTarget = (name == "Gemini") and "window 1" or string.format('first window whose name is "%s"', safeTitle)
                        local script = string.format([[
                            tell application "System Events" to tell (first process whose unix id is %d)
                                try
                                    set position of %s to {%d, %d}
                                end try
                            end tell
                        ]], pid, winTarget, math.floor(primaryFrame.x + 20), math.floor(primaryFrame.y + 40))
                        hs.applescript.applescript(script)
                    else
                        win:moveToScreen(primary, false, true)
                    end
                end
            end
        end
        obj.isRescued = true 
        hs.alert.show("Forced to Laptop Screen", 1.5)
    end
end

-- ==========================================
-- MENUBAR UI
-- ==========================================
local saveCountdown = obj.saveInterval
local timerMenu = hs.menubar.new()

function updateMenu()
    if timerMenu then
        local minutes = math.floor(saveCountdown / 60)
        local seconds = saveCountdown % 60
        timerMenu:setTitle(string.format("💠 %d:%02d", minutes, seconds))
        
        local screens = hs.screen.allScreens()
        local currentCount = #screens
        local modeIcon = (currentCount > 1) and "🖥️" or "💻"
        local modeName = (currentCount > 1) and "Docked" or "Laptop"
        
        local data = hs.json.read(obj.layoutFile)
        local lastTime = data and data.saveTime or "Never"
        local boldStyle = { font = { name = ".AppleSystemUIFontBold", size = 13 } }

        local menuTable = {
            { title = "📅 Last Saved: " .. lastTime, disabled = true },
            { title = "-" },
            { title = modeIcon .. " Status: " .. modeName .. " Mode", disabled = true },
            { title = "      ↳ 📺 Screens Detected: " .. currentCount, disabled = true },
            { title = "-" },
            { title = hs.styledtext.new("📸 Save Layout (⇧⌘W)", boldStyle), fn = function() obj.saveLayout(false); obj.resetCountdown() end },
            { title = hs.styledtext.new("🔄 Restore Layout (⇧⌘R)", boldStyle), fn = function() obj.restoreLayout() end },
            { title = hs.styledtext.new("🚀 Force Rescue: All Windows to Laptop (⇧⌘⌃L)", boldStyle), fn = function() obj.isRescued = false; obj.rescueWindowsToLaptop() end },
            { title = "-" },
            { title = "📦 Windows in Saved File:", disabled = true }
        }

        if data and data.windows then
            local seen, appNames = {}, {}
            for _, w in ipairs(data.windows) do 
                if not seen[w.appName] then 
                    table.insert(appNames, w.appName .. ", ") 
                    seen[w.appName] = true 
                end 
            end
            local appString = table.concat(appNames):gsub(", $", "")
            for _, line in ipairs(wrapText(appString, 45)) do 
                table.insert(menuTable, { title = line, disabled = true }) 
            end
        end

        timerMenu:setMenu(menuTable)
    end
end

function obj.resetCountdown()
    saveCountdown = obj.saveInterval
    updateMenu()
end

-- ==========================================
-- WATCHERS & HOTKEYS
-- ==========================================

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

obj.screenWatcher = hs.screen.watcher.new(function()
    local currentCount = #hs.screen.allScreens()
    if currentCount ~= obj.lastScreenCount then
        obj.isTransitioning = true 
        if currentCount > 1 then
            obj.isRescued = false 
            hs.timer.doAfter(4, function() obj.restoreLayout(); obj.isTransitioning = false end)
        else
            hs.timer.doAfter(2, function() obj.rescueWindowsToLaptop(); obj.isTransitioning = false end)
        end
        obj.lastScreenCount = currentCount
        updateMenu()
    end
end):start()

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

updateMenu()
return obj