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()
obj.lastSavedTime = "Never"

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
}

-- ==========================================
-- 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 hs.caffeinate.get("displayIdle") or obj.isTransitioning or #hs.screen.allScreens() ~= obj.lastScreenCount then 
        return 
    end
    
    local screens = hs.screen.allScreens()
    local layout = { 
        saveTime = os.date("%H:%M:%S"), 
        screenCount = #screens,
        windows = {} 
    }
    
    for _, win in ipairs(hs.window.allWindows()) do
        local app = win:application()
        local screen = win:screen()
        if app and screen and win:isVisible() and win:frame().w > 0 then
            local appName = app:name() or ""
            if not (ignoreListItems[appName] or appName:find("TheBoringNotch")) then
                table.insert(layout.windows, {
                    appName = appName,
                    bundleID = app:bundleID(),
                    winTitle = win:title(),
                    screenName = screen:name(),
                    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()
    local data = hs.json.read(obj.layoutFile)
    if not data then return end
    obj.isRescued = false 
    local appWindowTracker = {}

    for _, winData in ipairs(data.windows or {}) do
        local app = hs.application.get(winData.bundleID) or hs.application.get(winData.appName)
        if app then
            local appName = app:name()
            appWindowTracker[appName] = appWindowTracker[appName] or 1
            
            local validWins = {}
            for _, w in ipairs(app:allWindows()) do if w:isVisible() then table.insert(validWins, w) end end
            
            local win = validWins[appWindowTracker[appName]]
            if win then
                if win:isMinimized() then win:unminimize() end
                
                local x, y, w, h = winData.x, winData.y, winData.w, winData.h
                local isStubborn = stubbornAppsList[appName] or (app:path() or ""):find("Electron")

                local function moveAction()
                    if isStubborn then
                        local winTarget = (appName == "Gemini") and "window 1" or string.format('first window whose name is "%s"', win:title():gsub('"', '\\"'))
                        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
                        ]], app:pid(), winTarget, x, y, w, h)
                        hs.applescript.applescript(script)
                    else
                        win:setFrame({x=x, y=y, w=w, h=h}, 0)
                    end
                end

                moveAction()
                hs.timer.doAfter(0.5, moveAction)
                hs.timer.doAfter(1.5, moveAction)
                
                appWindowTracker[appName] = appWindowTracker[appName] + 1
            end
        end
    end
    hs.alert.show("Layout Restored", 1.5)
end

function obj.rescueWindowsToLaptop()
    local primary = hs.screen.primaryScreen()
    for _, win in ipairs(hs.window.allWindows()) do
        if win:isVisible() then win:moveToScreen(primary, false, true) end
    end
    obj.isRescued = true
    hs.alert.show("Rescued to Laptop", 1.5)
end

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

function updateMenu()
    if timerMenu then
        local screens = hs.screen.allScreens()
        local minutes, seconds = math.floor(saveCountdown / 60), saveCountdown % 60
        timerMenu:setTitle(string.format("💠 %d:%02d", minutes, seconds))
        
        local data = hs.json.read(obj.layoutFile)
        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 = obj.restoreLayout },
            { title = "🚀 Rescue Windows (Bring to Laptop) (⇧⌘⌃L)", fn = obj.rescueWindowsToLaptop },
            { title = "-" },
            { title = "📦 Saved Apps:", disabled = true }
        }
        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

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

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

obj.screenWatcher = hs.screen.watcher.new(function()
    obj.isTransitioning = true
    hs.timer.doAfter(30, 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