--
-- 50-session.lua
--
-- * Saves, loads, deletes and restarts Vim sessions stored under stdpath("state")/sessions.
-- * Augments :mksession with a JSON sidecar (dotfiles.session) for state it can't restore on its
--   own: tab names and man:// windows.
--
-- User commands:
--   `SessionSave`: save the session
--   `SessionLoad`: load the session
--   `SessionDelete`: delete the session
--   `SessionRestart`: save the session, then restart Neovim into it
--   `SessionExitSave`: save the session and quit
--   `SessionExitNoSave`: quit without saving the session
--

local session = require("dotfiles.session")

vim.opt.sessionoptions:remove("folds")

local session_dir = vim.fn.stdpath("state") .. "/sessions"
local session_default = session_dir .. "/default.vim"
if not vim.uv.fs_stat(session_dir) then
	vim.uv.fs_mkdir(session_dir, tonumber("755", 8))
	vim.notify("Sessions save directory created at " .. session_dir, vim.log.levels.INFO)
end

----------------------------------------------------------------------------------------------------
-- Sidecar providers
----------------------------------------------------------------------------------------------------

-- Capture each tab's t:tabname (:mksession drops tab-local variables). Stored positionally because
-- :mksession recreates tabs in their original order, making the index stable across reloads.
local function save_tabnames()
	local names = {}
	for i, tp in ipairs(vim.api.nvim_list_tabpages()) do
		-- pcall: nvim_tabpage_get_var errors when the variable isn't set on that tab
		local ok, name = pcall(vim.api.nvim_tabpage_get_var, tp, "tabname")
		names[i] = (ok and type(name) == "string") and name or ""
	end
	return names
end

-- Reapply the saved names to the recreated tabs by position, then redraw so they show at once.
local function restore_tabnames(names)
	for i, tp in ipairs(vim.api.nvim_list_tabpages()) do
		if names[i] and names[i] ~= "" then
			vim.api.nvim_tabpage_set_var(tp, "tabname", names[i])
		end
	end
	vim.cmd.redrawtabline()
end

-- man:// windows: :mksession can't restore them (they're :buftype=nofile, not backed by a file), so
-- we record them ourselves, nested as tab index -> window number -> { buffer name, cursor line }.
-- Both indices are positional for the same reason as tab names above: :mksession recreates tabs and
-- their windows in order, so the ordinals are stable across reloads. Nesting by tab matters because
-- window numbers restart at 1 in each tabpage and would otherwise collide. The indices are
-- stringified so the sparse table encodes as a JSON object, not a null-padded array.
--
-- Limitation: the cursor line is an index into the *rendered* man page, whose line wrapping depends
-- on MANWIDTH. If MANWIDTH differs on load, the page re-wraps and the saved line points elsewhere.
-- (It never changes in this setup, but the dependency is real.)
local function save_manpages()
	local pages = {}
	for ti, tp in ipairs(vim.api.nvim_list_tabpages()) do
		for _, win in ipairs(vim.api.nvim_tabpage_list_wins(tp)) do
			local name = vim.api.nvim_buf_get_name(vim.api.nvim_win_get_buf(win))
			if name:match("^man://") then
				local tk = tostring(ti)
				pages[tk] = pages[tk] or {}
				pages[tk][tostring(vim.api.nvim_win_get_number(win))] =
					{ name = name, line = vim.api.nvim_win_get_cursor(win)[1] }
			end
		end
	end
	return pages
end

-- Reopen the man:// buffers recorded by save_manpages(). Keys come back from JSON as strings, hence
-- the tonumber(). :edit man://... re-renders the page via the man plugin's BufReadCmd; we run it
-- window-scoped so the layout restored by :mksession is left untouched, then clamp the saved cursor
-- line to the (possibly re-wrapped) page.
local function restore_manpages(pages)
	local tabpages = vim.api.nvim_list_tabpages()
	for tk, wins in pairs(pages) do
		local tp = tabpages[tonumber(tk)]
		if tp then
			local by_number = {}
			for _, win in ipairs(vim.api.nvim_tabpage_list_wins(tp)) do
				by_number[vim.api.nvim_win_get_number(win)] = win
			end
			for wk, info in pairs(wins) do
				local win = by_number[tonumber(wk)]
				if win then
					vim.api.nvim_win_call(win, function()
						vim.cmd.edit({ args = { info.name } })
					end)
					local buf = vim.api.nvim_win_get_buf(win)
					local line =
						math.max(1, math.min(info.line or 1, vim.api.nvim_buf_line_count(buf)))
					vim.api.nvim_win_set_cursor(win, { line, 0 })
				end
			end
		end
	end
end

session.register("tabnames", save_tabnames, restore_tabnames)
session.register("manpages", save_manpages, restore_manpages)

----------------------------------------------------------------------------------------------------
-- Session operations
----------------------------------------------------------------------------------------------------

local function save_session(path)
	vim.cmd.mksession({ path, bang = true })
	session.write(path)
end

local function load_session(path)
	vim.cmd.source(path)
	session.read(path)
end

local function reload_session(path)
	save_session(path)
	vim.cmd.restart({ args = { "+qall", "SessionLoad", path } })
end

local function delete_session(path)
	vim.fs.rm(path)
	session.remove(path)
end

-- Complete session names (without the .vim extension) from the sessions directory.
local function session_completefunc(arg_lead, _, _)
	local completions = {}
	for path in vim.fs.dir(session_dir) do
		if string.match(path, "^" .. arg_lead) and string.match(path, ".vim$") then
			completions[#completions + 1] = path:sub(1, -5)
		end
	end
	return completions
end

-- Resolve a session name (or empty for the default) to an absolute `.vim` path, then run `op` on it.
local function session_op(base, op)
	local path = #base > 0 and base or session_default
	if not string.match(path, "^" .. session_dir) then
		path = session_dir .. "/" .. path
	end
	if not string.match(path, "%.vim$") then
		path = path .. ".vim"
	end
	op(path)
end

----------------------------------------------------------------------------------------------------
-- User commands
----------------------------------------------------------------------------------------------------

local function cmd_session_save(ev)
	session_op(ev.args, save_session)
end

local function cmd_session_load(ev)
	session_op(ev.args, load_session)
end

local function cmd_session_delete(ev)
	session_op(ev.args, delete_session)
end

local function cmd_session_restart(ev)
	session_op(ev.args, reload_session)
end

local function cmd_session_exit_save(ev)
	session_op(ev.args, save_session)
	vim.cmd.qall()
end

local function cmd_session_exit_no_save()
	vim.cmd.qall()
end

vim.api.nvim_create_user_command(
	"SessionSave",
	cmd_session_save,
	{ desc = "Save session", nargs = "?", complete = session_completefunc }
)
vim.api.nvim_create_user_command(
	"SessionLoad",
	cmd_session_load,
	{ desc = "Load session", nargs = "?", complete = session_completefunc }
)
vim.api.nvim_create_user_command(
	"SessionDelete",
	cmd_session_delete,
	{ desc = "Delete session", nargs = "?", complete = session_completefunc }
)
vim.api.nvim_create_user_command(
	"SessionRestart",
	cmd_session_restart,
	{ desc = "Reload session", nargs = "?", complete = session_completefunc }
)
vim.api.nvim_create_user_command(
	"SessionExitSave",
	cmd_session_exit_save,
	{ desc = "Save session and exit", nargs = "?", complete = session_completefunc }
)
vim.api.nvim_create_user_command(
	"SessionExitNoSave",
	cmd_session_exit_no_save,
	{ desc = "Exit without saving session", nargs = "?", complete = session_completefunc }
)