refactor: create a new buffer module to centralize buffer operation logic

lua/jj/core/buffer.lua

@@ -0,0 +1,300 @@
+--- @class jj.core.buffer
+local M = {}
+
+--- @class jj.core.buffer.opts
+--- @field name? string Buffer name
+--- @field split? "horizontal"|"vertical"|"tab"|"current" Split type (default: "horizontal")
+--- @field size? number Split size in lines/columns
+--- @field modifiable? boolean Whether buffer is modifiable (default: true)
+--- @field filetype? string Filetype to set
+--- @field buftype? string Buffer type (e.g., "nofile", "acwrite", etc. - optional, defaults to scratch buffer)
+--- @field on_exit? fun(buf: number) Callback when buffer is closed
+--- @field keymaps? jj.core.buffer.keymap[] Keymaps to set on the buffer
+
+--- @class jj.core.buffer.keymap
+--- @field modes? string|string[] Modes for the keymap (default: "n")
+--- @field mode? string Alias for modes
+--- @field lhs string Left-hand side of the keymap
+--- @field rhs string|fun() Right-hand side of the keymap (string or function
+--- @field opts? table Additional keymap options
+
+--- @class jj.core.buffer.float_opts
+--- @field width? number Window width (default: 80% of columns)
+--- @field height? number Window height (default: 80% of lines)
+--- @field row? number Window row position (default: centered)
+--- @field col? number Window column position (default: centered)
+--- @field relative? string Relative positioning (default: "editor")
+--- @field style? string Window style (default: "minimal")
+--- @field border? string Border style (default: "rounded")
+--- @field title? string Window title (default: none)
+--- @field title_pos? string Title position (default: "center")
+--- @field enter? boolean Whether to enter the window after creation (default: false)
+--- @field modifiable? boolean Whether buffer is modifiable (default: true)
+--- @field filetype? string Filetype to set
+--- @field buftype? string Buffer type (e.g., "nofile", "acwrite", etc. - optional, defaults to scratch buffer)
+--- @field bufhidden? string Buffer hidden behavior (default: "hide")
+--- @field on_exit? fun(buf: number) Callback when buffer is closed
+--- @field keymaps? jj.core.buffer.keymap[] Keymaps to set on the buffer
+--- @field win_options? table Window-specific options to set
+
+--- Create and configure a new buffer
+--- @param opts jj.core.buffer.opts Buffer configuration options
+--- @return number buf Buffer handle
+--- @return number? win Window handle (nil if using current window)
+function M.create(opts)
+	opts = opts or {}
+
+	local win = nil
+
+	-- Handle window/split creation
+	if opts.split == "vertical" then
+		local width = opts.size or math.floor(vim.o.columns / 2)
+		vim.cmd(string.format("vsplit | vertical resize %d", width))
+		win = vim.api.nvim_get_current_win()
+	elseif opts.split == "tab" then
+		vim.cmd("tabnew")
+		win = vim.api.nvim_get_current_win()
+	elseif opts.split == "current" then
+		win = vim.api.nvim_get_current_win()
+	else -- horizontal (default)
+		local height = opts.size or math.floor(vim.o.lines / 2)
+		vim.cmd(string.format("split | resize %d", height))
+		win = vim.api.nvim_get_current_win()
+	end
+
+	-- Create buffer
+	local buf = vim.api.nvim_create_buf(false, true)
+
+	-- Set buffer in window if we created/got a window
+	if win then
+		vim.api.nvim_win_set_buf(win, buf)
+	end
+
+	-- Set buffer name if provided (only if it doesn't already exist)
+	if opts.name then
+		pcall(vim.api.nvim_buf_set_name, buf, opts.name)
+	end
+
+	-- Set buffer options
+	if opts.buftype then
+		vim.bo[buf].buftype = opts.buftype
+	end
+	vim.bo[buf].modifiable = opts.modifiable ~= nil and opts.modifiable or true
+	vim.bo[buf].swapfile = false
+	vim.bo[buf].buflisted = false
+
+	-- Set filetype if provided
+	if opts.filetype then
+		vim.bo[buf].filetype = opts.filetype
+	end
+
+	-- Set keymaps if provided
+	if opts.keymaps then
+		M.set_keymaps(buf, opts.keymaps)
+	end
+
+	-- Set up cleanup autocmd if on_exit callback provided
+	if opts.on_exit then
+		vim.api.nvim_create_autocmd({ "BufWipeout", "BufDelete" }, {
+			buffer = buf,
+			once = true,
+			callback = function()
+				opts.on_exit(buf)
+			end,
+		})
+	end
+
+	return buf, win
+end
+
+--- Create and configure a floating window buffer
+--- @param opts jj.core.buffer.float_opts Floating window configuration options
+--- @return number buf Buffer handle
+--- @return number win Window handle
+function M.create_float(opts)
+	opts = opts or {}
+
+	-- Default config
+	local width = opts.width or math.floor(vim.o.columns * 0.8)
+	local height = opts.height or math.floor(vim.o.lines * 0.8)
+	local row = opts.row or math.floor((vim.o.lines - height) / 2)
+	local col = opts.col or math.floor((vim.o.columns - width) / 2)
+
+	local win_config = {
+		relative = opts.relative or "editor",
+		width = width,
+		height = height,
+		row = row,
+		col = col,
+		style = opts.style or "minimal",
+		border = opts.border or "rounded",
+	}
+
+	-- Add optional title
+	if opts.title then
+		win_config.title = opts.title
+		win_config.title_pos = opts.title_pos or "center"
+	end
+
+	-- Create buffer
+	local buf = vim.api.nvim_create_buf(false, true)
+
+	-- Create floating window
+	local win = vim.api.nvim_open_win(buf, opts.enter or false, win_config)
+
+	-- Set buffer options
+	if opts.buftype then
+		vim.bo[buf].buftype = opts.buftype
+	end
+	vim.bo[buf].modifiable = opts.modifiable ~= nil and opts.modifiable or true
+	vim.bo[buf].swapfile = false
+	vim.bo[buf].buflisted = false
+	if opts.bufhidden then
+		vim.bo[buf].bufhidden = opts.bufhidden
+	end
+
+	-- Set filetype if provided
+	if opts.filetype then
+		vim.bo[buf].filetype = opts.filetype
+	end
+
+	-- Set window options
+	if opts.win_options then
+		for option, value in pairs(opts.win_options) do
+			vim.wo[win][option] = value
+		end
+	end
+
+	-- Set keymaps if provided
+	if opts.keymaps then
+		M.set_keymaps(buf, opts.keymaps)
+	end
+
+	-- Set up cleanup autocmd if on_exit callback provided
+	if opts.on_exit then
+		vim.api.nvim_create_autocmd({ "BufWipeout", "BufDelete" }, {
+			buffer = buf,
+			once = true,
+			callback = function()
+				opts.on_exit(buf)
+			end,
+		})
+	end
+
+	-- Set up auto-resize on VimResized
+	vim.api.nvim_create_autocmd("VimResized", {
+		buffer = buf,
+		callback = function()
+			if not vim.api.nvim_win_is_valid(win) then
+				return true -- Remove autocmd if window is invalid
+			end
+
+			-- Recalculate dimensions
+			local new_width = opts.width or math.floor(vim.o.columns * 0.8)
+			local new_height = opts.height or math.floor(vim.o.lines * 0.8)
+			local new_row = opts.row or math.floor((vim.o.lines - new_height) / 2)
+			local new_col = opts.col or math.floor((vim.o.columns - new_width) / 2)
+
+			-- Update window configuration
+			vim.api.nvim_win_set_config(win, {
+				relative = opts.relative or "editor",
+				width = new_width,
+				height = new_height,
+				row = new_row,
+				col = new_col,
+			})
+		end,
+	})
+
+	return buf, win
+end
+
+--- Close/wipe a buffer safely
+--- @param buf number Buffer handle
+--- @param force? boolean Force close (default: true)
+function M.close(buf, force)
+	if not buf or not vim.api.nvim_buf_is_valid(buf) then
+		return
+	end
+
+	local cmd = force ~= false and "bwipeout!" or "bwipeout"
+	vim.cmd(cmd .. " " .. buf)
+end
+
+--- Add keymaps to a buffer
+--- @param buf number Buffer handle
+--- @param keymaps jj.core.buffer.keymap[] Array of keymap definitions
+function M.set_keymaps(buf, keymaps)
+	if not vim.api.nvim_buf_is_valid(buf) then
+		return
+	end
+
+	for _, keymap in ipairs(keymaps) do
+		local modes = keymap.modes or keymap.mode or "n"
+		local lhs = keymap.lhs or keymap[1]
+		local rhs = keymap.rhs or keymap[2]
+		local opts = vim.tbl_extend("force", {
+			buffer = buf,
+			noremap = true,
+			silent = true,
+		}, keymap.opts or {})
+		vim.keymap.set(modes, lhs, rhs, opts)
+	end
+end
+
+--- Remove keymaps from a buffer
+--- @param buf number Buffer handle
+--- @param keymaps jj.core.buffer.keymap[] Array of keymap definitions with modes and lhs
+function M.remove_keymaps(buf, keymaps)
+	if not vim.api.nvim_buf_is_valid(buf) then
+		return
+	end
+
+	for _, keymap in ipairs(keymaps) do
+		local modes = keymap.modes or keymap.mode or "n"
+		local lhs = keymap.lhs or keymap[1]
+
+		local modes_list = type(modes) == "table" and modes or { modes }
+
+		for _, mode in ipairs(modes_list) do
+			pcall(vim.keymap.del, mode, lhs, { buffer = buf })
+		end
+	end
+end
+
+--- Set buffer as modifiable or not
+--- @param buf number Buffer handle
+--- @param modifiable boolean Whether buffer should be modifiable
+function M.set_modifiable(buf, modifiable)
+	if not vim.api.nvim_buf_is_valid(buf) then
+		return
+	end
+	vim.bo[buf].modifiable = modifiable
+end
+
+--- Stop insert mode if in the given buffer if the cursor is currently in that buffer
+--- @param buf number Buffer handle
+function M.stop_insert(buf)
+	if not vim.api.nvim_buf_is_valid(buf) then
+		return
+	end
+	if vim.api.nvim_get_current_buf() ~= buf then
+		return
+	end
+
+	vim.cmd("stopinsert")
+end
+
+--- Start insert mode in the given buffer if the cursor is currently in that buffer
+--- @param buf number Buffer handle
+function M.start_insert(buf)
+	if not vim.api.nvim_buf_is_valid(buf) then
+		return
+	end
+	if vim.api.nvim_get_current_buf() ~= buf then
+		return
+	end
+	vim.cmd("startinsert")
+end
+
+return M

lua/jj/ui/editor.lua

@@ -1,6 +1,8 @@
 --- @class jj.ui.editor
 local M = {}
 
+local buffer = require("jj.core.buffer")
+
 --- @class jj.ui.editor.highlights
 ---@field added table Highlight settings for added lines
 ---@field modified table Highlight settings for modified lines
@@ -53,27 +55,11 @@ function M.open_editor(initial_text, on_done)
 	-- Initialize highlight groups once
 	init_highlights()
 
-	-- Create a horizontal split at the bottom, half the screen height
-	local height = math.floor(vim.o.lines / 2)
-	vim.cmd(string.format("%dsplit", height))
-
-	-- Create a new unlisted, scratch buffer
-	local buf = vim.api.nvim_create_buf(false, true)
-	vim.api.nvim_buf_set_name(buf, "jujutsu:///DESCRIBE_EDITMSG")
-	vim.api.nvim_buf_set_lines(buf, 0, -1, false, initial_text)
-	vim.api.nvim_win_set_buf(0, buf)
-
-	-- Configure buffer options
-	vim.bo[buf].buftype = "acwrite" -- Allow custom write handling
-	vim.bo[buf].bufhidden = "wipe" -- Automatically wipe buffer when hidden
-	vim.bo[buf].swapfile = false -- Disable swapfile
-	vim.bo[buf].modifiable = true -- Allow editing
-
 	-- Create a namespace for our highlights
 	local ns_id = vim.api.nvim_create_namespace("jj_describe_highlights")
 
 	-- Function to apply highlights to the buffer
-	local function apply_highlights()
+	local function apply_highlights(buf)
 		-- Clear existing highlights
 		vim.api.nvim_buf_clear_namespace(buf, ns_id, 0, -1)
 
@@ -131,13 +117,34 @@ function M.open_editor(initial_text, on_done)
 		end
 	end
 
+	-- Create buffer
+	local buf = buffer.create({
+		name = "jujutsu:///DESCRIBE_EDITMSG",
+		split = "horizontal",
+		size = math.floor(vim.o.lines / 2),
+		buftype = "acwrite",
+		modifiable = true,
+		keymaps = {
+			{ modes = "n", lhs = "q", rhs = "<cmd>close!<CR>", opts = { desc = "Close describe buffer" } },
+			{ modes = "n", lhs = "<Esc>", rhs = "<cmd>close!<CR>", opts = { desc = "Close describe buffer" } },
+		},
+	})
+
+	-- Set buffer content
+	vim.api.nvim_buf_set_lines(buf, 0, -1, false, initial_text)
+
+	-- Set bufhidden after creation
+	vim.bo[buf].bufhidden = "wipe"
+
 	-- Apply highlights initially
-	apply_highlights()
+	apply_highlights(buf)
 
 	-- Reapply highlights when text changes
 	vim.api.nvim_create_autocmd({ "TextChanged", "TextChangedI" }, {
 		buffer = buf,
-		callback = apply_highlights,
+		callback = function()
+			apply_highlights(buf)
+		end,
 	})
 
 	-- Handle :w and :wq commands
@@ -151,22 +158,6 @@ function M.open_editor(initial_text, on_done)
 			vim.bo[buf].modified = false
 		end,
 	})
