claude-baseline-1752125979

lua/nvim-claude/hooks.lua

@@ -6,7 +6,7 @@ M.pre_edit_commit = nil
 
 function M.setup()
   vim.notify('Hooks module loaded', vim.log.levels.DEBUG)
-  
+
   -- Auto-cleanup old Claude commits on startup
   vim.defer_fn(function()
     M.cleanup_old_commits()
@@ -16,62 +16,62 @@ end
 -- Pre-tool-use hook: Create baseline commit only if one doesn't exist
 function M.pre_tool_use_hook()
   vim.notify('Pre-hook called', vim.log.levels.INFO)
-  
+
   -- Debug log
-  local debug_msg = string.format('Pre-hook executing at %s', os.date('%Y-%m-%d %H:%M:%S'))
+  local debug_msg = string.format('Pre-hook executing at %s', os.date '%Y-%m-%d %H:%M:%S')
-  vim.fn.writefile({debug_msg}, '/tmp/claude-pre-hook-debug.log', 'a')
+  vim.fn.writefile({ debug_msg }, '/tmp/claude-pre-hook-debug.log', 'a')
-  
+
-  local utils = require('nvim-claude.utils')
+  local utils = require 'nvim-claude.utils'
-  
+
   -- Check if we're in a git repository
   local git_root = utils.get_project_root()
   if not git_root then
     vim.notify('Not in a git repository', vim.log.levels.WARN)
     return
   end
-  
+
   -- Check if we already have a baseline commit
   local baseline_file = '/tmp/claude-baseline-commit'
   local existing_baseline = utils.read_file(baseline_file)
-  
+
   if existing_baseline and existing_baseline ~= '' then
     existing_baseline = existing_baseline:gsub('%s+', '')
     -- Verify the baseline commit still exists
     local check_cmd = string.format('cd "%s" && git rev-parse --verify %s', git_root, existing_baseline)
     local check_result, check_err = utils.exec(check_cmd)
-    
+
     if not check_err then
       vim.notify('Using existing baseline commit: ' .. existing_baseline, vim.log.levels.INFO)
       return
     end
   end
-  
+
   -- Create new baseline commit
   local timestamp = os.time()
   local commit_msg = string.format('claude-baseline-%d', timestamp)
-  
+
   -- Stage all current changes (including untracked files)
   local add_cmd = string.format('cd "%s" && git add -A', git_root)
   local add_result, add_err = utils.exec(add_cmd)
-  
+
   if add_err then
     vim.notify('Failed to stage changes: ' .. add_err, vim.log.levels.ERROR)
     return
   end
-  
+
   -- Create a temporary commit
   local commit_cmd = string.format('cd "%s" && git commit -m "%s" --allow-empty', git_root, commit_msg)
   local commit_result, commit_err = utils.exec(commit_cmd)
-  
+
-  if commit_err and not commit_err:match('nothing to commit') then
+  if commit_err and not commit_err:match 'nothing to commit' then
     vim.notify('Failed to create baseline commit: ' .. commit_err, vim.log.levels.ERROR)
     return
   end
-  
+
   -- Store the actual commit hash instead of 'HEAD'
   local hash_cmd = string.format('cd "%s" && git rev-parse HEAD', git_root)
   local commit_hash, hash_err = utils.exec(hash_cmd)
-  
+
   if not hash_err and commit_hash and commit_hash ~= '' then
     commit_hash = commit_hash:gsub('%s+', '')
     M.baseline_commit = commit_hash
@@ -80,78 +80,78 @@ function M.pre_tool_use_hook()
     M.baseline_commit = 'HEAD'
     utils.write_file(baseline_file, 'HEAD')
   end
-  
+
   vim.notify('New baseline commit created: ' .. commit_msg, vim.log.levels.INFO)
 end
 
 -- Post-tool-use hook: Create stash of Claude's changes and trigger diff review
 function M.post_tool_use_hook()
   vim.notify('Post-hook called', vim.log.levels.INFO)
-  
+
   -- Debug log
-  local debug_msg = string.format('Post-hook executing at %s', os.date('%Y-%m-%d %H:%M:%S'))
+  local debug_msg = string.format('Post-hook executing at %s', os.date '%Y-%m-%d %H:%M:%S')
-  vim.fn.writefile({debug_msg}, '/tmp/claude-post-hook-debug.log', 'a')
+  vim.fn.writefile({ debug_msg }, '/tmp/claude-post-hook-debug.log', 'a')
-  
+
   vim.notify('Post-tool-use hook triggered', vim.log.levels.INFO)
-  
+
   -- Use vim.schedule to ensure we're in the main thread
   vim.schedule(function()
-    local utils = require('nvim-claude.utils')
+    local utils = require 'nvim-claude.utils'
-    
+
     -- Refresh all buffers to show Claude's changes
-    vim.cmd('checktime')
+    vim.cmd 'checktime'
-    
+
     -- Check if Claude made any changes
     local git_root = utils.get_project_root()
     if not git_root then
       vim.notify('Not in a git repository', vim.log.levels.WARN)
       return
     end
-    
+
     local status_cmd = string.format('cd "%s" && git status --porcelain', git_root)
     local status_result = utils.exec(status_cmd)
-    
+
     if not status_result or status_result == '' then
       vim.notify('No changes detected from Claude', vim.log.levels.INFO)
       return
     end
-    
+
     -- Get list of modified files
     local modified_files = {}
-    for line in status_result:gmatch('[^\n]+') do
+    for line in status_result:gmatch '[^\n]+' do
-      local file = line:match('^.M (.+)$') or line:match('^M. (.+)$') or line:match('^.. (.+)$')
+      local file = line:match '^.M (.+)$' or line:match '^M. (.+)$' or line:match '^.. (.+)$'
       if file then
         table.insert(modified_files, file)
       end
     end
-    
+
     -- Create a stash of Claude's changes (but keep them in working directory)
-    local timestamp = os.date('%Y-%m-%d %H:%M:%S')
+    local timestamp = os.date '%Y-%m-%d %H:%M:%S'
     local stash_msg = string.format('[claude-edit] %s', timestamp)
-    
+
     -- Use git stash create to create stash without removing changes
     local stash_cmd = string.format('cd "%s" && git stash create -u', git_root)
     local stash_hash, stash_err = utils.exec(stash_cmd)
-    
+
     if not stash_err and stash_hash and stash_hash ~= '' then
       -- Store the stash with a message
       stash_hash = stash_hash:gsub('%s+', '') -- trim whitespace
       local store_cmd = string.format('cd "%s" && git stash store -m "%s" %s', git_root, stash_msg, stash_hash)
       utils.exec(store_cmd)
-      
+
       -- Get the baseline commit reference
-      local baseline_ref = utils.read_file('/tmp/claude-baseline-commit')
+      local baseline_ref = utils.read_file '/tmp/claude-baseline-commit'
       if baseline_ref then
         baseline_ref = baseline_ref:gsub('%s+', '') -- trim whitespace
       end
-      
+
       -- Check if any modified files are currently open in buffers
-      local inline_diff = require('nvim-claude.inline-diff')
+      local inline_diff = require 'nvim-claude.inline-diff'
       local opened_inline = false
-      
+
       for _, file in ipairs(modified_files) do
         local full_path = git_root .. '/' .. file
-        
+
         -- Find buffer with this file
         for _, buf in ipairs(vim.api.nvim_list_bufs()) do
           if vim.api.nvim_buf_is_valid(buf) and vim.api.nvim_buf_is_loaded(buf) then
@@ -160,30 +160,32 @@ function M.post_tool_use_hook()
               -- Get the original content (from baseline)
               local baseline_cmd = string.format('cd "%s" && git show %s:%s 2>/dev/null', git_root, baseline_ref or 'HEAD', file)
               local original_content, orig_err = utils.exec(baseline_cmd)
-              
+
               if not orig_err and original_content then
                 -- Get current content
                 local current_lines = vim.api.nvim_buf_get_lines(buf, 0, -1, false)
                 local current_content = table.concat(current_lines, '\n')
-                
+
                 -- Show inline diff
                 inline_diff.show_inline_diff(buf, original_content, current_content)
                 opened_inline = true
-                
+
                 -- Switch to that buffer if it's not the current one
                 if buf ~= vim.api.nvim_get_current_buf() then
                   vim.api.nvim_set_current_buf(buf)
                 end
-                
+
                 break -- Only show inline diff for first matching buffer
               end
             end
           end
         end
-        
+
-        if opened_inline then break end
+        if opened_inline then
+          break
+        end
       end
-      
+
       -- If no inline diff was shown, fall back to regular diff review
       if not opened_inline then
         -- Trigger diff review - show Claude stashes against baseline
@@ -200,91 +202,143 @@ function M.post_tool_use_hook()
   end)
 end
 
+-- Test inline diff manually
+function M.test_inline_diff()
+  vim.notify('Testing inline diff manually...', vim.log.levels.INFO)
+
+  local utils = require 'nvim-claude.utils'
+  local git_root = utils.get_project_root()
+
+  if not git_root then
+    vim.notify('Not in git repository', vim.log.levels.ERROR)
+    return
+  end
+
+  -- Get current buffer
+  local bufnr = vim.api.nvim_get_current_buf()
+  local buf_name = vim.api.nvim_buf_get_name(bufnr)
+
+  if buf_name == '' then
+    vim.notify('Current buffer has no file', vim.log.levels.ERROR)
+    return
+  end
+
+  -- Get relative path
+  local relative_path = buf_name:gsub(git_root .. '/', '')
+  vim.notify('Testing inline diff for: ' .. relative_path, vim.log.levels.INFO)
+
+  -- Get baseline content
+  local baseline_ref = utils.read_file '/tmp/claude-baseline-commit' or 'HEAD'
+  baseline_ref = baseline_ref:gsub('%s+', '')
+
+  local baseline_cmd = string.format('cd "%s" && git show %s:%s 2>/dev/null', git_root, baseline_ref, relative_path)
+  local original_content, orig_err = utils.exec(baseline_cmd)
+
+  if orig_err then
+    vim.notify('Failed to get baseline content: ' .. orig_err, vim.log.levels.ERROR)
+    return
+  end
+
+  -- Get current content
+  local current_lines = vim.api.nvim_buf_get_lines(bufnr, 0, -1, false)
+  local current_content = table.concat(current_lines, '\n')
+
+  -- Show inline diff
+  local inline_diff = require 'nvim-claude.inline-diff'
+  inline_diff.show_inline_diff(bufnr, original_content, current_content)
+end
+
 -- Manual hook testing
 function M.test_hooks()
   vim.notify('=== Testing nvim-claude hooks ===', vim.log.levels.INFO)
-  
+
   -- Test pre-tool-use hook
   vim.notify('1. Testing pre-tool-use hook (creating snapshot)...', vim.log.levels.INFO)
   M.pre_tool_use_hook()
-  
+
   -- Simulate making a change
   vim.notify('2. Make some changes to test files now...', vim.log.levels.INFO)
-  
+
   -- Test post-tool-use hook after a delay
   vim.notify('3. Will trigger post-tool-use hook in 3 seconds...', vim.log.levels.INFO)
-  
+
   vim.defer_fn(function()
     M.post_tool_use_hook()
   end, 3000)
-  
+
   vim.notify('=== Hook testing started - make changes now! ===', vim.log.levels.INFO)
 end
 
 -- Install Claude Code hooks
 function M.install_hooks()
-  local utils = require('nvim-claude.utils')
+  local utils = require 'nvim-claude.utils'
-  
+
   -- Get project root
   local project_root = utils.get_project_root()
   if not project_root then
     vim.notify('Not in a git repository', vim.log.levels.ERROR)
     return
   end
-  
+
   -- Create .claude directory
   local claude_dir = project_root .. '/.claude'
   if not vim.fn.isdirectory(claude_dir) then
     vim.fn.mkdir(claude_dir, 'p')
   end
-  
+
   -- Create hooks configuration
   local server_name = vim.v.servername or 'NVIM'
-  local pre_command = string.format('nvim --headless --server %s --remote-send "<C-\\><C-N>:lua require(\'nvim-claude.hooks\').pre_tool_use_hook()<CR>" 2>/dev/null || true', server_name)
+  local pre_command = string.format(
-  local post_command = string.format('nvim --headless --server %s --remote-send "<C-\\><C-N>:lua require(\'nvim-claude.hooks\').post_tool_use_hook()<CR>" 2>/dev/null || true', server_name)
+    'nvim --headless --server %s --remote-send "<C-\\><C-N>:lua require(\'nvim-claude.hooks\').pre_tool_use_hook()<CR>" 2>/dev/null || true',
-  
+    server_name
+  )
+  local post_command = string.format(
+    'nvim --headless --server %s --remote-send "<C-\\><C-N>:lua require(\'nvim-claude.hooks\').post_tool_use_hook()<CR>" 2>/dev/null || true',
+    server_name
+  )
+
   local hooks_config = {
     hooks = {
       PreToolUse = {
         {
-          matcher = "Edit|Write|MultiEdit",  -- Only match file editing tools
+          matcher = 'Edit|Write|MultiEdit', -- Only match file editing tools
           hooks = {
             {
-              type = "command",
+              type = 'command',
-              command = pre_command
+              command = pre_command,
-            }
+            },
-          }
+          },
-        }
+        },
       },
       PostToolUse = {
         {
-          matcher = "Edit|Write|MultiEdit",  -- Only match file editing tools
+          matcher = 'Edit|Write|MultiEdit', -- Only match file editing tools
           hooks = {
             {
-              type = "command", 
+              type = 'command',
-              command = post_command
+              command = post_command,
-            }
+            },
-          }
+          },
-        }
+        },
-      }
+      },
-    }
+    },
   }
