-- ~~~~~~~~~~ CONFIGURATION & STATE ~~~~~~~~~~
local bento = {}
local k = _G.hyper or {"alt"}

_G.BentoState = _G.BentoState or {
    gridSize = "3x2",
    autoRestore = true,
    windowTracker = {},
    isSnapping = false  -- LOCK: Prevents the bounce
}

if _G.BentoResources then
    for _, item in ipairs(_G.BentoResources) do 
        if item.delete then item:delete() end
        if item.stop then item:stop() end
    end
end
_G.BentoResources = {}

hs.grid.setGrid(_G.BentoState.gridSize)
hs.grid.setMargins({5, 5})
hs.window.animationDuration = 0 -- Set to 0 to prevent movement lag during snaps

-- ~~~~~~~~~~ UTILITIES ~~~~~~~~~~

local function isWindowInGrid(win)
    local cell = hs.grid.get(win)
    local frame = win:frame()
    local screen = win:screen()
    local gridFrame = hs.grid.getCell(cell, screen)
    
    -- Using a wider 20px buffer to account for macOS window shadows/borders
    return math.abs(frame.x - gridFrame.x) < 20 and
           math.abs(frame.y - gridFrame.y) < 20 and
           math.abs(frame.w - gridFrame.w) < 20 and
           math.abs(frame.h - gridFrame.h) < 20
end

local function manage(action)
    local win = hs.window.focusedWindow()
    if not win then return end
    
    local id = win:id()
    if not _G.BentoState.windowTracker[id] then
        _G.BentoState.windowTracker[id] = win:frame()
    end
    
    -- ENABLE LOCK
    _G.BentoState.isSnapping = true
    
    if action == "ui" then 
        hs.grid.show() 
    else 
        action(win) 
    end

    -- RELEASE LOCK after a short delay
    hs.timer.doAfter(0.3, function()
        _G.BentoState.isSnapping = false
    end)
end

-- ~~~~~~~~~~ DYNAMIC MENU BAR ~~~~~~~~~~

local menu = hs.menubar.new()
table.insert(_G.BentoResources, menu)

local function updateMenu()
    return {
        { title = "Grid: " .. _G.BentoState.gridSize, disabled = true },
        { title = "-" },
        { title = "Standard (2x2)", fn = function() _G.BentoState.gridSize = "2x2"; hs.reload() end, checked = (_G.BentoState.gridSize == "2x2") },
        { title = "Bento (3x2)",    fn = function() _G.BentoState.gridSize = "3x2"; hs.reload() end, checked = (_G.BentoState.gridSize == "3x2") },
        { title = "-" },
        { title = "Reload Bento Only", fn = function() package.loaded["bento"] = nil; require("bento") end },
        { title = "Hard Reload", fn = hs.reload }
    }
end

menu:setTitle("⊞")
menu:setMenu(updateMenu)

-- ~~~~~~~~~~ BINDINGS & WATCHERS ~~~~~~~~~~

local function bind(key, fn)
    local b = hs.hotkey.bind(k, key, fn)
    if b then table.insert(_G.BentoResources, b) end
end

bind("G", function() manage("ui") end)
bind("Left",  function() manage(hs.grid.pushWindowLeft) end)
bind("Right", function() manage(hs.grid.pushWindowRight) end)
bind("Up",    function() manage(hs.grid.pushWindowUp) end)
bind("Down",  function() manage(hs.grid.pushWindowDown) end)

if _G.BentoState.autoRestore then
    local wf = hs.window.filter.new()
    wf:subscribe(hs.window.filter.windowMoved, function(win)
        -- ONLY trigger if we aren't currently snapping via keyboard
        if not _G.BentoState.isSnapping then
            local id = win:id()
            if _G.BentoState.windowTracker[id] and not isWindowInGrid(win) then
                win:setFrame(_G.BentoState.windowTracker[id])
                _G.BentoState.windowTracker[id] = nil
                hs.alert.show("Restored", 0.5)
            end
        end
    end)
    table.insert(_G.BentoResources, wf)
end

hs.alert.show("Bento Ready", 1)
return bento