local M = {}

---Toggle checkbox at cursor position and all children
-- Uses treesitter to find markdown list items with checkboxes
function M.toggle_checkbox()
  local bufnr = vim.api.nvim_get_current_buf()
  local cursor = vim.api.nvim_win_get_cursor(0)
  local cursor_row = cursor[1] - 1 -- 0-indexed
  local cursor_col = cursor[2]

  -- Get markdown parser
  local ok, parser = pcall(vim.treesitter.get_parser, bufnr, 'markdown')
  if not ok or not parser then
    vim.notify('Markdown parser not available', vim.log.levels.WARN)
    return
  end

  -- Parse and get tree
  local trees = parser:parse()
  if not trees or #trees == 0 then
    return
  end

  local tree = trees[1]
  local root = tree:root()

  -- Find list_item node at cursor position
  local target_node = nil
  for node in root:iter_children() do
    if node:range() then
      local start_row, start_col, end_row, end_col = node:range()
      -- end_row is exclusive, so use < not <=
      if cursor_row >= start_row and cursor_row < end_row then
        -- Check if this is a list_item or contains one
        target_node = M.find_list_item_at(node, cursor_row, cursor_col)
        if target_node then
          break
        end
      end
    end
  end

  if not target_node then
    return
  end

  -- Toggle the target and its children
  local new_state = M.toggle_list_item(bufnr, target_node)
  if new_state then
    M.toggle_children(bufnr, target_node, new_state)
  end
end

---Find list_item node containing cursor position
-- Searches children first to return the deepest matching node
function M.find_list_item_at(node, cursor_row, cursor_col)
  -- Check children first to find the deepest match
  for child in node:iter_children() do
    local result = M.find_list_item_at(child, cursor_row, cursor_col)
    if result then
      return result
    end
  end

  -- Then check this node
  if node:type() == 'list_item' then
    local start_row, start_col, end_row, end_col = node:range()
    -- end_row is exclusive, so use < not <=
    if cursor_row >= start_row and cursor_row < end_row then
      return node
    end
  end

  return nil
end

---Toggle a single list item's checkbox
-- Returns the new state ('x' or ' ') or nil if not a task
function M.toggle_list_item(bufnr, node)
  -- Find task_list_marker
  for child in node:iter_children() do
    local child_type = child:type()
    if child_type == 'task_list_marker_unchecked' then
      -- Replace [ ] with [X]
      local start_row, start_col, end_row, end_col = child:range()
      vim.api.nvim_buf_set_text(bufnr, start_row, start_col, end_row, end_col, { '[X]' })
      return 'X'
    elseif child_type == 'task_list_marker_checked' then
      -- Replace [X] with [ ]
      local start_row, start_col, end_row, end_col = child:range()
      vim.api.nvim_buf_set_text(bufnr, start_row, start_col, end_row, end_col, { '[ ]' })
      return ' '
    end
  end

  return nil
end

---Toggle all child list items to match parent state
function M.toggle_children(bufnr, parent_node, target_state)
  -- Find nested list
  for child in parent_node:iter_children() do
    if child:type() == 'list' then
      -- Iterate through list_items in the nested list
      for list_item in child:iter_children() do
        if list_item:type() == 'list_item' then
          -- Force this child to target state
          M.force_list_item_state(bufnr, list_item, target_state)
          -- Recursively handle grandchildren
          M.toggle_children(bufnr, list_item, target_state)
        end
      end
      break
    end
  end
end

---Force a list item to specific state without toggling
function M.force_list_item_state(bufnr, node, target_state)
  for child in node:iter_children() do
    local child_type = child:type()
    if child_type == 'task_list_marker_unchecked' and target_state == 'X' then
      -- Need to check it
      local start_row, start_col, end_row, end_col = child:range()
      vim.api.nvim_buf_set_text(bufnr, start_row, start_col, end_row, end_col, { '[X]' })
      return
    elseif child_type == 'task_list_marker_checked' and target_state == ' ' then
      -- Need to uncheck it
      local start_row, start_col, end_row, end_col = child:range()
      vim.api.nvim_buf_set_text(bufnr, start_row, start_col, end_row, end_col, { '[ ]' })
      return
    end
  end
end

return M