-  
+
   -- Write hooks configuration
   local settings_file = claude_dir .. '/settings.json'
   local success, err = utils.write_json(settings_file, hooks_config)
-  
+
   if success then
     -- Add .claude to gitignore if needed
     local gitignore_path = project_root .. '/.gitignore'
     local gitignore_content = utils.read_file(gitignore_path) or ''
-    
+
-    if not gitignore_content:match('%.claude/') then
+    if not gitignore_content:match '%.claude/' then
       local new_content = gitignore_content .. '\n# Claude Code hooks\n.claude/\n'
       utils.write_file(gitignore_path, new_content)
       vim.notify('Added .claude/ to .gitignore', vim.log.levels.INFO)
     end
-    
+
     vim.notify('Claude Code hooks installed successfully', vim.log.levels.INFO)
     vim.notify('Hooks configuration written to: ' .. settings_file, vim.log.levels.INFO)
   else
@@ -294,17 +348,17 @@ end
 
 -- Uninstall Claude Code hooks
 function M.uninstall_hooks()
-  local utils = require('nvim-claude.utils')
+  local utils = require 'nvim-claude.utils'
-  
+
   -- Get project root
   local project_root = utils.get_project_root()
   if not project_root then
     vim.notify('Not in a git repository', vim.log.levels.ERROR)
     return
   end
