local M = {}

local get_searchdirs = function(dirs)
  dirs = dirs or {}
  if type(dirs) == 'table' and #dirs > 0 then
    return dirs
  elseif type(dirs) == 'string' and (dirs == ':workspace') then
    dirs = vim.lsp.buf.list_workspace_folders()
  elseif type(dirs) == 'string' and not (dirs == '') then
    dirs = vim.split(dirs, ',')
  end
  -- skip asking if we're looking at known filetypes
  local ft = vim.api.nvim_buf_get_option(0, 'filetype')
  if #dirs == 0 and (ft == 'dirbuf' or ft == 'NvimTree' or ft == 'alpha') then
    dirs = { '%:p:h' }
  end
  -- ask user
  if #dirs == 0 then
    vim.ui.input({ prompt = 'Search directories: ', default = '%:p:h' }, function(result)
      dirs = { result }
    end)
  end
  -- default to something reasonable
  if #dirs == 0 then
    dirs = { '%:p:h' }
  end
  -- finally return
  return dirs
end

function onchoice(opts, cb)
  if not cb then
    return
  end
  opts.attach_mappings = function(prompt_bufnr)
    require('telescope.actions').select_default:replace(function()
      local selection = require('telescope.actions.state').get_selected_entry()
      if selection == nil then
        return
      end
      require('telescope.actions').close(prompt_bufnr)
      on_choice(selection.path)
    end)
    return true
  end
end

function search_shim(fn, search_dirs, opts, cb)
  opts = opts or {}
  opts.search_dirs = get_searchdirs(search_dirs)
  opts.prompt_title = vim.fn.expand(table.concat(opts.search_dirs, ','))
  onchoice(opts, cb)
  fn(opts)
end

function M.live_grep(search_dirs, opts, cb)
  search_shim(require('telescope.builtin').live_grep, search_dirs, opts, cb)
end

function M.find_files(search_dirs, opts, cb)
  -- convenience aliases
  if type(search_dirs) == 'string' and search_dirs == 'data' then
    search_dirs = vim.fn.stdpath('data')
  end
  search_shim(require('telescope.builtin').find_files, search_dirs, opts, cb)
end

function M.setup(opts)
  vim.keymap.set('n', ';f', function()
    vim.notify('use gsf', vim.log.levels.ERROR)
  end, { desc = 'find files' })
  vim.keymap.set('n', ';g', function()
    vim.notify('use gsg', vim.log.levels.ERROR)
  end, { desc = 'live grep' })

  vim.keymap.set('n', 'gsg', function()
    M.live_grep(':workspace')
  end, { desc = 'live grep in workspace' })
  vim.keymap.set('n', 'gsG', M.live_grep, { desc = 'live grep ask' })
  vim.keymap.set('n', 'gs-', function()
    M.live_grep('%:p:h')
  end, { desc = 'live grep in parent' })

  vim.keymap.set('n', 'gsf', function()
    M.find_files(':workspace')
  end, { desc = 'find files in workspace' })
  vim.keymap.set('n', 'gsF', M.find_files, { desc = 'find files ask' })

  vim.api.nvim_create_user_command('Grep', function(o)
    M.live_grep(table.concat(o.fargs))
  end, {
    desc = 'live grep (rg)',
    nargs = '*',
    complete = 'dir',
  })

  vim.api.nvim_create_user_command('Find', function(o)
    M.find_files(table.concat(o.fargs))
  end, {
    desc = 'find files (fd)',
    nargs = '*',
    complete = 'dir',
  })
end

return M