lite-xl gemini code assist plugin

gemini.lua

-- mod-version:3
local core = require "core"
local common = require "core.common"
local config = require "core.config"
local DocView = require "core.docview"
local command = require "core.command"
local keymap = require "core.keymap"
local www = require "libraries.www"
local json = require "libraries.dkjson"

-- Configuration
local API_KEY = "get yours at https://aistudio.google.com/apikey"
local API_URL = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent?key=" .. API_KEY
local DELAY_AFTER_TYPING = 2000
local CONTEXT_LINES = 20
local MAX_CONTEXT_LENGTH = 1500

-- Global state
local ai_state = {
  timer = nil,
  suggestion = nil,
  suggestion_lines = {},
  doc = nil,
  line = nil,
  col = nil,
  is_showing = false,
  is_requesting = false
}

-- Safe logging function
local function safe_log(message, data)
  local safe_message = tostring(message or "nil")
  if data then
    local data_str = tostring(data)
    if data_str and #data_str > 0 then
      safe_message = safe_message .. ": " .. data_str:sub(1, 200)
    end
  end
  core.log("AI: " .. safe_message)
end

-- String validation function
local function is_valid_string(str, min_length)
  min_length = min_length or 1
  return str and type(str) == "string" and #str >= min_length
end

-- Number validation function
local function is_valid_number(num, min_val)
  min_val = min_val or 0
  return num and type(num) == "number" and num >= min_val
end