-  
+
   local settings_file = project_root .. '/.claude/settings.json'
-  
+
   if vim.fn.filereadable(settings_file) then
     vim.fn.delete(settings_file)
     vim.notify('Claude Code hooks uninstalled', vim.log.levels.INFO)
@@ -318,103 +372,116 @@ function M.setup_commands()
   vim.api.nvim_create_user_command('ClaudeTestHooks', function()
     M.test_hooks()
   end, {
-    desc = 'Test Claude Code hooks'
+    desc = 'Test Claude Code hooks',
   })
-  
+
+  vim.api.nvim_create_user_command('ClaudeTestInlineDiff', function()
+    M.test_inline_diff()
+  end, {
+    desc = 'Test Claude inline diff manually',
+  })
+
+  vim.api.nvim_create_user_command('ClaudeTestKeymap', function()
+    require('nvim-claude.inline-diff').test_keymap()
+  end, {
+    desc = 'Test Claude keymap functionality',
+  })
+
   vim.api.nvim_create_user_command('ClaudeTestDiff', function()
-    local utils = require('nvim-claude.utils')
+    local utils = require 'nvim-claude.utils'
-    
+
     -- Check if we're in a git repository
     local git_root = utils.get_project_root()
     if not git_root then
       vim.notify('Not in a git repository', vim.log.levels.WARN)
       return
     end
-    
+
     -- Check if there are any changes
     local status_cmd = string.format('cd "%s" && git status --porcelain', git_root)
     local status_result = utils.exec(status_cmd)
-    
+
     if not status_result or status_result == '' then
       vim.notify('No changes to test', vim.log.levels.INFO)
       return
     end
-    
+
     -- Create test stash without restoring (to avoid conflicts)
-    local timestamp = os.date('%Y-%m-%d %H:%M:%S')
+    local timestamp = os.date '%Y-%m-%d %H:%M:%S'
     local stash_msg = string.format('[claude-test] %s', timestamp)
-    
+
     local stash_cmd = string.format('cd "%s" && git stash push -u -m "%s"', git_root, stash_msg)
     local stash_result, stash_err = utils.exec(stash_cmd)
-    
+
     if stash_err then
       vim.notify('Failed to create test stash: ' .. stash_err, vim.log.levels.ERROR)
       return
     end
-    
+
     -- Trigger diff review with the stash (no pre-edit ref for manual test)
-    local diff_review = require('nvim-claude.diff-review')
+    local diff_review = require 'nvim-claude.diff-review'
     diff_review.handle_claude_edit('stash@{0}', nil)
-    
+
     -- Pop the stash to restore changes
     vim.defer_fn(function()
       local pop_cmd = string.format('cd "%s" && git stash pop --quiet', git_root)
       utils.exec(pop_cmd)
-      vim.cmd('checktime')  -- Refresh buffers
+      vim.cmd 'checktime' -- Refresh buffers
     end, 100)
   end, {
-    desc = 'Test Claude diff review with current changes'
+    desc = 'Test Claude diff review with current changes',
   })