-
-	-- Add keymap to close the buffer with 'q' in normal mode
-	vim.keymap.set(
-		"n",
-		"q",
-		"<cmd>close!<CR>",
-		{ buffer = buf, noremap = true, silent = true, desc = "Close describe buffer" }
-	)
-
-	-- Add keymap to close the buffer with '<Esc>' in normal mode
-	vim.keymap.set(
-		"n",
-		"<Esc>",
-		"<cmd>close!<CR>",
-		{ buffer = buf, noremap = true, silent = true, desc = "Close describe buffer" }
-	)
 end
 
 return M

lua/jj/ui/terminal.lua

@@ -3,7 +3,10 @@ local M = {}
 
 local utils = require("jj.utils")
 local parser = require("jj.core.parser")
+local buffer = require("jj.core.buffer")
+local runner = require("jj.core.runner")
 
+--- @class jj.ui.terminal.state
 local state = {
 	-- The current terminal buffer for jj commands
 	--- @type integer|nil
@@ -29,28 +32,17 @@ local state = {
 	floating_job_id = nil,
 }
 
+-- Re-export
 M.state = state
 
 --- Close the current terminal buffer if it exists
 function M.close_terminal_buffer()
-	if not state.buf then
-		return
-	elseif state.buf and vim.api.nvim_buf_is_valid(state.buf) then
-		vim.cmd("bwipeout! " .. state.buf)
-	else
-		vim.cmd("close")
-	end
+	buffer.close(state.buf)
 end
 
 --- Close the current terminal buffer if it exists
 local function close_floating_buffer()
-	if not state.floating_buf then
-		return
-	elseif state.floating_buf and vim.api.nvim_buf_is_valid(state.floating_buf) then
-		vim.cmd("bwipeout! " .. state.floating_buf)
-	else
-		vim.cmd("close")
-	end
+	buffer.close(state.floating_buf)
 end
 
 --- Hide the current floating window
@@ -89,8 +81,6 @@ local function handle_status_restore()
 		return
 	end
 
-	local runner = require("jj.core.runner")
-
 	if file_info.is_rename then
 		-- For renamed files, remove the new file and restore the old one from parent revision
 		local rm_cmd = "rm " .. vim.fn.shellescape(file_info.new_path)
@@ -131,7 +121,6 @@ local function handle_log_enter(ignore_immut)
 		return
 	end
 
-	local runner = require("jj.core.runner")
 	-- If we found a revision, edit it.
 
 	-- Build command parts.
@@ -172,8 +161,6 @@ local function handle_log_new(flag, ignore_immut)
 		return
 	end
 
-	local runner = require("jj.core.runner")
-
 	-- Mapping for flag-specific options and messages.
 	local flag_map = {
 		after = {
@@ -235,45 +222,6 @@ local function handle_log_describe()
 	end
 end
 
---- Create a floating window for terminal output
---- @param config table Window configuration options
---- @param enter boolean Whether to enter the window after creation
---- @return integer buf Buffer number
---- @return integer win Window number
-local function create_floating_window(config, enter)
-	local default_config = {
-		width = math.floor(vim.o.columns * 0.8),
-		height = math.floor(vim.o.lines * 0.8),
-		row = math.floor((vim.o.lines - math.floor(vim.o.lines * 0.8)) / 2),
-		col = math.floor((vim.o.columns - math.floor(vim.o.columns * 0.8)) / 2),
-		relative = "editor",
-		style = "minimal",
-		border = "rounded",
-		title = " JJ Diff ",
-		title_pos = "center",
-	}
-
-	local merged_config = vim.tbl_extend("force", default_config, config or {})
-
-	-- Create buffer
-	local buf = vim.api.nvim_create_buf(false, true)
-
-	-- Create window
-	local win = vim.api.nvim_open_win(buf, enter or false, merged_config)
-
-	-- Set buffer options
-	vim.bo[buf].bufhidden = "hide"
-
-	-- Set window options
-	vim.wo[win].wrap = true
-	vim.wo[win].number = false
-	vim.wo[win].relativenumber = false
-	vim.wo[win].cursorline = false
-	vim.wo[win].signcolumn = "no"
-
-	return buf, win
-end
-
 --- Run the command in a floating window
 --- @param cmd string The command to run in the floating window
 function M.run_floating(cmd)
@@ -303,7 +251,32 @@ function M.run_floating(cmd)
 	end
 
 	-- Create new floating buffer
-	local buf, win = create_floating_window({}, true)
+	local buf, win = buffer.create_float({
+		title = " JJ Diff ",
+		title_pos = "center",
+		enter = true,
+		bufhidden = "hide",
+		win_options = {
+			wrap = true,
+			number = false,
+			relativenumber = false,
+			cursorline = false,
+			signcolumn = "no",
+		},
+		on_exit = function(b)
+			if state.floating_buf == b then
+				state.floating_buf = nil
+			end
+			if state.floating_chan then
+				vim.fn.chanclose(state.floating_chan)
+				state.floating_chan = nil
+			end
+			if state.floating_job_id then
+				vim.fn.jobstop(state.floating_job_id)
+				state.floating_job_id = nil
+			end
+		end,
+	})
 	state.floating_buf = buf
 
 	-- Create new terminal channel
@@ -329,7 +302,7 @@ function M.run_floating(cmd)
 			DFT_BACKGROUND = "light",
 		},
 		on_stdout = function(_, data)
-			if not vim.api.nvim_buf_is_valid(state.floating_buf) then
+			if not state.floating_buf or not vim.api.nvim_buf_is_valid(state.floating_buf) then
 				return
 			end
 			local output = table.concat(data, "\n")
@@ -337,75 +310,42 @@ function M.run_floating(cmd)
 		end,
 		on_exit = function(_, _) --[[ exit_code ]]
 			vim.schedule(function()
-				if vim.api.nvim_buf_is_valid(state.floating_buf) then
-					vim.bo[state.floating_buf].modifiable = false
-					if vim.api.nvim_get_current_buf() == state.floating_buf then
-						vim.cmd("stopinsert")
-					end
-				end
+				buffer.set_modifiable(state.floating_buf, false)
+				buffer.stop_insert(state.floating_buf)
 			end)
 		end,
 	})
 
