dotfiles/.config/nvim/lua/func/worktree.lua

--[[ func/worktree.lua ]]
--
-- Telescope picker over `git worktree list`. Picking an entry cds to that
-- worktree and re-points every open buffer at the same relative path inside it.
--
--   <Leader>gw   pick a worktree
--   :Worktree [path]
--   :Telescope worktree

local M = {}

local uv = vim.uv or vim.loop

local function notify(msg, level)
  vim.notify(msg, level or vim.log.levels.INFO, { title = 'worktree' })
end

local function normalise(path)
  if path == nil or path == '' then
    return nil
  end
  local resolved = uv.fs_realpath(path) or path
  local trimmed = resolved:gsub('/+$', '')
  return trimmed
end

local function contains(root, path)
  return path == root or path:sub(1, #root + 1) == root .. '/'
end

-- git commands run against the current buffer's directory, not just the cwd,
-- so the picker still works after wandering into another repo.
local function context_dir()
  local name = vim.api.nvim_buf_get_name(0)
  if name ~= '' and vim.bo.buftype == '' then
    local dir = vim.fn.fnamemodify(name, ':p:h')
    if vim.fn.isdirectory(dir) == 1 then
      return dir
    end
  end
  return vim.fn.getcwd()
end

local function git(dir, args)
  local cmd = { 'git', '-C', dir }
  vim.list_extend(cmd, args)
  local out = vim.fn.systemlist(cmd)
  if vim.v.shell_error ~= 0 then
    return nil, table.concat(out, ' ')
  end
  return out
end

function M.list(dir)
  dir = dir or context_dir()
  local out, err = git(dir, { 'worktree', 'list', '--porcelain' })
  if not out then
    return nil, err
  end

  local trees, current = {}, nil
  for _, line in ipairs(out) do
    if line:match('^worktree ') then
      current = { path = normalise(line:sub(10)) }
      table.insert(trees, current)
    elseif current then
      if line == 'bare' then
        current.bare = true
      elseif line == 'detached' then
        current.detached = true
      elseif line:match('^HEAD ') then
        current.head = line:sub(6)
      elseif line:match('^branch ') then
        current.branch = line:sub(8):gsub('^refs/heads/', '')
      elseif line:match('^locked') then
        current.locked = true
      end
    end
  end
  return trees
end

function M.label(tree)
  if tree.bare then
    return '(bare)'
  end
  if tree.branch then
    return tree.branch
  end
  if tree.head then
    return '(detached ' .. tree.head:sub(1, 7) .. ')'
  end
  return '(unknown)'
end

local function detect_root(dir)
  local out = git(dir, { 'rev-parse', '--show-toplevel' })
  if out and out[1] and out[1] ~= '' then
    return normalise(out[1])
  end
  return nil
end

local function listed_file_buffers()
  local bufs = {}
  for _, buf in ipairs(vim.api.nvim_list_bufs()) do
    if vim.api.nvim_buf_is_valid(buf) and vim.bo[buf].buflisted and vim.bo[buf].buftype == '' then
      local name = vim.api.nvim_buf_get_name(buf)
      if name ~= '' then
        table.insert(bufs, { buf = buf, name = normalise(name) or name })
      end
    end
  end
  return bufs
end

local function windows_for(buf)
  local wins = {}
  for _, win in ipairs(vim.api.nvim_list_wins()) do
    if vim.api.nvim_win_is_valid(win) and vim.api.nvim_win_get_buf(win) == buf then
      table.insert(wins, win)
    end
  end
  return wins
end

--- Fit `cursor` to `buf`, which may be shorter or empty. Rows are 1-based.
function M.clamp_cursor(buf, cursor)
  local row = math.max(1, math.min(cursor[1], vim.api.nvim_buf_line_count(buf)))
  local line = vim.api.nvim_buf_get_lines(buf, row - 1, row, false)[1] or ''
  return { row, math.max(0, math.min(cursor[2], #line)) }
end

local function restore_cursor(win, buf, cursor)
  pcall(vim.api.nvim_win_set_cursor, win, M.clamp_cursor(buf, cursor))
end

--- Every worktree a buffer could be carried away from, longest path first so a
--- worktree nested inside another one wins the prefix match.
local function source_roots(new_root, from_root)
  if from_root then
    return { from_root }
  end

  local roots = {}
  for _, tree in ipairs(M.list(new_root) or {}) do
    if tree.path and not tree.bare and tree.path ~= new_root then
      table.insert(roots, tree.path)
    end
  end

  if #roots == 0 then
    local fallback = detect_root(context_dir()) or normalise(vim.fn.getcwd())
    if fallback and fallback ~= new_root then
      roots = { fallback }
    end
  end

  table.sort(roots, function(a, b)
    return #a > #b
  end)
  return roots
end

local function root_of(path, roots)
  for _, root in ipairs(roots) do
    if path ~= root and contains(root, path) then
      return root
    end
  end
  return nil
end

local function names(list)
  local out = {}
  for _, item in ipairs(list) do
    table.insert(out, item.rel)
  end
  return out
end

--- Switch the session to `target` and re-point buffers living in its sibling worktrees.
--- Returns a summary table, or nil when the switch was refused.
function M.switch(target, opts)
  opts = opts or {}
  local new_root = normalise(target)
  if not new_root or vim.fn.isdirectory(new_root) == 0 then
    notify('not a directory: ' .. tostring(target), vim.log.levels.ERROR)
    return nil
  end

  local roots = source_roots(new_root, normalise(opts.from_root))
  local bufs = listed_file_buffers()

  local moved, missing, blocked, kept = {}, {}, {}, {}
  for _, entry in ipairs(bufs) do
    local root = root_of(entry.name, roots)
    if root then
      local rel = entry.name:sub(#root + 2)
      local item = { buf = entry.buf, root = root, rel = rel, dest = new_root .. '/' .. rel }
      local exists = vim.fn.filereadable(item.dest) == 1
      if vim.bo[entry.buf].modified then
        table.insert(exists and blocked or kept, item)
      elseif exists then
        table.insert(moved, item)
      else
        table.insert(missing, item)
      end
    end
  end

  local matched = #moved + #missing + #blocked + #kept
  if matched == 0 and normalise(vim.fn.getcwd()) == new_root then
    notify('already in ' .. new_root)
    return { root = new_root, moved = {}, missing = {}, kept = {} }
  end

  if #blocked > 0 then
    notify(
      'unsaved changes would be replaced; write them first: ' .. table.concat(names(blocked), ', '),
      vim.log.levels.ERROR
    )
    return nil
  end

  vim.cmd('cd ' .. vim.fn.fnameescape(new_root))

  for _, item in ipairs(moved) do
    local wins = windows_for(item.buf)
    local cursors = {}
    for _, win in ipairs(wins) do
      cursors[win] = vim.api.nvim_win_get_cursor(win)
    end

    local last
    if wins[1] then
      last = cursors[wins[1]]
    else
      local mark = vim.api.nvim_buf_get_mark(item.buf, '"')
      if mark[1] > 0 then
        last = mark
      end
    end

    local newbuf = vim.fn.bufadd(item.dest)
    vim.fn.bufload(newbuf)
    vim.bo[newbuf].buflisted = true

    -- The `"` mark is what a hidden buffer is re-opened at, so carry it too.
    if last then
      local pos = M.clamp_cursor(newbuf, last)
      pcall(vim.api.nvim_buf_set_mark, newbuf, '"', pos[1], pos[2], {})
    end

    for _, win in ipairs(wins) do
      vim.api.nvim_win_set_buf(win, newbuf)
      restore_cursor(win, newbuf, cursors[win])
    end

    if newbuf ~= item.buf and vim.api.nvim_buf_is_valid(item.buf) then
      pcall(vim.api.nvim_buf_delete, item.buf, { force = false })
    end
  end

  -- Buffers with no counterpart keep pointing at the old worktree; flag them
  -- so it is obvious why they did not follow.
  for _, list in ipairs({ missing, kept }) do
    for _, item in ipairs(list) do
      if vim.api.nvim_buf_is_valid(item.buf) then
        vim.b[item.buf].worktree_orphan = item.root
      end
    end
  end

  local summary = {
    root = new_root,
    sources = roots,
    moved = names(moved),
    missing = names(missing),
    kept = names(kept),
  }

  local msg = ('switched to %s (%d buffer(s) moved)'):format(new_root, #moved)
  local level = vim.log.levels.INFO
  if matched == 0 and #bufs > 0 then
    msg = msg .. ('; %d open buffer(s) live outside this repository'):format(#bufs)
    level = vim.log.levels.WARN
  end
  if #missing > 0 then
    msg = msg .. '; no counterpart there: ' .. table.concat(summary.missing, ', ')
    level = vim.log.levels.WARN
  end
  if #kept > 0 then
    msg = msg .. '; left unsaved on old worktree: ' .. table.concat(summary.kept, ', ')
    level = vim.log.levels.WARN
  end
  notify(msg, level)

  return summary
end

local function selectable(dir)
  local trees, err = M.list(dir)
  if not trees then
    notify('no git worktrees here: ' .. (err or 'not a git repository'), vim.log.levels.ERROR)
    return nil
  end
  local usable = {}
  for _, tree in ipairs(trees) do
    if not tree.bare and tree.path and vim.fn.isdirectory(tree.path) == 1 then
      table.insert(usable, tree)
    end
  end
  if #usable == 0 then
    notify('no checked-out worktrees to switch to', vim.log.levels.WARN)
    return nil
  end
  return usable
end

function M.pick(opts)
  opts = opts or {}
  local usable = selectable(opts.cwd or context_dir())
  if not usable then
    return
  end

  local ok, pickers = pcall(require, 'telescope.pickers')
  if not ok then
    vim.ui.select(usable, {
      prompt = 'Git worktree',
      format_item = function(tree)
        return M.label(tree) .. '  ' .. vim.fn.fnamemodify(tree.path, ':~')
      end,
    }, function(choice)
      if choice then
        M.switch(choice.path)
      end
    end)
    return
  end

  local finders = require('telescope.finders')
  local conf = require('telescope.config').values
  local actions = require('telescope.actions')
  local action_state = require('telescope.actions.state')
  local entry_display = require('telescope.pickers.entry_display')

  local width = 0
  for _, tree in ipairs(usable) do
    width = math.max(width, #M.label(tree))
  end
  local displayer = entry_display.create({
    separator = '  ',
    items = { { width = width }, { remaining = true } },
  })

  pickers
    .new(opts, {
      prompt_title = 'Git Worktrees',
      finder = finders.new_table({
        results = usable,
        entry_maker = function(tree)
          return {
            value = tree,
            path = tree.path,
            ordinal = M.label(tree) .. ' ' .. tree.path,
            display = function(entry)
              return displayer({
                M.label(entry.value),
                vim.fn.fnamemodify(entry.value.path, ':~'),
              })
            end,
          }
        end,
      }),
      sorter = conf.generic_sorter(opts),
      attach_mappings = function(prompt_bufnr)
        actions.select_default:replace(function()
          local entry = action_state.get_selected_entry()
          actions.close(prompt_bufnr)
          if entry then
            M.switch(entry.value.path)
          end
        end)
        return true
      end,
    })
    :find()
end

vim.api.nvim_create_user_command('Worktree', function(cmd)
  if cmd.args ~= '' then
    M.switch(cmd.args)
  else
    M.pick()
  end
end, { nargs = '?', complete = 'dir', desc = 'Switch git worktree and follow open buffers' })

vim.keymap.set('n', '<Leader>gw', function()
  M.pick()
end, { silent = true, desc = 'Git worktree picker' })

pcall(function()
  require('telescope').load_extension('worktree')
end)

return M