-  
+
   vim.api.nvim_create_user_command('ClaudeInstallHooks', function()
     M.install_hooks()
   end, {
-    desc = 'Install Claude Code hooks for this project'
+    desc = 'Install Claude Code hooks for this project',
   })
-  
+
   vim.api.nvim_create_user_command('ClaudeUninstallHooks', function()
     M.uninstall_hooks()
   end, {
-    desc = 'Uninstall Claude Code hooks for this project'
+    desc = 'Uninstall Claude Code hooks for this project',
   })
 end
 
 -- Cleanup old Claude commits and temp files
 function M.cleanup_old_commits()
-  local utils = require('nvim-claude.utils')
+  local utils = require 'nvim-claude.utils'
-  
+
   local git_root = utils.get_project_root()
   if not git_root then
     return
   end
-  
+
   -- Clean up old temp files
   local temp_files = {
     '/tmp/claude-pre-edit-commit',
     '/tmp/claude-baseline-commit',
     '/tmp/claude-last-snapshot',
-    '/tmp/claude-hook-test.log'
+    '/tmp/claude-hook-test.log',
   }
-  
+
   for _, file in ipairs(temp_files) do
     if vim.fn.filereadable(file) == 1 then
       vim.fn.delete(file)
     end
   end
-  
+
   -- Clean up old Claude commits (keep only the last 5)
