--- stolen and adapted from https://github.com/MLFlexer/smart_workspace_switcher.wezterm/blob/main/plugin/init.lua

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

local zoxide_path = 'zoxide'

local function workspace_formatter( label )
  return wezterm.format {
    { Text = '󱂬 : ' .. label, },
  }
end

local run_child_process = function ( args )
  local success, stdout, stderr = wezterm.run_child_process( args )

  if not success then
    wezterm.log_error( "Child process '" .. table.concat( args, ' ' ) .. "' failed with stderr: '" .. stderr .. "'" )
  end
  return stdout or ''
end

local function get_workspace_elements( choice_table, include_instances )
  local workspace_ids = {}
  for _, workspace in ipairs( mux.get_workspace_names() ) do
    table.insert( choice_table, {
      id = workspace,
      label = workspace_formatter( workspace ),
    } )
    workspace_ids[workspace] = true
  end

  if include_instances then
    local current_pid = instances.current_pid()
    local seen = {}
    for _, row in ipairs( instances.list() or {} ) do
      local key = tostring( row.gui_pid ) .. ':' .. (row.workspace or 'default')
      if row.gui_pid ~= current_pid and not seen[key] then
        local id = instances.target_id( row.gui_pid, row.pane_id )
        local title = wezterm.truncate_right( row.title or 'WezTerm', 40 )
        table.insert( choice_table, {
          id = id,
          label = workspace_formatter( (row.workspace or 'default') .. '  ' .. title ),
        } )
        workspace_ids[id] = true
        seen[key] = true
      end
    end
  end

  return choice_table, workspace_ids
end

