local wez = require 'wezterm'
local act = wez.action
local instances = require 'bugabinga.instances'

local names = {
  aider = 'Aider',
  claude = 'Claude',
  codex = 'Codex',
  copilot = 'Copilot',
  gemini = 'Gemini',
  opencode = 'OpenCode',
  pi = 'Pi',
}

local function clean( value )
  if not value or value == '' then return nil end
  value = value:gsub( '^%s+', '' ):gsub( '%s+$', '' )
  return value ~= '' and value or nil
end

local function basename( value )
  value = clean( value )
  if not value then return nil end
  value = value:gsub( '\\', '/' )
  return value:match '([^/]+)$' or value
end

local function program( user_vars )
  local command = clean( user_vars and user_vars.WEZTERM_PROG )
  if not command then return nil end
  command = command:gsub( '^%^', '' )
  local executable = basename( command:match '^([^%s]+)' )
  return executable and executable:lower():gsub( '%.exe$', '' ):gsub( '%.cmd$', '' ):gsub( '%.bat$', '' ) or nil
end

local function identity( user_vars )
  local explicit = clean( user_vars and user_vars.agent )
  if explicit and names[explicit:lower()] then return names[explicit:lower()] end
  return names[program( user_vars )]
end

local function detect( title, process, user_vars )
  local explicit = identity( user_vars )
  if explicit then return explicit end

  local executable = basename( process )
  if executable then
    executable = executable:lower():gsub( '%.exe$', '' )
    if names[executable] then return names[executable] end
  end

  if clean( user_vars and user_vars.WEZTERM_STATE ) == 'idle' then return nil end

  title = clean( title )
  if not title then return nil end
  if title:sub( 1, #'π' ) == 'π' then return 'Pi' end

  local lower = title:lower()
  for executable_name, display_name in pairs( names ) do
    if lower:match( '^' .. executable_name .. '%f[%W]' ) then return display_name end
  end

  return nil
end

local function state( progress )
  if progress == 'Indeterminate' then return 'working', 1 end
  if type( progress ) == 'table' and progress.Error ~= nil then return 'error', 0 end
  if type( progress ) == 'table' and progress.Percentage ~= nil then return 'working', 1 end
  return 'idle', 2
end

local function badge( progress )
  local progress_state = state( progress )
  if progress_state == 'error' then return '✗', progress_state end
  if progress_state == 'working' then return '◉', progress_state end
  return nil, progress_state
end

local function call( object, method )
  local fn = object and object[method]
  if type( fn ) ~= 'function' then return nil end
  local ok, value = pcall( fn, object )
  return ok and value or nil
end

local function sort_entries( entries )
  table.sort( entries, function ( left, right )
    if left.priority ~= right.priority then return left.priority < right.priority end
    if left.workspace ~= right.workspace then return left.workspace:lower() < right.workspace:lower() end
    if left.agent ~= right.agent then return left.agent < right.agent end
    return left.id < right.id
  end )
  return entries
end

local function scan( windows )
  local panes = {}

  for _, mux_window in ipairs( windows or {} ) do
    local workspace = call( mux_window, 'get_workspace' ) or 'default'
    for _, tab in ipairs( call( mux_window, 'tabs' ) or {} ) do
      for _, pane_info in ipairs( call( tab, 'panes_with_info' ) or {} ) do
        local pane = pane_info.pane
        local pane_id = call( pane, 'pane_id' ) or pane_info.pane_id
        local title = call( pane, 'get_title' ) or ''
        local process = call( pane, 'get_foreground_process_name' )
        local agent = detect( title, process, call( pane, 'get_user_vars' ) )
        local progress_state, priority
        if agent then progress_state, priority = state( call( pane, 'get_progress' ) ) end
        table.insert( panes, {
          id = tostring( pane_id ),
          pane_id = pane_id,
          pane = pane,
          workspace = workspace,
          title = title,
          process = process,
          agent = agent,
          state = progress_state,
          priority = priority,
        } )
      end
    end
  end

  return panes
end

local function collect( windows )
  local entries = {}
  for _, pane in ipairs( scan( windows ) ) do
    if pane.agent then table.insert( entries, pane ) end
  end
  return sort_entries( entries )
end

local function from_instances( rows )
  local entries = {}
  for _, row in ipairs( rows or {} ) do
    local agent = row.agent
    if agent then
      local state = row.state or 'unknown'
      local priority = tonumber( row.priority ) or ({ error = 0, working = 1, idle = 2, unknown = 3, })[state] or 3
      table.insert( entries, {
        id = instances.target_id( row.gui_pid, row.pane_id ),
        gui_pid = row.gui_pid,
        pane_id = row.pane_id,
        workspace = row.workspace or 'default',
        title = row.title or '',
        agent = agent,
        state = state,
        priority = priority,
      } )
    end
  end
  return sort_entries( entries )
end

local last_snapshot = nil

local function write_snapshot()
  if not instances.enabled() then return end

  local ok, windows = pcall( wez.mux.all_windows )
  if not ok then return end

  local panes = {}
  for _, pane in ipairs( scan( windows ) ) do
    table.insert( panes, {
      pane_id = pane.pane_id,
      workspace = pane.workspace,
      title = pane.title,
      process = pane.process,
      agent = pane.agent,
      state = pane.state,
      priority = pane.priority,
    } )
  end

  local payload = wez.json_encode { panes = panes, }
  if payload == last_snapshot then return end

  local pid = instances.current_pid()
  local runtime = os.getenv 'XDG_RUNTIME_DIR'
  if not pid or not runtime then return end

  local path = string.format( '%s/wezterm/instances-%d.json', runtime, pid )
  local temporary = path .. '.tmp'
  local file, open_error = io.open( temporary, 'w' )
  if not file then
    wez.log_error( 'failed to open agent snapshot: ' .. tostring( open_error ) )
    return
  end

  local wrote, write_error = file:write( payload )
  file:close()
  if not wrote then
    wez.log_error( 'failed to write agent snapshot: ' .. tostring( write_error ) )
    return
  end

  local renamed, rename_error = os.rename( temporary, path )
  if not renamed then
    wez.log_error( 'failed to publish agent snapshot: ' .. tostring( rename_error ) )
    return
  end
  last_snapshot = payload
end

local function all_entries()
  if instances.enabled() then
    write_snapshot()
    local rows = instances.list()
    if rows then return from_instances( rows ) end
  end

  local ok, windows = pcall( wez.mux.all_windows )
  return collect( ok and windows or {} )
end

local function state_marker( state )
  if state == 'error' then return '✗' end
  if state == 'working' then return '●' end
  if state == 'idle' then return '○' end
  return '?'
end

local function switcher()
  return wez.action_callback( function ( window, pane )
    local entries = all_entries()
    if #entries == 0 then
      window:toast_notification( 'WezTerm', 'No agent panes found', nil, 3000 )
      return
    end

    local targets = {}
    local choices = {}
    for _, entry in ipairs( entries ) do
      targets[entry.id] = entry
      local title = clean( entry.title ) or entry.agent
      table.insert( choices, {
        id = entry.id,
        label = string.format(
          '%s %s %-7s  %s  %s',
          state_marker( entry.state ),
          entry.agent,
          entry.state,
          entry.workspace,
          wez.truncate_right( title, 48 )
        ),
      } )
    end

    window:perform_action(
      act.InputSelector {
        title = 'Agent panes',
        description = 'Failed and working agents first',
        fuzzy_description = 'Agent: ',
        fuzzy = true,
        choices = choices,
        action = wez.action_callback( function ( inner_window, inner_pane, id )
          local target = id and targets[id] or nil
          if not target then return end

          if target.gui_pid then
            if not instances.activate( target.gui_pid, target.pane_id ) then
              inner_window:toast_notification( 'WezTerm', 'Could not activate agent pane', nil, 3000 )
            end
            return
          end

          inner_window:perform_action(
            act.SwitchToWorkspace { name = target.workspace, },
            inner_pane
          )
          target.pane:activate()
        end ),
      },
      pane
    )
  end )
end

return {
  badge = badge,
  identity = identity,
  program = program,
  write_snapshot = write_snapshot,
  switcher = switcher,
}