-  local log_cmd = string.format('cd "%s" && git log --oneline --grep="claude-" --grep="claude-baseline" --grep="claude-pre-edit" --all --max-count=10', git_root)
+  local log_cmd =
+    string.format('cd "%s" && git log --oneline --grep="claude-" --grep="claude-baseline" --grep="claude-pre-edit" --all --max-count=10', git_root)
   local log_result = utils.exec(log_cmd)
-  
+
   if log_result and log_result ~= '' then
     local commits = {}
-    for line in log_result:gmatch('[^\n]+') do
+    for line in log_result:gmatch '[^\n]+' do
-      local hash = line:match('^(%w+)')
+      local hash = line:match '^(%w+)'
       if hash then
         table.insert(commits, hash)
       end
     end
-    
+
     -- Keep only the last 5 Claude commits, remove the rest
     if #commits > 5 then
       for i = 6, #commits do
@@ -426,4 +493,5 @@ function M.cleanup_old_commits()
   end
 end
 
-return M
+return M
+

lua/nvim-claude/inline-diff.lua

@@ -105,7 +105,7 @@ function M.parse_diff(diff_text)
   return hunks
 end
 
--- Apply visual indicators for diff
+-- Apply visual indicators for diff with line highlights
 function M.apply_diff_visualization(bufnr)
   local diff_data = M.active_diffs[bufnr]
   if not diff_data then return end