-	-- Set keymaps only if they haven't been set for this buffer
-	if not vim.b[state.floating_buf].jj_keymaps_set then
-		vim.keymap.set(
-			{ "n", "v" },
-			"i",
-			function() end,
-			{ buffer = state.floating_buf, noremap = true, silent = true }
-		)
-		vim.keymap.set(
-			{ "n", "v" },
-			"c",
-			function() end,
-			{ buffer = state.floating_buf, noremap = true, silent = true }
-		)
-		vim.keymap.set(
-			{ "n", "v" },
-			"a",
-			function() end,
-			{ buffer = state.floating_buf, noremap = true, silent = true }
-		)
-		vim.keymap.set(
-			{ "n", "v" },
-			"q",
-			close_floating_buffer,
-			{ buffer = state.floating_buf, noremap = true, silent = true, desc = "Close the floating buffer" }
-		)
-		vim.keymap.set(
-			{ "n" },
-			"<ESC>",
-			hide_floating_window,
-			{ buffer = state.floating_buf, noremap = true, silent = true, desc = "Hide the buffer" }
-		)
-		vim.b[state.floating_buf].jj_keymaps_set = true
+	if jid <= 0 then
+		vim.api.nvim_chan_send(chan, "Failed to start command: " .. cmd .. "\r\n")
+		state.floating_chan = nil
+	else
+		state.floating_job_id = jid
 	end
 
