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

-- ==========================================
-- CONFIGURATION
-- ==========================================
selector.storageDir = os.getenv("HOME") .. "/.hammerspoon/layouts/"
selector.hotkeys = {} 

if not hs.fs.attributes(selector.storageDir) then hs.fs.mkdir(selector.storageDir) end

local stubbornAppsList = {
    ["Gemini"] = true,
    ["AFFiNE"] = true,
    ["Terminal"] = true,
    ["System Settings"] = 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
}

-- ==========================================
-- INTERNAL LOGIC
-- ==========================================

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

local function captureCurrentLayout()
    local screens = hs.screen.allScreens()
    local layout = { 
        saveTime = os.date("%Y-%m-%d %H:%M:%S"), 
        screenCount = #screens,
        mode = (#screens > 1) and "Docked" or "Laptop",
        windows = {} 
    }
    
    for _, win in ipairs(hs.window.allWindows()) do
        local app = win:application()
        local screen = win:screen()
        local frame = win:frame()
        if app and screen and win:isVisible() 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(), 
                    screenName = screen:name(), 
                    x = math.floor(frame.x), 
                    y = math.floor(frame.y),
                    w = math.floor(frame.w), 
                    h = math.floor(frame.h)
                })
            end
        end
    end
    return layout
end

local function executeRestore(filePath, layoutName)
    local data = hs.json.read(filePath)
    if not data then return end
    local windowList = data.windows or data
    local launchedAny = false

    for _, winData in ipairs(windowList) do
        local app = hs.application.get(winData.bundleID) or hs.application.get(winData.appName)
        if not app then
            if winData.bundleID and winData.bundleID ~= "" then 
                hs.application.launchOrFocusByBundleID(winData.bundleID)
            else 
                hs.application.launchOrFocus(winData.appName) 
            end
            launchedAny = true
        end
    end

    local function moveWindows()
        for _, winData in ipairs(windowList) do
            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
                            -- FIX: Target window by title to prevent stacking
                            local safeTitle = win:title():gsub('"', '\\"')
                            local script = string.format([[
                                tell application "System Events" to tell (first process whose unix id is %d)
                                    try
                                        set targetWin to (first window whose name is "%s")
                                        set position of targetWin to {%d, %d}
                                        set size of targetWin to {%d, %d}
                                    end try
                                end tell
                            ]], pid, safeTitle, 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
        hs.alert.show("Restored: " .. layoutName, 1.5)
    end

    if launchedAny then
        hs.alert.show("Launching apps...", 3.5)
        hs.timer.doAfter(4.5, moveWindows)
    else
        moveWindows()
    end
end

local function minimizeAll()
    local allWindows = hs.window.filter.new():getWindows()
    for _, win in ipairs(allWindows) do
        local app = win:application()
        if app and not ignoreListItems[app:name()] then
            win:minimize()
        end
    end
    hs.alert.show("All Windows Minimized", 1)
end

-- ==========================================
-- MENU BAR & HOTKEYS
-- ==========================================
selector.barItem = hs.menubar.new()

function selector.refreshMenu()
    for _, hk in pairs(selector.hotkeys) do hk:delete() end
    selector.hotkeys = {}
    
    local screens = hs.screen.allScreens()
    local isDocked = #screens > 1
    
    selector.barItem:setTitle(isDocked and "🖥️" or "💻")

    selector.hotkeys["save"] = hs.hotkey.bind({"shift", "ctrl", "alt", "cmd"}, "S", function()
        local _, name = hs.dialog.textPrompt("New Layout", "Enter name:", "", "Save", "Cancel")
        if name and name ~= "" then
            hs.json.write(captureCurrentLayout(), selector.storageDir .. name .. ".json", true, true)
            selector.refreshMenu()
        end
    end)
    selector.hotkeys["minimize"] = hs.hotkey.bind({"shift", "ctrl", "alt", "cmd"}, "M", minimizeAll)

    local menuTable = {
        { title = "📸 Save New Layout... (⇧⌃⌥⌘S)", fn = function() hs.eventtap.keyStroke({"shift", "ctrl", "alt", "cmd"}, "S") end },
        { title = "📉 Minimize All Windows (⇧⌃⌥⌘M)", fn = minimizeAll },
        { title = "-" }
    }
    
    local files = {}
    local iter, dir_obj = hs.fs.dir(selector.storageDir)
    if iter then for f in iter, dir_obj do if f:find("%.json$") then table.insert(files, f) end end end
    table.sort(files)

    local boldStyle = { font = { name = ".AppleSystemUIFontBold", size = 13 } }

    for i, file in ipairs(files) do
        local path, label = selector.storageDir .. file, file:gsub("%.json$", "")
        local data = hs.json.read(path)
        
        if i <= 9 then selector.hotkeys[i] = hs.hotkey.bind({"ctrl", "alt", "cmd"}, tostring(i), function() executeRestore(path, label) end) end

        table.insert(menuTable, { 
            title = hs.styledtext.new("Layout: " .. label .. (i<=9 and " (⌃⌥⌘"..i..")" or ""), boldStyle), 
            fn = function() executeRestore(path, label) end 
        })
        
        if data and data.mode then
            local contextStr = (data.mode == "Docked") 
                and string.format("      🖥️ Docked (%d Screens)", data.screenCount or 0)
                or "      💻 Laptop Mode"
            table.insert(menuTable, { title = contextStr, disabled = true })
        end

        if data and data.saveTime then table.insert(menuTable, { title = "      📅 Saved: " .. data.saveTime, disabled = true }) end
        
        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
            for _, line in ipairs(wrapText(table.concat(appNames, " "), 50)) do table.insert(menuTable, { title = line:gsub(",$", ""), disabled = true }) end
        end

        table.insert(menuTable, { 
            title = hs.styledtext.new("      🔄 Update " .. label, boldStyle), 
            fn = function() 
                hs.json.write(captureCurrentLayout(), path, true, true)
                selector.refreshMenu()
                hs.alert.show("Updated: " .. label) 
            end 
        })

        table.insert(menuTable, { title = "      🧹 Restore & Cleanup", fn = function() executeRestore(path, label) end })
        table.insert(menuTable, { title = "      📝 Rename", fn = function() 
            local _, n = hs.dialog.textPrompt("Rename", "New name:", label, "Rename", "Cancel")
            if n and n ~= "" then os.rename(path, selector.storageDir .. n .. ".json"); selector.refreshMenu() end
        end })
        table.insert(menuTable, { title = "      🗑️ Delete", fn = function() 
            if hs.dialog.blockAlert("Delete", "Delete "..label.."?", "Delete", "Cancel", "critical") == "Delete" then 
                os.remove(path); selector.refreshMenu() 
            end 
        end })
        table.insert(menuTable, { title = "-" })
    end
    selector.barItem:setMenu(menuTable)
end

selector.screenWatcher = hs.screen.watcher.new(function()
    selector.refreshMenu()
end):start()

selector.refreshMenu()
hs.alert.show("Layout Selector Loaded", 1.5)
return selector