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

local async = utils.async
local await = utils.await
local awaitAll = utils.awaitAll
local exec = utils.exec

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

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

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

local function getAllWindows()
  return exec(
    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

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

-- Build app icons string from app list
local function buildAppIcons(apps)
  local appkeys = {}
  for app in pairs(apps) do
    table.insert(appkeys, app)
  end
  table.sort(appkeys)

  if #appkeys == 0 then return '' end

  local icons = ''
  for _, app in ipairs(appkeys) do
    local lookup = app_icons[app]
    local icon = lookup or app_icons['Default']
    icons = icons .. ' ' .. icon
  end
  return icons
end

-- Fetch and build workspace state
local fetchState = async(function()
  if state.updating then return end
  state.updating = true

  local ok, err = pcall(function()
    local all, visible, windows = table.unpack(awaitAll {
      getAllWorkspaces(),
      getVisibleWorkspaces(),
      getAllWindows(),
    })

    local newstate = { workspaces = {} }

    -- Initialize all workspaces
    for workspaceid, _ in pairs(spaces) do
      newstate.workspaces[workspaceid] = {
        id = workspaceid,
        monitor = 1,
        active = false,
        empty = true,
        apps = {},
      }
    end

    -- Set monitor IDs
    for _, ws in ipairs(all) do
      local id = ws['workspace']
      if newstate.workspaces[id] then newstate.workspaces[id].monitor = getMonitorId(ws) end
    end

    -- Mark visible workspaces as active
    for _, ws in ipairs(visible) do
      local id = ws['workspace']
      if newstate.workspaces[id] then newstate.workspaces[id].active = true end
    end

    -- Collect apps per workspace
    for _, win in ipairs(windows) do
      local id = win['workspace']
      local app = win['app-name']
      if newstate.workspaces[id] then
        newstate.workspaces[id].apps[app] = true
        newstate.workspaces[id].empty = false
      end
    end

    -- Build app icons
    for _, ws in pairs(newstate.workspaces) do
      ws.appicons = buildAppIcons(ws.apps)
    end

    state.workspaces = newstate.workspaces
  end)

  state.updating = false
  if not ok then print('Error fetching state: ' .. tostring(err)) end
end)

-- Sync UI with state
local function syncState()
  sbar.animate('tanh', 10, function()
    for id, ws in pairs(state.workspaces) do
      local shouldShow = not ws.empty or ws.active

      spaces[id]:set {
        drawing = shouldShow,
        display = ws.monitor,
        label = {
          string = ws.appicons,
          highlight = ws.active,
        },
        icon = {
          highlight = ws.active,
        },
      }
      space_paddings[id]:set { drawing = shouldShow }
    end
  end)
end

local function updateStateAndSync()
  fetchState()
  -- Use a small delay to let fetchState complete before syncing
  sbar.exec('sleep 0.1', syncState)
end

-- Handle workspace change event
local function onActiveSpaceChange(env)
  local focused = env.FOCUSED_WORKSPACE
  local prev = env.PREV_WORKSPACE

  -- Immediately update UI for responsiveness
  if spaces[focused] then
    sbar.animate('tanh', 10, function()
      spaces[focused]:set {
        drawing = true,
        icon = { highlight = true },
        label = { highlight = true },
      }
      space_paddings[focused]:set { drawing = true }

      if spaces[prev] then
        spaces[prev]:set {
          icon = { highlight = false },
          label = { highlight = false },
        }
        if state.workspaces[prev] and state.workspaces[prev].empty then
          spaces[prev]:set { drawing = false }
          space_paddings[prev]:set { drawing = false }
        end
      end
    end)
  end

  updateStateAndSync()
end

-- Setup
local setup = async(function()
  local workspaces = await(getAllWorkspaces())

  -- Create space items
  for _, ws in ipairs(workspaces) do
    local id = ws['workspace']
    local display = getMonitorId(ws)

    local space = sbar.add('item', 'space.' .. id, {
      drawing = false,
      updates = 'when_shown',
      display = display,
      icon = {
        string = id,
        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 ' .. id,
    })
    spaces[id] = space

    local padding = sbar.add('space', 'space.padding.' .. space.name, {
      drawing = false,
      updates = 'when_shown',
      display = display,
      script = '',
      width = settings.space_paddings,
    })
    space_paddings[id] = padding
  end

  -- Register events
  sbar.add('event', 'aerospace_workspace_change')

  local observer = sbar.add('item', {
    drawing = false,
    updates = true,
  })

  observer:subscribe('aerospace_workspace_change', onActiveSpaceChange)
  observer:subscribe('space_windows_change', updateStateAndSync)
  observer:subscribe('system_woke', updateStateAndSync)
  observer:subscribe('front_app_switched', updateStateAndSync)

  -- Initial state sync
  updateStateAndSync()
end)

setup()