-- 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.fg,
            highlight_color = colors.red,
          },
          label = {
            padding_right = 12,
            color = colors.fg,
            highlight_color = colors.blue,
            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()