persisted.nvim/lua/persisted/init.lua

local utils = require("persisted.utils")

local M = {}

local config
local start_args = vim.fn.argc() > 0 or vim.g.started_with_stdin

local e = vim.fn.fnameescape
local uv = vim.uv or vim.loop

---Fire an event
---@param event string
function M.fire(event)
  vim.api.nvim_exec_autocmds("User", { pattern = "Persisted" .. event })
end

---Get the current session for the current working directory and git branch
---@param opts? {branch?: boolean}
---@return string
function M.current(opts)
  opts = opts or {}
  local name = utils.make_fs_safe(vim.fn.getcwd())

  if config.use_git_branch and opts.branch ~= false then
    local branch = M.branch()
    if branch then
      branch = utils.make_fs_safe(branch)
      name = name .. "@@" .. branch
    end
  end

  return config.save_dir .. name .. ".vim"
end

---Automatically load the session for the current dir
---@param opts? { force?: boolean }
function M.autoload(opts)
  opts = opts or {}

  if not opts.force and start_args then
    return
  end

  if config.autoload and M.allowed_dir() then
    M.load({ autoload = true })
  end
end

---Load a session
---@param opts? { last?: boolean, autoload?: boolean, session?: string }
function M.load(opts)
  opts = opts or {}
  local session

  vim.api.nvim_create_autocmd("SessionLoadPost", {
    callback = function()
      if type(config.load_post) == "function" then
        config.load_post()
      end
      M.fire("LoadPost")
      vim.api.nvim_notify("Session Loaded\n" .. session, vim.log.levels.OFF, {})
      return true -- returning deletes autocmd after fired
    end,
  })
  if opts.last then
    session = M.last()
  elseif opts.session then
    session = opts.session
  else
    session = M.current()
    if vim.fn.filereadable(session) == 0 then
      session = M.current({ branch = false })
    end
  end

  if session and vim.fn.filereadable(session) ~= 0 then
    vim.g.persisting_session = not config.follow_cwd and session or nil
    vim.g.persisted_loaded_session = session

    M.fire("LoadPre")
    if type(config.load_pre) == "function" then
      config.load_pre()
    end
    vim.cmd("silent! source " .. e(session))
  elseif opts.autoload and type(config.on_autoload_no_session) == "function" then
    config.on_autoload_no_session()
  end

  if config.autostart and M.allowed_dir() and not start_args then
    M.start()
  end
end

---Start a session
function M.start()
  vim.api.nvim_create_autocmd("VimLeavePre", {
    group = vim.api.nvim_create_augroup("Persisted", { clear = true }),
    callback = function()
      M.save()
    end,
  })

  vim.g.persisting = true
  M.fire("Start")
end

---Stop a session
function M.stop()
  vim.g.persisting = false
  pcall(vim.api.nvim_del_augroup_by_name, "Persisted")
  M.fire("Stop")
end

---Save the session
---@param opts? { force?: boolean, session?: string }
function M.save(opts)
  opts = opts or {}

  -- Do not save the session if should_save evals to false...unless it's forced
  if type(config.should_save) == "function" and not config.should_save() and not opts.force then
    M.fire("SavePost")
    return
  end

  vim.api.nvim_create_autocmd("SessionWritePost", {
    callback = function()
      if type(config.save_post) == "function" then
        config.save_post()
      end
      vim.api.nvim_notify("Session Saved", vim.log.levels.OFF, {})
      M.fire("SavePost")
      return true -- returning true deletes autocmd after fired
    end,
  })

  M.fire("SavePre")
  if type(config.save_pre) == "function" then
    config.save_pre()
  end
  vim.api.nvim_command("wa")
  vim.api.nvim_command("mksession! " .. e(opts.session or vim.g.persisting_session or M.current()))
end

---Delete the current session
---@param opts? { session?: string }
function M.delete(opts)
  opts = opts or {}
  local session = opts.session or M.current()

  if session and uv.fs_stat(session) ~= 0 then
    M.fire("DeletePre")
    vim.schedule(function()
      M.stop()
      vim.fn.delete(vim.fn.expand(session))
    end)
    M.fire("DeletePost")
  end
end

---Get the current Git branch
---@return string?
function M.branch()
  if uv.fs_stat(".git") then
    local branch = vim.fn.systemlist("git branch --show-current")[1]
    return vim.v.shell_error == 0 and branch or nil
  end
end

-- Switch sessions
---@param session_file_path string
function M.switch(session_file_path)
  vim.api.nvim_create_autocmd("User", {
    pattern = "PersistedSavePost",
    callback = function()
      for _, buf in ipairs(vim.api.nvim_list_bufs()) do
        if vim.api.nvim_buf_is_loaded(buf) then
          vim.api.nvim_buf_delete(buf, {})
        end
      end
      M.load({ session = session_file_path })
      return true
    end,
  })

  M.save({ session = vim.g.persisted_loaded_session })
end

---Select a session to load
function M.select()
  ---@type { session: string, dir: string, branch?: string }[]
  local items = {}
  local found = {} ---@type table<string, boolean>
  for _, session in ipairs(M.list()) do
    if uv.fs_stat(session) then
      local file = session:sub(#M.config.save_dir + 1, -5)
      local dir, branch = unpack(vim.split(file, "@@", { plain = true }))
      dir = dir:gsub("%%", "/")
      if jit.os:find("Windows") then
        dir = dir:gsub("^(%w)/", "%1:/")
      end
      if not found[dir .. (branch or "")] then
        found[dir .. (branch or "")] = true
        items[#items + 1] = { session = session, dir = dir, branch = branch }
      end
    end
  end
  vim.ui.select(items, {
    prompt = "Select a session: ",
    format_item = function(item)
      local name = vim.fn.fnamemodify(item.dir, ":p:~")
      if item.branch then
        name = name .. " (" .. item.branch .. ")"
      end
      return name
    end,
  }, function(item)
    if item then
      M.switch(item.session)
    end
  end)
end

---Determines whether to load, start or stop a session
function M.toggle()
  M.fire("Toggle")
  if vim.g.persisting == nil then
    return M.load()
  end
  if vim.g.persisting then
    return M.stop()
  end
  return M.start()
end

---Allow autosaving and autoloading for the given dir?
---@param opts? {dir?: string}
---@return boolean
function M.allowed_dir(opts)
  opts = opts or {}
  local dir = opts.dir or vim.fn.getcwd()

  return (next(config.allowed_dirs) and utils.dirs_match(dir, config.allowed_dirs) or true)
    and not (next(config.ignored_dirs) and utils.dirs_match(dir, config.ignored_dirs) or false)
end

---Get an ordered list of sessions, sorted by modified time
---@return string[]
function M.list()
  local sessions = vim.fn.glob(config.save_dir .. "*.vim", true, true)

  table.sort(sessions, function(a, b)
    return uv.fs_stat(a).mtime.sec > uv.fs_stat(b).mtime.sec
  end)

  return sessions
end

---Get the last session that was saved
---@return string
function M.last()
  return M.list()[1]
end

---Setup the plugin
---@param opts? table
function M.setup(opts)
  config = vim.tbl_deep_extend("force", require("persisted.config"), opts or {})
  M.config = config

  vim.fn.mkdir(config.save_dir, "p")

  if config.autostart and M.allowed_dir() and vim.g.persisting == nil and not start_args then
    M.start()
  end
end

return M