-	-- Set up cleanup autocmd only once per buffer
-	if not vim.b[state.floating_buf].jj_cleanup_set then
-		vim.api.nvim_create_autocmd({ "BufWipeout", "BufDelete" }, {
-			buffer = state.floating_buf,
-			callback = function()
-				if state.floating_buf and vim.api.nvim_buf_is_valid(state.floating_buf) then
-					state.floating_buf = nil
-				end
-				if state.floating_chan then
-					vim.fn.chanclose(chan)
-				end
-				if jid then
-					vim.fn.jobstop(jid)
-				end
-			end,
+	-- Set keymaps only if they haven't been set for this buffer
+	if not vim.b[state.floating_buf].jj_keymaps_set then
+		buffer.set_keymaps(state.floating_buf, {
+			{ modes = { "n", "v" }, lhs = "i", rhs = function() end },
+			{ modes = { "n", "v" }, lhs = "c", rhs = function() end },
+			{ modes = { "n", "v" }, lhs = "a", rhs = function() end },
+			{
+				modes = { "n", "v" },
+				lhs = "q",
+				rhs = close_floating_buffer,
+				opts = { desc = "Close the floating buffer" },
+			},
+			{ modes = "n", lhs = "<ESC>", rhs = hide_floating_window, opts = { desc = "Hide the buffer" } },
 		})
-		vim.b[state.floating_buf].jj_cleanup_set = true
+		vim.b[state.floating_buf].jj_keymaps_set = true
 	end
 end
 
 --- Run a command and show it's output in a terminal buffer
 --- If a previous command already existed it smartly reuses the buffer cleaning the previous output
---- @param cmd string|string[]
-function M.run(cmd)
+--- @param cmd string|string[] The command to run in the terminal buffer
+--- @param keymaps jj.core.buffer.keymap[]|nil Additional keymaps to set for this command buffer
+function M.run(cmd, keymaps)
 	if type(cmd) == "string" then
 		cmd = { cmd }
 	end
@@ -437,12 +377,26 @@ function M.run(cmd)
 	end
 
 	-- Create new terminal buffer
-	local height = math.floor(vim.o.lines / 2)
-	vim.cmd(string.format("%dsplit", height))
+	state.buf = buffer.create({
+		split = "horizontal",
+		size = math.floor(vim.o.lines / 2),
+		on_exit = function(buf)
+			if state.buf == buf then
+				state.buf = nil
+			end
+			if state.chan then
+				vim.fn.chanclose(state.chan)
+				state.chan = nil
+			end
+			if state.job_id then
+				vim.fn.jobstop(state.job_id)
+				state.job_id = nil
+			end
+			state.buf_cmd = nil
+		end,
+	})
 
 	local win = vim.api.nvim_get_current_win()
-	state.buf = vim.api.nvim_create_buf(false, true)
-	vim.api.nvim_win_set_buf(win, state.buf)
 	vim.bo[state.buf].bufhidden = "wipe"
 
 	-- Create new terminal channel
@@ -474,7 +428,7 @@ function M.run(cmd)
 			DFT_BACKGROUND = "light",
 		},
 		on_stdout = function(_, data)
-			if not vim.api.nvim_buf_is_valid(state.buf) or not state.chan then
+			if not state.buf or not vim.api.nvim_buf_is_valid(state.buf) or not state.chan then
 				return
 			end
 			local output = table.concat(data, "\n")
@@ -483,12 +437,8 @@ function M.run(cmd)
 		on_exit = function(_, exit_code)
 			vim.schedule(function()
 				-- Make the buffer not modifiable
-				if vim.api.nvim_buf_is_valid(state.buf) then
-					vim.bo[state.buf].modifiable = false
-					if vim.api.nvim_get_current_buf() == state.buf then
-						vim.cmd("stopinsert")
-					end
-				end
+				buffer.set_modifiable(state.buf, false)
+				buffer.stop_insert(state.buf)
 				-- Store the subcommand on successful exit
 				if exit_code == 0 then
 					state.buf_cmd = cmd[2] or nil
@@ -507,114 +457,123 @@ function M.run(cmd)
 	-- Set keymaps only if they haven't been set for this buffer
 	-- Set base keymaps only if they haven't been set for this buffer yet
 	if not vim.b[state.buf].jj_keymaps_set then
-		vim.keymap.set({ "n", "v" }, "i", function() end, { buffer = state.buf, noremap = true, silent = true })
-		vim.keymap.set({ "n", "v" }, "c", function() end, { buffer = state.buf, noremap = true, silent = true })
-		vim.keymap.set({ "n", "v" }, "a", function() end, { buffer = state.buf, noremap = true, silent = true })
-		vim.keymap.set(
-			{ "n", "v" },
-			"q",
-			M.close_terminal_buffer,
-			{ buffer = state.buf, noremap = true, silent = true, desc = "Close the terminal buffer" }
-		)
-		vim.keymap.set(
-			{ "n" },
-			"<ESC>",
-			M.close_terminal_buffer,
-			{ buffer = state.buf, noremap = true, silent = true, desc = "Close the terminal buffer" }
-		)
+		buffer.set_keymaps(state.buf, {
+			-- Disable insert, command and append modes
+			{ modes = { "n", "v" }, lhs = "i", rhs = function() end },
+			{ modes = { "n", "v" }, lhs = "c", rhs = function() end },
+			{ modes = { "n", "v" }, lhs = "a", rhs = function() end },
+			-- Close terminal buffer
+			{
+				modes = { "n", "v" },
+				lhs = "q",
+				rhs = M.close_terminal_buffer,
+				opts = { desc = "Close the terminal buffer" },
+			},
+			-- Close terminal buffer with ESC
+			{
+				modes = "n",
+				lhs = "<ESC>",
+				rhs = M.close_terminal_buffer,
+				opts = { desc = "Close the terminal buffer" },
+			},
+		})
 
 		vim.b[state.buf].jj_keymaps_set = true
 	end
 
 	-- Remove command-specific keymaps from previous runs
 	if vim.b[state.buf].jj_command_keymaps then
-		for _, map in ipairs(vim.b[state.buf].jj_command_keymaps) do
-			local modes = map.modes
-			if type(modes) ~= "table" then
-				modes = { modes }
-			end
-			for _, mode in ipairs(modes) do
-				pcall(vim.keymap.del, mode, map.lhs, { buffer = state.buf })
-			end
-		end
+		buffer.remove_keymaps(state.buf, vim.b[state.buf].jj_command_keymaps)
 		vim.b[state.buf].jj_command_keymaps = nil
 	end
 
 	-- Add command-specific keymaps for jj buffers
 	local new_command_keymaps = {}
-	local function register_command_keymap(modes, lhs, rhs, opts)
-		local normalized_modes = type(modes) == "table" and vim.deepcopy(modes) or { modes }
-		opts = opts or {}
-		opts.buffer = state.buf
-		if opts.noremap == nil then
-			opts.noremap = true
+
+	-- Append the given keymaps
+	if keymaps and #keymaps > 0 then
+		for _, km in ipairs(keymaps) do
+			table.insert(new_command_keymaps, km)
 		end
-		if opts.silent == nil then
-			opts.silent = true
-		end
-		vim.keymap.set(modes, lhs, rhs, opts)
-		table.insert(new_command_keymaps, { modes = normalized_modes, lhs = lhs })
 	end
 
 	-- Add Enter key mapping for status buffers to open files
 	if cmd[2] == "st" or cmd[2] == "status" then
-		register_command_keymap({ "n" }, "<CR>", handle_status_enter, { desc = "Open file under cursor" })
-		register_command_keymap({ "n" }, "X", handle_status_restore, { desc = "Restore file under cursor" })
+		new_command_keymaps = {
+			{ modes = "n", lhs = "<CR>", rhs = handle_status_enter, opts = { desc = "Open file under cursor" } },
+			{ modes = "n", lhs = "X", rhs = handle_status_restore, opts = { desc = "Restore file under cursor" } },
+		}
 	elseif cmd[2] == "log" then
-		-- Edit
-		register_command_keymap({ "n" }, "<CR>", function()
-			handle_log_enter(false)
-		end, { desc = "Edit change under cursor" })
-		register_command_keymap({ "n" }, "<S-CR>", function()
-			handle_log_enter(true)
-		end, { desc = "Edit change under cursor ignoring immutability" })
-		-- Diff
-		register_command_keymap({ "n" }, "d", handle_log_diff, { desc = "Diff change under cursor" })
-		-- New
-		register_command_keymap({ "n" }, "n", handle_log_new, { desc = "New change off the change under cursor" })
-		register_command_keymap({ "n" }, "<C-n>", function()
-			handle_log_new("after")
-		end, { desc = "New change after the change under cursor" })
-		register_command_keymap({ "n" }, "<S-n>", function()
-			handle_log_new("after", true)
-		end, { desc = "New change after the change under cursor ignoring immutability" })
-		-- Undo/Redo
-		register_command_keymap({ "n" }, "u", function()
-			require("jj.cmd").undo()
-		end, { desc = "Undo last operation" })
-		register_command_keymap({ "n" }, "r", function()
-			require("jj.cmd").redo()
-		end, { desc = "Redo last operation" })
-		register_command_keymap({ "n" }, "D", handle_log_describe, { desc = "Describe change under cursor" })
+		new_command_keymaps = {
+			-- Edit
+			{
+				modes = "n",
+				lhs = "<CR>",
+				rhs = function()
+					handle_log_enter(false)
+				end,
+				opts = { desc = "Edit change under cursor" },
+			},
+			{
+				modes = "n",
+				lhs = "<S-CR>",
+				rhs = function()
+					handle_log_enter(true)
+				end,
+				opts = { desc = "Edit change under cursor ignoring immutability" },
+			},
+			-- Diff
+			{ modes = "n", lhs = "d", rhs = handle_log_diff, opts = { desc = "Diff change under cursor" } },
+			-- New
+			{
+				modes = "n",
+				lhs = "n",
+				rhs = handle_log_new,
+				opts = { desc = "New change off the change under cursor" },
+			},
+			{
+				modes = "n",
+				lhs = "<C-n>",
+				rhs = function()
+					handle_log_new("after")
+				end,
+				opts = { desc = "New change after the change under cursor" },
+			},
+			{
+				modes = "n",
+				lhs = "<S-n>",
+				rhs = function()
+					handle_log_new("after", true)
+				end,
+				opts = { desc = "New change after the change under cursor ignoring immutability" },
+			},
+			-- Undo/Redo
+			{
+				modes = "n",
+				lhs = "u",
+				rhs = function()
+					require("jj.cmd").undo()
+				end,
+				opts = { desc = "Undo last operation" },
+			},
+			{
+				modes = "n",
+				lhs = "r",
+				rhs = function()
+					require("jj.cmd").redo()
+				end,
+				opts = { desc = "Redo last operation" },
+			},
+			{ modes = "n", lhs = "D", rhs = handle_log_describe, opts = { desc = "Describe change under cursor" } },
+		}
 	end
 
 	if #new_command_keymaps > 0 then
+		buffer.set_keymaps(state.buf, new_command_keymaps)
 		vim.b[state.buf].jj_command_keymaps = new_command_keymaps
 	end
 
 	vim.cmd("stopinsert")
-
-	-- Set up cleanup autocmd only once per buffer
-	if not vim.b[state.buf].jj_cleanup_set then
-		vim.api.nvim_create_autocmd({ "BufWipeout", "BufDelete" }, {
-			buffer = state.buf,
-			callback = function()
-				if state.buf and vim.api.nvim_buf_is_valid(state.buf) then
-					state.buf = nil
-				end
-				if state.chan then
-					vim.fn.chanclose(state.chan)
-					state.chan = nil
-				end
-				if state.job_id then
-					vim.fn.jobstop(state.job_id)
-					state.job_id = nil
-				end
-				state.buf_cmd = nil
-			end,
-		})
-		vim.b[state.buf].jj_cleanup_set = true
-	end
 end
 
 return M