-- jj-mini.diff Neovim plugin

local M = {}

-- Default configuration
local config = {
  signs = {
    add = { text = "│", texthl = "JjDiffAdd", numhl = "JjDiffAdd" },
    change = { text = "│", texthl = "JjDiffChange", numhl = "JjDiffChange" },
    delete = { text = "─", texthl = "JjDiffDelete", numhl = "JjDiffDelete" },
  },
  autocmd_events = { "BufReadPost", "BufWritePost", "CursorHold" },
}

-- Helper function to run jj commands
local function _run_jj_command(args)
  local cmd = "jj " .. table.concat(args, " ")
  local handle = io.popen(cmd)
  if not handle then
    return nil, "Failed to run command: " .. cmd
  end
  local output = handle:read("*a")
  local status = handle:close()
  if not status then
    return nil, "Command failed: " .. cmd
  end
  return output
end

-- Define Neovim signs for diff
local function _define_signs()
  vim.fn.sign_define("JjDiffAdd", config.signs.add)
  vim.fn.sign_define("JjDiffChange", config.signs.change)
  vim.fn.sign_define("JjDiffDelete", config.signs.delete)
end

-- Get jj diff output for the current buffer
local function _get_jj_diff_for_buffer()
  local file_path = vim.api.nvim_buf_get_name(0)
  if file_path == "" then
    return nil, "Current buffer is not associated with a file."
  end

  -- Run 'jj diff' for the specific file
  local diff_output, err = _run_jj_command({ "diff", "--color=never", file_path })
  if err then
    return nil, err
  end
  return diff_output
end

-- Parse jj diff output
local function _parse_diff_output(diff_output)
  local added_lines = {}
  local changed_lines = {}
  local deleted_lines = {}

  local current_line_num = 0
  local lines = vim.split(diff_output, "\n", { plain = true })
  local prev_line_was_deleted = false

  for _, line in ipairs(lines) do
    if line:match("^@@ .- +(%d+)") then
      -- Extract new_start_line from hunk header
      local _, _, new_start_line_str = line:find("^@@ .- +(%d+)")
      current_line_num = tonumber(new_start_line_str) - 1
      prev_line_was_deleted = false
    elseif line:match("^[ ]") then
      current_line_num = current_line_num + 1
      prev_line_was_deleted = false
    elseif line:match("^[+]") then
      current_line_num = current_line_num + 1
      if prev_line_was_deleted then
        table.insert(changed_lines, current_line_num)
      else
        table.insert(added_lines, current_line_num)
      end
      prev_line_was_deleted = false
    elseif line:match("^[-]") then
      -- For deleted lines, we mark the line *before* the deletion as changed, 
      -- or if it's the first line, we can't mark it.
      -- This is a simplification for now.
      if current_line_num > 0 then
        table.insert(deleted_lines, current_line_num)
      end
      prev_line_was_deleted = true
    end
  end

  return added_lines, changed_lines, deleted_lines
end

-- Place signs in the current buffer
local function _place_signs_in_buffer()
  local buf_nr = vim.api.nvim_get_current_buf()
  vim.fn.sign_unplace("jj_mini_diff", { buffer = buf_nr }) -- Clear existing signs

  local diff_output = _get_jj_diff_for_buffer()
  if not diff_output then
    return
  end

  local added, changed, deleted = _parse_diff_output(diff_output)

  local sign_id = 1
  for _, line_num in ipairs(added) do
    vim.fn.sign_place(sign_id, "jj_mini_diff", "JjDiffAdd", buf_nr, { lnum = line_num })
    sign_id = sign_id + 1
  end
  for _, line_num in ipairs(changed) do
    vim.fn.sign_place(sign_id, "jj_mini_diff", "JjDiffChange", buf_nr, { lnum = line_num })
    sign_id = sign_id + 1
  end
  for _, line_num in ipairs(deleted) do
    vim.fn.sign_place(sign_id, "jj_mini_diff", "JjDiffDelete", buf_nr, { lnum = line_num })
    sign_id = sign_id + 1
  end
end

-- Check if current buffer's file is in a jj repo
local function _is_jj_repo()
  local file_path = vim.api.nvim_buf_get_name(0)
  if file_path == "" then
    return false
  end
  local dir = vim.fn.fnamemodify(file_path, ":h")
  while dir ~= "" and dir ~= "/" do
    if vim.fn.isdirectory(dir .. "/.jj") == 1 then
      return true
    end
    dir = vim.fn.fnamemodify(dir, ":h")
  end
  return false
end

function M.refresh_signs()
  if _is_jj_repo() then
    _place_signs_in_buffer()
  end
end

function M.setup(opts)
  opts = opts or {}
  _define_signs() -- Call to define signs

  -- Autocommands to update signs
  vim.api.nvim_create_autocmd({ "BufReadPost", "BufWritePost", "CursorHold" }, {
    group = vim.api.nvim_create_augroup("JjMiniDiff", { clear = true }),
    callback = function()
      if _is_jj_repo() then
        _place_signs_in_buffer()
      end
    end,
  })
end

return M