-- save and restore current window layout
-- inspired by:
-- - https://github.com/dedowsdi/.vim/blob/master/autoload/ddd/layout.vim
-- - https://vi.stackexchange.com/a/22545/53081
--
-- added cursor position handling

local M = {}
local layouts = {}
local resize_cmds = {}
local cursor_positions = {}

function M.save(name)
  layouts[name] = vim.fn.winlayout()
  resize_cmds[name] = vim.fn.winrestcmd()
  cursor_positions[name] = {}
  M.add_buf_to_layout(layouts[name], cursor_positions[name])
end

function M.add_buf_to_layout(layout, cursor_position)
  if layout[1] == "leaf" then
    local win_id = layout[2]
    table.insert(layout, vim.fn.winbufnr(win_id))
    table.insert(cursor_position, vim.api.nvim_win_get_cursor(win_id))
  else
    for _, child_layout in ipairs(layout[2]) do
      M.add_buf_to_layout(child_layout)
    end
  end
end

function M.restore(name)
  if not layouts[name] then
    return
  end

  -- Create clean window
  vim.cmd("new")
  vim.cmd("wincmd o")

  -- Recursively restore buffers and cursor positions
  M.apply_layout(layouts[name], cursor_positions[name], name)

  -- Resize
  vim.cmd(resize_cmds[name])
end

function M.apply_layout(layout, cursor_position, name)
  if layout[1] == "leaf" then
    -- Load buffer for leaf
    if vim.fn.bufexists(layout[3]) == 1 then
      vim.cmd(string.format("b %d", layout[3]))
      -- Restore cursor position
      if cursor_position then
        vim.api.nvim_win_set_cursor(vim.fn.win_getid(), cursor_position[1])
      end
    end
  else
    -- Split cols or rows, split n-1 times
    local split_method = layout[1] == "col" and "rightbelow split" or "rightbelow vsplit"
    local wins = { vim.fn.win_getid() }
    for i = 2, #layout[2] do
      vim.cmd(split_method)
      table.insert(wins, vim.fn.win_getid())
    end

    -- Recursive into child windows
    for index, win_id in ipairs(wins) do
      vim.fn.win_gotoid(win_id)
      M.apply_layout(layout[2][index], name)
    end
  end
end

return M