diff --git a/lua/jsonfly/insert.lua b/lua/jsonfly/insert.lua
new file mode 100644
index 0000000..afc866d
--- /dev/null
+++ b/lua/jsonfly/insert.lua
@@ -0,0 +1,172 @@
+local utils = require"jsonfly.utils"
+
+
+local M = {};
+
+-- https://stackoverflow.com/a/24823383/9878135
+function table.slice(tbl, first, last, step)
+  local sliced = {}
+
+  for i = first or 1, last or #tbl, step or 1 do
+    sliced[#sliced+1] = tbl[i]
+  end
+
+  return sliced
+end
+
+---@param entry Entry
+---@param key string
+---@param index number
+local function check_key_equal(entry, key, index)
+    local splitted = utils:split_by_char(entry.key, ".")
+
+    return splitted[index] == key
+end
+
+---Find the entry in `entries` with the most matching keys at the beginning based on the `keys`.
+---Returns the index of the entry
+---@param entries Entry[]
+---@param keys KeyDescription[]
+---@return number|nil
+local function find_best_fitting_entry(entries, keys)
+    local entry_index
+    local current_indexes = {1, #entries}
+
+    for kk=1, #keys do
+        local key = keys[kk].key
+
+        local start_index = current_indexes[1]
+        local end_index = current_indexes[2]
+
+        current_indexes = {nil, nil}
+
+        for ii=start_index, end_index do
+            if check_key_equal(entries[ii], key, kk) then
+                if current_indexes[1] == nil then
+                    current_indexes[1] = ii
+                end
+
+                current_indexes[2] = ii
+            end
+        end
+
+        if current_indexes[1] == nil then
+            -- No entries found
+            break
+        else
+            entry_index = current_indexes[1]
+        end
+    end
+
+    return entry_index
+end
+
+---@param keys KeyDescription
+---@param index number - Index of the key
+local function write_keys(keys, index)
+    local lines = {}
+
+    if index == #keys then
+        return {
+            { "\"" .. keys[index].key .. "\": \"\""},
+            true
+        }
+    end
+
+    local insertions = write_keys(keys, index + 1)
+
+    lines[#lines + 1] = "\"" .. keys[index].key .. "\": {"
+
+    for ii=1, #insertions do
+        lines[#lines + 1] = insertions[ii]
+    end
+
+    lines[#lines + 1] = "}"
+
+    return lines
+end
+
+---@param buffer number
+---@param insertion_line number
+local function add_comma(buffer, insertion_line)
+    local BUFFER_SIZE = 5
+
+    -- Find next non-empty character in reverse
+    for ii=insertion_line, 0, -BUFFER_SIZE do
+        local previous_lines = vim.api.nvim_buf_get_lines(buffer, ii - BUFFER_SIZE, ii, false)
+
+        if #previous_lines == 0 then
+            return
+        end
+
+        for jj=#previous_lines, 0, -1 do
+            local line = previous_lines[jj]
+
+            for char_index=#line, 0, -1 do
+                local char = line:sub(char_index, char_index)
+
+                if char ~= " " and char ~= "\t" and char ~= "\n" and char ~= "\r" then
+                    if char == "," then
+                        return
+                    end
+
+
+                    -- Insert comma at position
+                    vim.api.nvim_buf_set_text(buffer, ii - 1, char_index, ii - 1, char_index, {","})
+                    return
+                end
+            end
+        end
+    end
+end
+
+---@param entries Entry[]
+---@param keys KeyDescription[]
+---@param buffer number
+function M:insert_new_key(entries, keys, buffer)
+    -- Close current buffer
+    vim.cmd [[quit!]]
+
+    local entry_index = find_best_fitting_entry(entries, keys) or 0
+    local entry = entries[entry_index]
+    local existing_keys_depth = #utils:split_by_char(entry.key, ".") + 1
+    local remaining_keys = table.slice(keys, existing_keys_depth, #keys)
+
+    local _writes = write_keys(remaining_keys, 1)
+    local writes = {}
+
+    for ii=1, #_writes do
+        if _writes[ii] == true then
+            -- Unwrap table
+            writes[#writes] = writes[#writes][1]
+        else
+            writes[#writes + 1] = _writes[ii]
+        end
+    end
+
+    vim.api.nvim_win_set_cursor(0, {entry.position.line_number, entry.position.value_start})
+    -- Hacky way to jump to end of object
+    vim.cmd [[execute "normal %"]]
+    local start_line = vim.api.nvim_win_get_cursor(0)[1] - 1
+
+    -- Add comma to previous line
+    add_comma(buffer, start_line)
+    --
+    -- Insert new lines
+    vim.api.nvim_buf_set_lines(buffer, start_line, start_line, false, writes)
+    --
+    -- -- -- Format lines
+    vim.api.nvim_win_set_cursor(0, {start_line, 1})
+    vim.cmd('execute "normal =' .. #writes .. 'j"')
+    --
+    -- -- Jump to the key
+    vim.api.nvim_win_set_cursor(0, {start_line + math.ceil(#writes / 2), 0})
+    vim.cmd [[execute "normal $a"]] 
+
+    -- vim.schedule(function()
+    --     vim.cmd [[%]]
+    -- end)
+
+end
+
+return M;
diff --git a/lua/telescope/_extensions/jsonfly.lua b/lua/telescope/_extensions/jsonfly.lua
index f0c8c81..3c93d3d 100644
--- a/lua/telescope/_extensions/jsonfly.lua
+++ b/lua/telescope/_extensions/jsonfly.lua
@@ -23,6 +23,7 @@
 local parsers = require"jsonfly.parsers"
 local utils = require"jsonfly.utils"
 local cache = require"jsonfly.cache"
+local insert = require"jsonfly.insert"
 
 local json = require"jsonfly.json"
 local finders = require "telescope.finders"
@@ -54,141 +55,6 @@ local opts = {
     use_cache = 500,
 }
 
--- https://stackoverflow.com/a/24823383/9878135
-function table.slice(tbl, first, last, step)
-  local sliced = {}
-
-  for i = first or 1, last or #tbl, step or 1 do
-    sliced[#sliced+1] = tbl[i]
-  end
-
-  return sliced
-end
-
----@param entry Entry
----@param key string
----@param index number
-local function check_key_equal(entry, key, index)
-    local splitted = utils:split_by_char(entry.key, ".")
-
-    return splitted[index] == key
-end
-
----@param entries Entry[]
----@param keys KeyDescription[]
----@return [Entry, string[]]
-local function find_remaining_keys(entries, keys)
-    local start_index = 1
-
-    local existing_keys_depth = nil
-
-    local potential_keys_indexes = {}
-
-    for kk=1, #keys do
-        local found_result = false
-        local key = keys[kk].key
-
-        for ii=start_index, #entries do
-            if check_key_equal(entries[ii], key, kk) then
-                found_result = true
-                start_index = ii
-                potential_keys_indexes[#potential_keys_indexes + 1] = ii
-                -- Not sure if this can be used for optimization
-                -- break
-            end
-        end
-
-        if not found_result then
-            existing_keys_depth = kk
-            break
-        end
-    end
-
-    if existing_keys_depth == nil then
-        existing_keys_depth = #keys
-    end
-
-    local existing_keys = table.slice(keys, 1, existing_keys_depth - 1)
-    local remaining_keys = table.slice(keys, existing_keys_depth, #keys)
-    local existing_key_str = ""
-
-    for ii=1, #existing_keys do
-        existing_key_str = existing_key_str .. existing_keys[ii].key
-        if ii < #existing_keys then
-            existing_key_str = existing_key_str .. "."
-        end
-    end
-
-    -- Find key that matches perfectly
-    for ii=1, #potential_keys_indexes do
-        local index = potential_keys_indexes[ii]
-        local entry = entries[index]
-
-        if entry.key == existing_key_str then
-            return {entry, remaining_keys}
-        end
-    end
-
-    -- Weird case, return the first one in hope that it's correct
-    return {potential_keys_indexes[1], remaining_keys}
-end
-
----@param keys KeyDescription
----@param index number - Index of the key
-local function write_keys(keys, index)
-    local lines = {}
-
-    if index == #keys then
-        return { "\"" .. keys[index].key .. "\": \"\""}
-    end
-
-    local insertions = write_keys(keys, index + 1)
-
-    lines[#lines + 1] = "\"" .. keys[index].key .. "\": {"
-
-    for ii=1, #insertions do
-        lines[#lines + 1] = insertions[ii]
-    end
-
-    lines[#lines + 1] = "}"
-
-    return lines
-end
-
----@param entries Entry[]
----@param keys KeyDescription[]
----@param buffer number
-local function insert_new_key(entries, keys, buffer)
-    -- Close current buffer
-    vim.cmd [[quit!]]
-
-    local _result = find_remaining_keys(entries, keys)
-    local entry = _result[1]
-    local remaining_keys = _result[2]
-
-    local writes = write_keys(remaining_keys, 1)
-    vim.api.nvim_buf_set_lines(buffer, start_line, start_line, false, writes)
-
-    print(vim.inspect(entry))
-    print(vim.inspect(remaining_keys))
-    print(vim.inspect(writes))
-
-    -- for parts=#keys, 1, -1 do
-    --     ---@type Entries[]
-    --     local sub_keys = table.slice(keys, 1, parts)
-    --     local path = ""
-    --     
-    --     for ii=1, #sub_keys do
-    --         path = path .. sub_keys[ii].key
-    --         if ii < #sub_keys then
-    --             path = path .. "."
-    --         end
-    --     end
-    --
-    --     print(path)
-    -- end
-end
-
 ---@param entries Entry[]
 ---@param buffer number
 local function show_picker(entries, buffer)
@@ -220,7 +86,7 @@ local function show_picker(entries, buffer)
 
                   local key_descriptor = utils:extract_key_description(input)
 
-                  insert_new_key(entries, key_descriptor, buffer)
+                  insert:insert_new_key(entries, key_descriptor, buffer)
             end)
 
             return true