-- Safe string split function
local function split_string(str, delimiter)
  if not is_valid_string(str) or not is_valid_string(delimiter) then
    safe_log("split_string: Invalid input", "str=" .. tostring(str) .. ", delimiter=" .. tostring(delimiter))
    return {}
  end
  
  local result = {}
  local pattern = "(.-)" .. delimiter
  local last_end = 1
  
  local max_iterations = 1000
  local iterations = 0
  
  local s, e, cap = str:find(pattern, 1)
  
  while s and iterations < max_iterations do
    iterations = iterations + 1
    
    if s ~= 1 or cap ~= "" then
      table.insert(result, cap or "")
    end
    last_end = e + 1
    s, e, cap = str:find(pattern, last_end)
  end
  
  if last_end <= #str then
    cap = str:sub(last_end)
    table.insert(result, cap or "")
  end
  
  safe_log("split_string: Split into " .. #result .. " parts")
  return result
end

-- Parse JSON using dkjson with checks
local function parse_json(str)
  safe_log("parse_json: Starting", "length=" .. tostring(str and #str or 0))
  
  if not is_valid_string(str, 2) then
    safe_log("parse_json: Invalid input string")
    return nil
  end
  
  local trimmed = str:gsub("^%s*", ""):gsub("%s*$", "")
  if not (trimmed:sub(1, 1) == "{" or trimmed:sub(1, 1) == "[") then
    safe_log("parse_json: String doesn't look like JSON")
    return nil
  end
  
  safe_log("parse_json: JSON preview", trimmed:sub(1, 100))
  
  local ok, result, pos, err = pcall(json.decode, str, 1, nil)
  
  if not ok then
    safe_log("parse_json: pcall failed", tostring(result))
    return nil
  end
  
  if result and type(result) == "table" then 
    safe_log("parse_json: JSON parsed successfully")
    
    if not result.candidates then
      safe_log("parse_json: No candidates field")
      return nil
    end
    
    if type(result.candidates) ~= "table" then
      safe_log("parse_json: candidates is not a table")
      return nil
    end
    
    if #result.candidates == 0 then
      safe_log("parse_json: Empty candidates array")
      return nil
    end
    
    safe_log("parse_json: Found " .. #result.candidates .. " candidates")
    return result 
  else
    safe_log("parse_json: Parse failed", "pos=" .. tostring(pos) .. ", err=" .. tostring(err))
    return nil
  end
end

-- Clean markdown code blocks from response
local function clean_code_response(text)
  if not is_valid_string(text) then
    return text
  end
  
  -- Remove markdown code blocks (```language ... ```)
  local cleaned = text:gsub("```[%w]*\n?(.-)```", "%1")
  
  -- Remove single backticks
  cleaned = cleaned:gsub("`([^`]*)`", "%1")
  
  -- Remove any leading/trailing whitespace
  cleaned = cleaned:gsub("^%s*", ""):gsub("%s*$", "")
  
  return cleaned
end

-- Extract suggestion from Gemini response with validations
local function extract_suggestion(response_body)
  safe_log("extract_suggestion: Starting")
  
  if not is_valid_string(response_body, 10) then
    safe_log("extract_suggestion: Invalid response body")
    return nil
  end
  
  local parsed_json = parse_json(response_body)
  if not parsed_json then
    safe_log("extract_suggestion: Failed to parse JSON")
    return nil
  end
  
  if not parsed_json.candidates or type(parsed_json.candidates) ~= "table" or #parsed_json.candidates == 0 then
    safe_log("extract_suggestion: Invalid candidates")
    return nil
  end
  
  local candidate = parsed_json.candidates[1]
  if not candidate or type(candidate) ~= "table" then
    safe_log("extract_suggestion: Invalid first candidate")
    return nil
  end
  
  if not candidate.content or type(candidate.content) ~= "table" then
    safe_log("extract_suggestion: Invalid candidate content")
    return nil
  end
  
  if not candidate.content.parts or type(candidate.content.parts) ~= "table" or #candidate.content.parts == 0 then
    safe_log("extract_suggestion: Invalid content parts")
    return nil
  end
  
  local part = candidate.content.parts[1]
  if not part or type(part) ~= "table" then
    safe_log("extract_suggestion: Invalid first part")
    return nil
  end
  
  local text = part.text
  if not is_valid_string(text, 1) then
    safe_log("extract_suggestion: No valid text found")
    return nil
  end
  
  -- Clean markdown formatting
  text = clean_code_response(text)
  
  safe_log("extract_suggestion: Success", "length=" .. #text)
  return text
end

-- Get current code context
local function get_code_context(doc, line, col)
  safe_log("get_code_context: Starting", "line=" .. tostring(line) .. ", col=" .. tostring(col))
  
  if not doc or not doc.lines then
    safe_log("get_code_context: Invalid document")
    return nil
  end
  
  if not is_valid_number(line, 1) or not is_valid_number(col, 1) then
    safe_log("get_code_context: Invalid line/col")
    return nil
  end
  
  if line > #doc.lines then
    safe_log("get_code_context: Line beyond document")
    return nil
  end
  
  -- Build context with surrounding lines
  local context = ""
  local start_line = math.max(1, line - CONTEXT_LINES)
  local end_line = math.min(#doc.lines, line + 5)
  
  for i = start_line, end_line do
    local line_text = doc.lines[i] or ""
    
    if i == line then
      -- Mark the cursor position
      if col > #line_text + 1 then
        context = context .. line_text .. "<CURSOR>"
      else
        local before_cursor = line_text:sub(1, col - 1)
        local after_cursor = line_text:sub(col)
        context = context .. before_cursor .. "<CURSOR>" .. after_cursor
      end
    else
      context = context .. line_text
    end
    
    if i < end_line then
      context = context .. "\n"
    end
  end
  
  -- Limit context size
  if #context > MAX_CONTEXT_LENGTH then
    context = context:sub(-MAX_CONTEXT_LENGTH)
  end
  
  if not is_valid_string(context, 5) then
    safe_log("get_code_context: Context too small")
    return nil
  end
  
  safe_log("get_code_context: Success", "length=" .. #context)
  return context
end

-- Create prompt for AI with filename
local function create_prompt(context, filename)
  local filename_info = filename and filename ~= "" and filename or "untitled"
  
  local prompt = string.format([[You are an AI code completion assistant. Complete the code at the <CURSOR> position.

Filename: %s

IMPORTANT INSTRUCTIONS:
- Return ONLY the code completion, absolutely NO markdown formatting
- Do NOT wrap your response in ```code blocks``` or backticks
- Do NOT include any explanations, comments, or descriptions
- Provide only the raw code that should be inserted at the cursor position
- Maintain proper indentation and formatting that matches the existing code
- Consider the existing code context and patterns
- Make the completion logical and syntactically correct
- Detect the programming language from the filename and context
- If it's a function, complete the entire function body
- If it's a statement, complete it appropriately
- Keep completions concise but functional
- Remove the <CURSOR> marker from your response

Code:
%s]], filename_info, context)

  return prompt
end

-- Get current cursor position (updated in real-time)
local function get_current_cursor_position(doc)
  if not doc then
    return nil, nil
  end
  
  local line, col = doc:get_selection()
  return line, col
end

-- Make request to Gemini
local function request_ai_completion(doc, line, col)
  safe_log("request_ai_completion: Starting")
  
  if ai_state.is_requesting then
    safe_log("request_ai_completion: Already requesting, skipping")
    return
  end
  
  if not doc or not is_valid_number(line, 1) or not is_valid_number(col, 1) then
    safe_log("request_ai_completion: Invalid parameters")
    return
  end
  
  -- Get the most current cursor position when making the request
  local current_line, current_col = get_current_cursor_position(doc)
  if not current_line or not current_col then
    safe_log("request_ai_completion: Failed to get current cursor position")
    return
  end
  
  -- Use the current position instead of the potentially outdated ones
  line, col = current_line, current_col
  
  local context = get_code_context(doc, line, col)
  if not context then
    safe_log("request_ai_completion: Failed to get context")
    return
  end
  
  if not is_valid_string(API_KEY, 10) then
    safe_log("request_ai_completion: Invalid API key")
    return
  end
  
  ai_state.is_requesting = true
  safe_log("request_ai_completion: Creating payload")
  
  local filename = doc.filename or ""
  local prompt = create_prompt(context, filename)
  
  local ok, escaped_prompt = pcall(json.encode, prompt)
  if not ok then
    safe_log("request_ai_completion: Failed to encode prompt")
    ai_state.is_requesting = false
    return
  end
  
  local payload = string.format([[{
    "contents": [
      {
        "parts": [
          {
            "text": %s
          }
        ]
      }
    ],
    "generationConfig": {
      "temperature": 0.3,
      "maxOutputTokens": 1000,
      "topP": 0.8,
      "topK": 40
    }
  }]], escaped_prompt)
  
  if not is_valid_string(payload, 50) then
    safe_log("request_ai_completion: Invalid payload")
    ai_state.is_requesting = false
    return
  end
  
  safe_log("request_ai_completion: Making HTTP request")
  
  core.add_thread(function()
    local request_ok, res = pcall(function()
      return www.request({
        url = API_URL,
        method = "POST",
        headers = {
          ["Content-Type"] = "application/json",
          ["Connection"] = "close"
        },
        body = payload
      })
    end)
    
    ai_state.is_requesting = false
    
    if not request_ok then
      safe_log("request_ai_completion: HTTP request failed", tostring(res))
      return
    end
    
    if not res or type(res) ~= "table" then
      safe_log("request_ai_completion: Invalid response object")
      return
    end
    
    safe_log("request_ai_completion: Response received", "code=" .. tostring(res.code))
    
    if not is_valid_number(res.code) then
      safe_log("request_ai_completion: Invalid response code")
      return
    end
    
    if not is_valid_string(res.body, 1) then
      safe_log("request_ai_completion: Invalid response body")
      return
    end
    
    safe_log("request_ai_completion: Body length", #res.body)
    
    if res.code == 200 then
      local suggestion = extract_suggestion(res.body)
      if suggestion and is_valid_string(suggestion, 1) then
        -- Process text and remove cursor marker
        local processed = suggestion:gsub("^%s*", ""):gsub("%s*$", "")
        processed = processed:gsub("\\n", "\n")
        processed = processed:gsub("<CURSOR>", "")
        
        if is_valid_string(processed, 1) then
          -- Get the current cursor position again to ensure accuracy
          local final_line, final_col = get_current_cursor_position(doc)
          if final_line and final_col then
            ai_state.suggestion = processed
            ai_state.suggestion_lines = split_string(processed, "\n")
            ai_state.doc = doc
            ai_state.line = final_line
            ai_state.col = final_col
            ai_state.is_showing = true
            
            safe_log("request_ai_completion: Success", "lines=" .. #ai_state.suggestion_lines)
          else
            safe_log("request_ai_completion: Failed to get final cursor position")
          end
        else
          safe_log("request_ai_completion: Processed text invalid")
        end
      else
        safe_log("request_ai_completion: No valid suggestion")
      end
    else
      safe_log("request_ai_completion: HTTP error", "code=" .. res.code)
    end
  end)
end

-- Clear suggestion state
local function clear_suggestion()
  ai_state.suggestion = nil
  ai_state.suggestion_lines = {}
  ai_state.is_showing = false
end

-- Reset timer - cancels existing and creates new one
local function reset_timer(doc, line, col)
  safe_log("reset_timer: Starting")
  
  if not doc or not is_valid_number(line, 1) or not is_valid_number(col, 1) then
    safe_log("reset_timer: Invalid parameters")
    return
  end
  
  -- Cancel any existing timer
  if ai_state.timer then
    safe_log("reset_timer: Cancelling previous timer")
    ai_state.timer = nil
  end
  
  -- Create new timer that waits for DELAY_AFTER_TYPING
  ai_state.timer = core.add_thread(function()
    local start_time = system.get_time()
    local delay_seconds = DELAY_AFTER_TYPING / 1000
    
    -- Wait for the delay period
    while system.get_time() - start_time < delay_seconds do
      coroutine.yield(0.1) -- Check every 100ms
      
      -- If timer was cancelled (new typing occurred), exit
      if not ai_state.timer then
        safe_log("reset_timer: Timer was cancelled during wait")
        return
      end
    end
    
    -- Timer completed, make request if still valid
    if ai_state.timer then
      safe_log("reset_timer: Timer expired, making request")
      -- Get the current cursor position at the time of request
      local current_line, current_col = get_current_cursor_position(doc)
      if current_line and current_col then
        request_ai_completion(doc, current_line, current_col)
      else
        safe_log("reset_timer: Failed to get current cursor position for request")
      end
      ai_state.timer = nil
    else
      safe_log("reset_timer: Timer was cancelled before completion")
    end
  end)
  
  safe_log("reset_timer: New timer created for 2 seconds")
end

-- Hook on text input
local original_on_text_input = DocView.on_text_input

function DocView:on_text_input(text)
  if not self or not self.doc then
    safe_log("on_text_input: Invalid self or doc")
    return
  end
  
  if not is_valid_string(text, 1) then
    safe_log("on_text_input: Invalid text input")
    return
  end
  
  local ok, result = pcall(original_on_text_input, self, text)
  if not ok then
    safe_log("on_text_input: Original function failed", tostring(result))
    return
  end
  
  clear_suggestion()
  
  -- Start/reset timer for any file (no extension filtering)
  local doc = self.doc
  if doc then
    -- Get current cursor position after the text input
    local line, col = get_current_cursor_position(doc)
    if line and col then
      reset_timer(doc, line, col)
    else
      safe_log("on_text_input: Invalid cursor position")
    end
  end
end

-- Hook to clear suggestion when cursor moves or keys are pressed
local original_on_mouse_pressed = DocView.on_mouse_pressed
function DocView:on_mouse_pressed(button, x, y, clicks)
  if not self then
    safe_log("on_mouse_pressed: Invalid self")
    return
  end
  
  local ok, result = pcall(original_on_mouse_pressed, self, button, x, y, clicks)
  if not ok then
    safe_log("on_mouse_pressed: Original function failed", tostring(result))
  end
  
  clear_suggestion()
  -- Cancel timer on mouse interaction
  if ai_state.timer then
    ai_state.timer = nil
  end
  
  return result
end

-- Hook key presses
local original_on_key_pressed = DocView.on_key_pressed
function DocView:on_key_pressed(key, x, y)
  if not self then
    safe_log("on_key_pressed: Invalid self")
    return
  end
  
  local ok, result = pcall(original_on_key_pressed, self, key, x, y)
  if not ok then
    safe_log("on_key_pressed: Original function failed", tostring(result))
    return
  end
  
  -- Clear suggestion and cancel timer on navigation/editing keys
  if key == "up" or key == "down" or key == "left" or key == "right" or 
     key == "home" or key == "end" or key == "pageup" or key == "pagedown" or
     key == "return" or key == "backspace" or key == "delete" then
    clear_suggestion()
    if ai_state.timer then
      ai_state.timer = nil
    end
  end
  
  return result
end

-- Enhanced ghost text rendering
local original_draw_line_body = DocView.draw_line_body

function DocView:draw_line_body(line, x, y)
  if not self or not is_valid_number(line, 1) or not is_valid_number(x) or not is_valid_number(y) then
    safe_log("draw_line_body: Invalid parameters")
    if self and original_draw_line_body then
      return original_draw_line_body(self, line or 1, x or 0, y or 0)
    end
    return
  end
  
  local ok, result = pcall(original_draw_line_body, self, line, x, y)
  if not ok then
    safe_log("draw_line_body: Original function failed", tostring(result))
    return
  end
  
  -- Draw ghost suggestion
  if ai_state.is_showing and ai_state.suggestion and #ai_state.suggestion_lines > 0 and 
     ai_state.doc == self.doc and ai_state.line == line then
    
    -- Get current cursor position to ensure accuracy
    local current_line, current_col = get_current_cursor_position(self.doc)
    if line == current_line and is_valid_number(current_col, 1) then
      local line_text = self.doc.lines[line] or ""
      local before_cursor = ""
      
      if ai_state.col and is_valid_number(ai_state.col, 1) and ai_state.col <= #line_text + 1 then
        before_cursor = line_text:sub(1, ai_state.col - 1)
      end
      
      local font = self:get_font()
      if not font then
        safe_log("draw_line_body: No font available")
        return result
      end
      
      local text_width = font:get_width(before_cursor or "")
      local line_height = self:get_line_height()
      
      if not is_valid_number(text_width) or not is_valid_number(line_height, 1) then
        safe_log("draw_line_body: Invalid font metrics")
        return result
      end
      
      -- Draw each line of the suggestion
      for i, sug_line in ipairs(ai_state.suggestion_lines) do
        if is_valid_string(sug_line) then
          local draw_y, draw_x
          
          if i == 1 then
            draw_y = y
            draw_x = x + text_width
          else
            draw_y = y + (i - 1) * line_height
            draw_x = x
          end
          
          if is_valid_number(draw_x) and is_valid_number(draw_y) then
            local sug_width = font:get_width(sug_line)
            if is_valid_number(sug_width, 0) then
              -- Draw subtle background
              renderer.draw_rect(draw_x, draw_y, sug_width, line_height, {50, 50, 50, 40})
              -- Draw border
              renderer.draw_rect(draw_x, draw_y, sug_width, 1, {100, 100, 100, 80})
              renderer.draw_rect(draw_x, draw_y + line_height - 1, sug_width, 1, {100, 100, 100, 80})
            end
            
            -- Draw suggestion text with ghost effect
            renderer.draw_text(font, sug_line, draw_x + 0.5, draw_y, {180, 180, 180, 160})
            renderer.draw_text(font, sug_line, draw_x, draw_y, {160, 160, 160, 180})
          end
        end
      end
    end
  end
  
  return result
end

-- Function to check if there's an active AI suggestion
local function has_active_suggestion(dv)
  return ai_state.is_showing and ai_state.suggestion and ai_state.doc and ai_state.doc == dv.doc
end

-- Commands for accepting/rejecting suggestions
command.add(DocView, {
  ["ai-assist:accept"] = function(dv)
    safe_log("accept: Starting")
    
    if not dv or not dv.doc then
      safe_log("accept: Invalid DocView or doc")
      return
    end
    
    if not ai_state.suggestion or not ai_state.doc or ai_state.doc ~= dv.doc then
      safe_log("accept: No valid suggestion")
      return
    end
    
    if not is_valid_string(ai_state.suggestion, 1) then
      safe_log("accept: Invalid suggestion text")
      clear_suggestion()
      return
    end
    
    local doc = dv.doc
    local current_line, current_col = get_current_cursor_position(doc)
    
    if not current_line or not current_col then
      safe_log("accept: Invalid cursor position")
      clear_suggestion()
      return
    end
    
    if current_line > #doc.lines then
      safe_log("accept: Line beyond document")
      clear_suggestion()
      return
    end
    
    local line_text = doc.lines[current_line] or ""
    if current_col > #line_text + 1 then
      safe_log("accept: Column beyond line")
      clear_suggestion()
      return
    end
    
    -- Insert the suggestion
    local insert_ok, insert_err = pcall(function()
      doc:insert(current_line, current_col, ai_state.suggestion)
    end)
    
    if insert_ok then
      safe_log("accept: Suggestion inserted successfully")
    else
      safe_log("accept: Insert failed", tostring(insert_err))
    end
    
    clear_suggestion()
  end,
  
  ["ai-assist:reject"] = function()
    safe_log("reject: Clearing suggestion")
    clear_suggestion()
    -- Also cancel any pending timer
    if ai_state.timer then
      ai_state.timer = nil
    end
  end,
  
  -- New command that handles tab intelligently
  ["ai-assist:smart-tab"] = function(dv)
    -- If there's an active AI suggestion, accept it
    if has_active_suggestion(dv) then
      command.perform("ai-assist:accept", dv)
    else
      -- Otherwise, perform normal tab functionality
      command.perform("doc:indent", dv)
    end
  end
})

-- Keybindings - usando o smart-tab que decide o comportamento
keymap.add {
  ["tab"] = "ai-assist:smart-tab",
  ["escape"] = "ai-assist:reject"
}

safe_log("AI Assistant plugin loaded successfully")