local function display_path( path )
  if path:sub( 1, #wezterm.home_dir ) == wezterm.home_dir then
    return '~' .. path:sub( #wezterm.home_dir + 1 )
  end
  return path
end

local function add_path_element( choice_table, workspace_ids, path, label )
  if not path or path == '' then return end
  local workspace = display_path( path )
  if not workspace_ids[workspace] then
    table.insert( choice_table, {
      id = path,
      label = label or workspace,
    } )
    workspace_ids[workspace] = true
  end
end

local function get_zoxide_elements( choice_table, workspace_ids )
  local stdout = run_child_process { zoxide_path, 'query', '-l', }

  for _, path in ipairs( wezterm.split_by_newlines( stdout ) ) do
    add_path_element( choice_table, workspace_ids, path )
  end
  return choice_table
end

local function parse_worktrees( output )
  local worktrees = {}
  local current = nil

  local function finish()
    if current and current.path then table.insert( worktrees, current ) end
    current = nil
  end

  for raw_line in (output .. '\n'):gmatch '([^\n]*)\n' do
    local line = raw_line:gsub( '\r$', '' )
    if line == '' then
      finish()
    else
      local path = line:match '^worktree (.+)$'
      local branch = line:match '^branch refs/heads/(.+)$'
      if path then
        finish()
        current = { path = path, branch = 'detached', }
      elseif branch and current then
        current.branch = branch
      elseif line == 'bare' and current then
        current.branch = 'bare'
      end
    end
  end
  finish()

  return worktrees
end

local function pane_cwd( pane )
  local ok, cwd = pcall( pane.get_current_working_dir, pane )
  if not ok or not cwd then return nil end

  local ok_path, path = pcall( function () return cwd.file_path end )
  if not ok_path or not path then
    path = tostring( cwd ):match '^file://[^/]*(/.*)$'
  end
  if not path then return nil end
  if wezterm.target_triple:find( 'windows', 1, true ) and path:match '^/[A-Za-z]:/' then
    path = path:sub( 2 )
  end
  return path
end

local function get_worktree_elements( choice_table, workspace_ids, pane )
  local cwd = pane_cwd( pane )
  if not cwd then return choice_table end

  local success, stdout = wezterm.run_child_process { 'git', '-C', cwd, 'worktree', 'list', '--porcelain', }
  if not success then return choice_table end

  for _, worktree in ipairs( parse_worktrees( stdout ) ) do
    add_path_element(
      choice_table,
      workspace_ids,
      worktree.path,
      'git ' .. worktree.branch .. '  ' .. display_path( worktree.path )
    )
  end
  return choice_table
end

local function is_dir( path )
  -- wezterm.read_dir returns children only if path itself is a directory.
  local ok, result = pcall( wezterm.read_dir, path )
  return ok and result ~= nil
end

local function get_workspace_dir_elements( choice_table, workspace_ids )
  local workspace_dir = wezterm.home_dir .. '/Workspace'
  local ok, entries = pcall( wezterm.read_dir, workspace_dir )
  if not ok or not entries then return choice_table end

  local paths = {}
  for _, entry in ipairs( entries ) do
    if is_dir( entry ) then
      table.insert( paths, entry )
    end
  end
  table.sort( paths )
  for _, path in ipairs( paths ) do
    add_path_element( choice_table, workspace_ids, path )
  end
  return choice_table
end

local function get_choices( pane, include_instances )
  local choices = {}
  local _, workspace_ids = get_workspace_elements( choices, include_instances )
  get_worktree_elements( choices, workspace_ids, pane )
  get_zoxide_elements( choices, workspace_ids )
  get_workspace_dir_elements( choices, workspace_ids )
  return choices
end

local function zoxide_chosen( window, pane, path, label_path )
  local default_prog = window:effective_config().default_prog
  window:perform_action(
    act.SwitchToWorkspace {
      name = label_path,
      spawn = {
        args = default_prog,
        label = 'Workspace: ' .. label_path,
        cwd = path,
      },
    },
    pane
  )
  -- increment zoxide path score
  run_child_process { zoxide_path, 'add', path, }
end

local function workspace_chosen( window, pane, workspace )
  window:perform_action(
    act.SwitchToWorkspace {
      name = workspace,
    },
    pane
  )
end

local function spawn_workspace_window( _window, workspace, cwd )
  -- `mux.spawn_window()` and `wezterm cli spawn --new-window` create mux
  -- windows, which can appear in the workspace HUD without necessarily mapping
  -- a new OS/WM window.  For the workflow-window use-case, force a fresh GUI.
  local args = { os.getenv 'WEZTERM_EXECUTABLE' or 'wezterm', 'start', '--always-new-process', '--workspace', workspace, }
  if cwd then
    table.insert( args, '--cwd' )
    table.insert( args, cwd )
  end

  wezterm.background_child_process( args )
end

local function zoxide_chosen_in_new_window( window, path, label_path )
  spawn_workspace_window( window, label_path, path )
  -- increment zoxide path score
  run_child_process { zoxide_path, 'add', path, }
end

local function workspace_chosen_in_new_window( window, workspace )
  spawn_workspace_window( window, workspace )
end

local function workspace_exists( workspace )
  for _, workspace_name in ipairs( mux.get_workspace_names() ) do
    if workspace == workspace_name then
      return true
    end
  end
  return false
end

local function choose_workspace( title, fuzzy_description, include_instances, handler )
  return wezterm.action_callback( function ( window, pane )
    local choices = get_choices( pane, include_instances )

    window:perform_action(
      act.InputSelector {
        action = wezterm.action_callback( function ( inner_window, inner_pane, id )
          if id then handler( inner_window, inner_pane, id ) end
        end ),
        title = title,
        description = 'Select a workspace and press Enter = accept, Esc = cancel, / = filter',
        fuzzy_description = fuzzy_description,
        choices = choices,
        fuzzy = true,
      },
      pane
    )
  end )
end

local function switch_workspace()
  return choose_workspace( 'Choose Workspace', 'Workspace to switch: ', true, function ( window, pane, id )
    local gui_pid, pane_id = instances.parse_target_id( id )
    if gui_pid then
      instances.activate( gui_pid, pane_id )
    elseif workspace_exists( id ) then
      workspace_chosen( window, pane, id )
    else
      zoxide_chosen( window, pane, id, display_path( id ) )
    end
  end )
end

local function switch_workspace_in_new_window()
  return choose_workspace(
    'Choose Workspace in New Window',
    'Workspace for new window: ',
    false,
    function ( window, pane, id )
      if workspace_exists( id ) then
        workspace_chosen_in_new_window( window, id )
      else
        zoxide_chosen_in_new_window( window, id, display_path( id ) )
      end
    end
  )
end

return {
  switch_workspace = switch_workspace,
  switch_workspace_in_new_window = switch_workspace_in_new_window,
}