@@ -113,55 +113,88 @@ function M.apply_diff_visualization(bufnr)
   -- Clear existing highlights
   vim.api.nvim_buf_clear_namespace(bufnr, ns_id, 0, -1)
   
+  -- Get current buffer lines for reference
+  local buf_lines = vim.api.nvim_buf_get_lines(bufnr, 0, -1, false)
+  
   -- Apply highlights for each hunk
   for i, hunk in ipairs(diff_data.hunks) do
-    local line_num = hunk.old_start - 1 -- 0-indexed
+    -- Track additions and deletions separately with proper line mapping
-    
+    local current_new_line = hunk.new_start - 1 -- 0-indexed, tracks position in current buffer
-    -- Track lines to highlight
+    local current_old_line = hunk.old_start - 1 -- 0-indexed, tracks position in old content
-    local del_lines = {}
+    local deletions = {}
-    local add_lines = {}
-    local current_line = line_num
     
     for _, diff_line in ipairs(hunk.lines) do
-      if diff_line:match('^%-') then
+      if diff_line:match('^%+') then
-        -- Deletion
+        -- This is an added line - highlight it in the current buffer
-        table.insert(del_lines, current_line)
+        if current_new_line >= 0 and current_new_line < #buf_lines then
-        current_line = current_line + 1
+          vim.api.nvim_buf_set_extmark(bufnr, ns_id, current_new_line, 0, {
-      elseif diff_line:match('^%+') then
+            line_hl_group = 'DiffAdd',
-        -- Addition (shown as virtual text)
+            id = 4000 + i * 1000 + current_new_line
-        table.insert(add_lines, {
+          })
-          line = current_line - 1,
+        end
-          text = diff_line:sub(2)
+        current_new_line = current_new_line + 1
+        -- Don't advance old_line for additions
+      elseif diff_line:match('^%-') then
+        -- This is a deleted line - show as virtual text above current position
+        table.insert(deletions, {
+          line = current_new_line, -- Show deletion above current position
+          text = diff_line:sub(2),
         })
+        current_old_line = current_old_line + 1
+        -- Don't advance new_line for deletions
       else
-        -- Context line
+        -- Context line - advance both
-        current_line = current_line + 1
+        current_new_line = current_new_line + 1
+        current_old_line = current_old_line + 1
       end
     end
     
-    -- Apply deletion highlights
+    -- Show deletions as virtual text above their position with full-width background
-    for _, line in ipairs(del_lines) do
+    for j, del in ipairs(deletions) do
-      vim.api.nvim_buf_add_highlight(bufnr, ns_id, 'DiffDelete', line, 0, -1)
+      if del.line >= 0 and del.line <= #buf_lines then
+        -- Calculate full width for the deletion line
+        local text = '- ' .. del.text
+        local win_width = vim.api.nvim_win_get_width(0)
+        local padding = string.rep(' ', math.max(0, win_width - vim.fn.strdisplaywidth(text)))
+        
+        vim.api.nvim_buf_set_extmark(bufnr, ns_id, del.line, 0, {
+          virt_lines = {{
+            {'- ' .. del.text .. padding, 'DiffDelete'}
+          }},
+          virt_lines_above = true,
+          id = 3000 + i * 100 + j
+        })
+      end
     end
     
-    -- Add virtual text for additions
+    -- Add sign in gutter for hunk
-    for _, add in ipairs(add_lines) do
+    local sign_line = hunk.new_start - 1
-      vim.api.nvim_buf_set_extmark(bufnr, ns_id, add.line, 0, {
+    local sign_text = '>'
-        virt_lines = {{
+    local sign_hl = 'DiffAdd'
-          {' + ' .. add.text, 'DiffAdd'}
+    
-        }},
+    -- If hunk has deletions, use different sign
-        virt_lines_above = false
+    if #deletions > 0 then
+      sign_text = '~'
+      sign_hl = 'DiffChange'
+    end
+    
+    if sign_line >= 0 and sign_line < #buf_lines then
+      vim.api.nvim_buf_set_extmark(bufnr, ns_id, sign_line, 0, {
+        sign_text = sign_text,
+        sign_hl_group = sign_hl,
+        id = 2000 + i
       })
     end
     
-    -- Add hunk header as virtual text
+    -- Add subtle hunk info at end of first line
-    vim.api.nvim_buf_set_extmark(bufnr, ns_id, line_num, 0, {
+    local info_line = hunk.new_start - 1
-      virt_lines = {{
+    if info_line >= 0 and info_line < #buf_lines then
-        {hunk.header .. ' [Hunk ' .. i .. '/' .. #diff_data.hunks .. ']', 'Comment'}
+      vim.api.nvim_buf_set_extmark(bufnr, ns_id, info_line, 0, {
-      }},
+        virt_text = {{' [Hunk ' .. i .. '/' .. #diff_data.hunks .. ']', 'Comment'}},
-      virt_lines_above = true,
+        virt_text_pos = 'eol',
-      id = 1000 + i -- Unique ID for hunk headers
+        id = 1000 + i
-    })
+      })
+    end
   end
 end
 
@@ -261,12 +294,65 @@ function M.reject_current_hunk(bufnr)
   local diff_data = M.active_diffs[bufnr]
   if not diff_data then return end
   
+  local hunk = diff_data.hunks[diff_data.current_hunk]
+  if not hunk then return end
+  
+  -- Revert the hunk by applying original content
+  M.revert_hunk_changes(bufnr, hunk)
+  
   vim.notify(string.format('Rejected hunk %d/%d', diff_data.current_hunk, #diff_data.hunks), vim.log.levels.INFO)
   
   -- Move to next hunk
   M.next_hunk(bufnr)
 end
 
+-- Revert hunk changes (restore original content)
+function M.revert_hunk_changes(bufnr, hunk)
+  -- Get current buffer lines
+  local lines = vim.api.nvim_buf_get_lines(bufnr, 0, -1, false)
+  local original_content = M.original_content[bufnr]
+  
+  if not original_content then
+    vim.notify('No original content available for rejection', vim.log.levels.ERROR)
+    return
+  end
+  
+  -- Split original content into lines
+  local original_lines = vim.split(original_content, '\n')
+  
+  -- Build new lines by reverting this hunk
+  local new_lines = {}
+  local buffer_line = 1
+  local applied = false
+  
+  while buffer_line <= #lines do
+    if buffer_line >= hunk.new_start and buffer_line < hunk.new_start + hunk.new_count and not applied then
+      -- Revert this section by using original lines
+      local orig_start = hunk.old_start
+      local orig_end = hunk.old_start + hunk.old_count - 1
+      
+      for orig_line = orig_start, orig_end do
+        if orig_line <= #original_lines then
+          table.insert(new_lines, original_lines[orig_line])
+        end
+      end
+      
+      -- Skip the modified lines in current buffer
+      buffer_line = hunk.new_start + hunk.new_count
+      applied = true
+    else
+      -- Copy unchanged line
+      if buffer_line <= #lines then
+        table.insert(new_lines, lines[buffer_line])
+      end
+      buffer_line = buffer_line + 1
+    end
+  end
+  
+  -- Update buffer
+  vim.api.nvim_buf_set_lines(bufnr, 0, -1, false, new_lines)
+end
+
 -- Apply hunk changes to buffer
 function M.apply_hunk_changes(bufnr, hunk)
   -- Get current buffer lines
@@ -357,4 +443,18 @@ function M.has_active_diff(bufnr)
   return M.active_diffs[bufnr] ~= nil
 end
 
+-- Test keymap functionality
+function M.test_keymap()
+  local bufnr = vim.api.nvim_get_current_buf()
+  vim.notify('Testing keymap for buffer: ' .. bufnr, vim.log.levels.INFO)
+  vim.notify('Available diff data: ' .. vim.inspect(vim.tbl_keys(M.active_diffs)), vim.log.levels.INFO)
+  
+  if M.active_diffs[bufnr] then
+    vim.notify('Diff data found! Calling reject function...', vim.log.levels.INFO)
+    M.reject_current_hunk(bufnr)
+  else
+    vim.notify('No diff data for this buffer', vim.log.levels.ERROR)
+  end
+end
+
 return M