diff --git a/lua/neogen.lua b/lua/neogen.lua
index 87f0614..8e52492 100644
--- a/lua/neogen.lua
+++ b/lua/neogen.lua
@@ -1,5 +1,4 @@
 local ts_utils = require("nvim-treesitter.ts_utils")
-local ts_query  = require("nvim-treesitter.query")
 
 local M = {}
 
@@ -10,7 +9,7 @@ M.generate = function ()
     (function (return_statement) @return)
     (function_definition (parameters) @params)
     ]]
-    
+
     -- Try to find the upper function
     local cursor = ts_utils.get_node_at_cursor(0)
     local function_node = cursor
@@ -21,6 +20,12 @@ M.generate = function ()
     end
     local line = ts_utils.get_node_range(function_node)
 
+    -- find the starting position in the line function
+    local line_content = vim.api.nvim_buf_get_lines(0, line, line+1, false)[1]
+    local offset = line_content:match("^%s+") or ""
+
+    local return_comment = offset .. "---@return "
+    local param_comment = offset .. "---@param "
 
     -- Parse and iterate over each found query
     local returned = vim.treesitter.parse_query("lua", query)
@@ -31,18 +36,18 @@ M.generate = function ()
             local params = ts_utils.get_node_text(node)[1]:sub(2,-2)
             for p in string.gmatch(params, '[^,]+') do
                 p = p:gsub("%s+", "") -- remove trailing spaces
-                table.insert(comment, "---@param " .. p .. " ")
+                table.insert(comment, param_comment .. p .. " ")
             end
         end
 
         -- Try to add return statement
         if returned.captures[id] == "return" then
-            table.insert(comment, "---@return ")
+            table.insert(comment, return_comment)
         end
     end
 
     -- At the end, add description annotation
-    table.insert(comment, 1, "---")
+    table.insert(comment, 1, offset .. "---")
 
     if #comment == 0 then return end
 
@@ -50,17 +55,16 @@ M.generate = function ()
     vim.fn.append(line, comment)
     vim.fn.cursor(line+1, #comment[1])
     vim.api.nvim_command('startinsert!')
-
 end
 
 function M.generate_command()
-  vim.api.nvim_command('command! -range -bar Neogen lua require("neogen").generate()')
+    vim.api.nvim_command('command! -range -bar Neogen lua require("neogen").generate()')
 end
 
 M.setup = function(opts)
     local config = opts or {}
     if config.enabled == true then
-       M.generate_command()
+        M.generate_command()
     end
 end