local Promise = require 'promise'
local colors = require 'colors'
local utils = require 'utils'
local settings = require 'settings'
local app_icons = require 'app_icons'

local aerospace_cmd = '/opt/homebrew/bin/aerospace'

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

local function getVisibleWorkspaces()
  return utils.sbarExecP(
    aerospace_cmd .. " 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_cmd
      .. " 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 tonumber(obj['monitor-appkit-nsscreen-screens-id']) or 1 end
  return tonumber(obj['monitor-id']) or 1
end

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

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

  -- Pre-initialize all workspaces from the spaces table
  for workspaceid, _ in pairs(spaces) do
    newstate.workspaces[workspaceid] = {
      id = workspaceid,
      monitor = 1,
      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']
      if newstate.workspaces[workspaceid] then newstate.workspaces[workspaceid]['monitor'] = getMonitorId(workspace) end
    end

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

    for _, window in ipairs(apps) do
      local workspaceid = window['workspace']
      local appname = window['app-name']
      if newstate.workspaces[workspaceid] then
        newstate.workspaces[workspaceid]['apps'][appname] = true
        newstate.workspaces[workspaceid]['empty'] = false
      end
    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
    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)
      :catch(function(err)
        state.updating = false
        print('Error updating state: ' .. tostring(err))
      end)
  end
  return Promise.resolve()
end

local function syncState()
  sbar.animate('tanh', 10, function()
    for workspaceid, workspacestate in pairs(state.workspaces) do
      if not workspacestate['empty'] or workspacestate['active'] then
        spaces[workspaceid]:set {
          drawing = true,
          display = workspacestate['monitor'],
          label = {
            string = workspacestate['appicons'],
            highlight = workspacestate['active'],
          },
          icon = {
            highlight = workspacestate['active'],
          },
        }
        space_paddings[workspaceid]:set { drawing = true }
      else
        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

local function onActiveSpaceChange(env)
  local focused_workspace = env.FOCUSED_WORKSPACE
  local prev_workspace = env.PREV_WORKSPACE

  -- Immediately show the focused workspace with animation
  if spaces[focused_workspace] then
    sbar.animate('tanh', 10, function()
      spaces[focused_workspace]:set {
        drawing = true,
        icon = { highlight = true },
        label = { highlight = true },
      }
      space_paddings[focused_workspace]:set { drawing = true }

      if spaces[prev_workspace] then
        spaces[prev_workspace]:set {
          icon = { highlight = false },
          label = { highlight = false },
        }
        -- Hide previous workspace if it's empty
        if state.workspaces[prev_workspace] and state.workspaces[prev_workspace]['empty'] then
          spaces[prev_workspace]:set { drawing = false }
          space_paddings[prev_workspace]:set { drawing = false }
        end
      end
    end)
  end

  updateStateAndSync()
end

local 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,
          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,
          },
          padding_left = 1,
          padding_right = 1,
          click_script = aerospace_cmd .. ' 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()
      sbar.add('event', 'aerospace_workspace_change')

      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()