-- https://github.com/Tnixc/nix-config/blob/main/home/programs/aerospace-sketchybar/sbar-config-libs/items/aerospaces.lua

local Promise = require("promise")
local colors = require("colors")
local utils = require("utils")
local settings = require("settings")
local app_icons = require("app_icons")

local function getAllWorkspaces()
   return utils.sbarExecP("aerospace list-workspaces --all --format '%{workspace}%{monitor-appkit-nsscreen-screens-id}%{monitor-id}%{monitor-name}' --json")
end

local function getVisibleWorkspaces()
   return utils.sbarExecP("aerospace list-workspaces --visible --monitor all --format '%{workspace}%{monitor-appkit-nsscreen-screens-id}%{monitor-id}%{monitor-name}' --json")
end

local function getAllWindows()
	return utils.sbarExecP(
		"aerospace list-windows --all --format '%{app-name}%{window-title}%{workspace}%{monitor-id}%{monitor-appkit-nsscreen-screens-id}%{monitor-name}' --json"
	)
end

local function getMonitorId(obj)
   if obj["monitor-name"] then
		if obj["monitor-name"] == "ZOWIE XL LCD" then
			return "2"
		elseif obj["monitor-name"] == "LG ULTRAWIDE" then
			return "1"
		end
	end
   if obj["monitor-appkit-nsscreen-screens-id"] then
      return obj["monitor-appkit-nsscreen-screens-id"]
   end
   return obj["monitor-id"]
end

local spaces = {}
local space_paddings = {}
local brackets = {}
local state = {
	workspaces = {},
	updating = false,
}

function getState()
   local newstate = {
      workspaces = {},
   }

   for workspaceid, space in pairs(spaces) do
      newstate.workspaces[workspaceid] = {
         monitor = 0,
         active = false,
         empty = true,
         apps = {},
         appicons = "",
      }
   end

   return Promise.all({ getAllWorkspaces(), getVisibleWorkspaces(), getAllWindows() }):thenCall(function(values)
         local all, visible, apps = values[1], values[2], values[3]
         for _, workspace in ipairs(all) do
            local workspaceid = workspace["workspace"]
            newstate.workspaces[workspaceid]["id"] = workspaceid
            newstate.workspaces[workspaceid]["monitor"] = getMonitorId(workspace)
         end

         for _, workspace in ipairs(visible) do
            local workspaceid = workspace["workspace"]
            newstate.workspaces[workspaceid]["active"] = true
         end

         for _, window in ipairs(apps) do
			local workspaceid = window["workspace"]
			local appname = window["app-name"]
			newstate.workspaces[workspaceid]["apps"][appname] = true
			newstate.workspaces[workspaceid]["empty"] = false
         end

         for workspaceid, workspacestate in pairs(newstate.workspaces) do
			local appkeys = {}
			for app in pairs(workspacestate["apps"]) do
				table.insert(appkeys, app)
			end
			table.sort(appkeys)
			if #appkeys > 0 then
				for _, app in ipairs(appkeys) do
					local lookup = app_icons[app]
					local icon = ((lookup == nil) and app_icons["Default"] or lookup)
					workspacestate["appicons"] = workspacestate["appicons"] .. " " .. icon
				end
			else
				workspacestate["appicons"] = ""
			end
            -- print(utils.dump(workspacestate))
		end

         return newstate
   end)
end

local function updateState()
	if not state.updating then
		state.updating = true
		return getState():thenCall(function(newstate)
			state.workspaces = newstate.workspaces
			state.updating = false
		end)
	end
	return Promise.reject("State is already updating")
end


local function highlightSpace(space, space_padding, space_bracket, selected)
    space:set({
          drawing = true,
          icon = { highlight = selected, },
          label = { highlight = selected },
          -- background = { border_color = selected and colors.white or colors.bg2 }
    })
	space_padding:set({
		drawing = true,
	})
    if space_bracket then
       space_bracket:set({
	       -- background = { border_color = selected and colors.grey or colors.bg2 },
       })
    end
end

local function onActiveSpaceChange(env)
	local focused_workspace = env.FOCUSED_WORKSPACE
	local last_workspace = env.PREV_WORKSPACE
	-- print("aerospace_workspace_change from " .. last_workspace .. " to " .. focused_workspace)

	local space = spaces[focused_workspace]
	local space_padding = space_paddings[focused_workspace]

	local prev_space = spaces[last_workspace]
	local prev_space_padding = space_paddings[last_workspace]

	sbar.animate("tanh", 10, function()
		highlightSpace(space, space_padding, nil, true)
		if state.workspaces[last_workspace]["monitor"] == state.workspaces[focused_workspace]["monitor"] then
            highlightSpace(prev_space, prev_space_padding, nil, false)
		end
	end)

	updateState()
end

local function syncState()
   sbar.animate("tanh", 10, function()
       for workspaceid, workspacestate in pairs(state.workspaces) do
          if not workspacestate["empty"] then
             spaces[workspaceid]:set({
                   drawing = true,
                   display = workspacestate["monitor"],
                   -- label = {
                   --    string = workspaceid,
                   --    highlight = workspacestate["active"],
                   -- },
                   -- icon = {
                   --    string = workspaceid,
                   --    color = colors.white,
                   --    highlight = workspacestate["active"],
                   -- },
                   label = {
						string = workspacestate["appicons"],
						highlight = workspacestate["active"],
					},
					icon = {
						highlight = workspacestate["active"],
					},
             })
             space_paddings[workspaceid]:set({ drawing = true })
          else
             -- These should be hidden
             spaces[workspaceid]:set({
                   drawing = false,
                   display = workspacestate["monitor"],
                   label = workspacestate["appicons"],
             })
             space_paddings[workspaceid]:set({ drawing = false })
          end
       end
   end)
end

local function updateStateAndSync()
	return updateState():thenCall(syncState)
end

function setup()
   getAllWorkspaces():thenCall(function(workspaces)
        for _, workspace in ipairs(workspaces) do
            local workspaceid = workspace["workspace"]
            local display = getMonitorId(workspace)

            local space = sbar.add("item", "space." .. workspaceid, {
                drawing = false, -- default to not showing the space -- we'll show if it has windows or is activated
                updates = "when_shown",
                display = display,
                icon = {
                     string = workspaceid,
                     color = colors.white,
                     highlight_color = colors.red,
                },
                label = { 
                   padding_right = 12,
                   color = colors.grey,
                   highlight_color = colors.white,
                   font = "sketchybar-app-font:Regular:14.0",
                   y_offset = -1,
                   -- drawing = false
                },
                padding_left = 1,
                padding_right = 1,
                click_script = "aerospace workspace " .. workspaceid,
            })

            spaces[workspaceid] = space

            local padding = sbar.add("space", "space.padding." .. space.name, {
                drawing = false,
                updates = "when_shown",
                display = display,
                script = "",
                width = settings.space_paddings,
            })
			space_paddings[workspaceid] = padding
        end
   end)
   :thenCall(function()
        local space_window_observer = sbar.add("item", {
            drawing = false,
            updates = true,
        })

        space_window_observer:subscribe("aerospace_workspace_change", onActiveSpaceChange)
        space_window_observer:subscribe("space_windows_change", updateStateAndSync)
        space_window_observer:subscribe("system_woke", updateStateAndSync)
        space_window_observer:subscribe("front_app_switched", updateStateAndSync)
   end)
   :thenCall(updateStateAndSync)
end

setup()