tia-lsp.nvim/lua/tia/variables.lua

-- SCL Variable Extractor
-- Extracts local variables from VAR sections and provides them with # prefix

local M = {}

local ts = vim.treesitter

-- Query to find variable declarations with their types
-- Simplified query that captures all var_item nodes regardless of parent type
local var_query_string = [[
  (var_item
    name: (identifier) @var_name
    type: (_) @var_type)
]]

-- Check if cursor is inside a VAR section
local function is_in_var_section()
  local cursor_node = ts.get_node()
  if not cursor_node then
    return false
  end
  
  local node = cursor_node
  while node do
    local type = node:type()
    if type == "var_declaration" then
      return true
    end
    if type == "organization_block" or type == "function_block" then
      -- Check if we're in the var declarations part (before BEGIN)
      local parent = node
      local children = {}
      for child in parent:iter_children() do
        table.insert(children, child)
      end
      
      for _, child in ipairs(children) do
        if child:type() == "var_declaration" then
          -- Check if cursor is before this var_declaration
          local start_row = child:range()
          local cursor_row = vim.api.nvim_win_get_cursor(0)[1] - 1
          if cursor_row <= start_row then
            return true
          end
        end
        if child:type() == "BEGIN" then
          local start_row = child:range()
          local cursor_row = vim.api.nvim_win_get_cursor(0)[1] - 1
          if cursor_row < start_row then
            return true
          end
        end
      end
      return false
    end
    node = node:parent()
  end
  return false
end

-- Extract all local variables from the buffer with their types
function M.extract_variables(bufnr)
  bufnr = bufnr or vim.api.nvim_get_current_buf()

  local variables = {}

  -- Try treesitter first
  local ok, parser = pcall(ts.get_parser, bufnr, "scl")
  if ok and parser then
    local tree = parser:parse()[1]
    if tree then
      local root = tree:root()
      local query_ok, query = pcall(ts.query.parse, "scl", var_query_string)
      if query_ok then
        for _, match, _ in query:iter_matches(root, bufnr, 0, -1) do
          local var_name = nil
          local var_type = nil
          
          for id, nodes in pairs(match) do
            -- nodes can be a single node or a table of nodes
            local node_list = type(nodes) == "table" and nodes or { nodes }
            for _, node in ipairs(node_list) do
              if node and type(node.range) == "function" then
                local text = ts.get_node_text(node, bufnr)
                local capture_name = query.captures[id]
                
                if capture_name == "var_name" then
                  var_name = text
                elseif capture_name == "var_type" then
                  var_type = text
                end
              end
            end
          end
          
          if var_name then
            table.insert(variables, {
              name = var_name,
              type = var_type or "",
            })
          end
        end
        return variables
      end
    end
  end

  -- Fallback: simple regex-based extraction
  local lines = vim.api.nvim_buf_get_lines(bufnr, 0, -1, false)
  local in_var = false
  for _, line in ipairs(lines) do
    if line:match("^%s*VAR") or line:match("^%s*VAR_INPUT") or line:match("^%s*VAR_OUTPUT") or line:match("^%s*VAR_IN_OUT") then
      in_var = true
    elseif line:match("^%s*END_VAR") then
      in_var = false
    elseif in_var then
      -- Match variable declarations like: "name : Type;" or "name : \"UdtType\";"
      -- Handle both basic types and quoted UDT types
      local var_name, var_type = line:match("^%s*([%w_]+)%s*:%s*([^;]+)")
      if var_name then
        -- Clean up type (remove attributes, default values, comments)
        if var_type then
          -- Remove attributes { ... }
          var_type = var_type:gsub("%s*{.-}%s*", " ")
          -- Remove default value := ...
          var_type = var_type:gsub("%s*:=.*$", "")
          -- Remove trailing comments
          var_type = var_type:gsub("%s*//.*$", "")
          -- Trim whitespace
          var_type = var_type:gsub("^%s*", ""):gsub("%s*$", "")
        end
        table.insert(variables, {
          name = var_name,
          type = var_type,
        })
      end
    end
  end

  return variables
end

-- Get the type of a specific variable
function M.get_variable_type(var_name, bufnr)
  local variables = M.extract_variables(bufnr)
  for _, var in ipairs(variables) do
    if var.name == var_name then
      return var.type
    end
  end
  return nil
end

-- Get all variables and their types as a map
function M.get_variable_types(bufnr)
  local variables = M.extract_variables(bufnr)
  local types = {}
  for _, var in ipairs(variables) do
    types[var.name] = var.type
  end
  return types
end

-- Get completion items based on context
function M.get_completions(bufnr)
  local variables = M.extract_variables(bufnr)
  local in_var_section = is_in_var_section()
  
  local items = {}
  for _, var in ipairs(variables) do
    if in_var_section then
      -- Inside VAR section: offer plain name
      table.insert(items, {
        label = var.name,
        kind = vim.lsp.protocol.CompletionItemKind.Variable,
        detail = var.type and (var.type .. " - Local variable") or "Local variable",
      })
    else
      -- Outside VAR section: offer with # prefix
      table.insert(items, {
        label = "#" .. var.name,
        insertText = "#" .. var.name,
        kind = vim.lsp.protocol.CompletionItemKind.Variable,
        detail = var.type and (var.type .. " - Local variable (prefixed)") or "Local variable (prefixed)",
      })
    end
  end
  
  return items
end

-- Check if we should provide completions
function M.should_complete()
  local ft = vim.bo.filetype
  if ft ~= "scl" then
    return false
  end
  
  -- Check if treesitter parser is available
  local ok, parser = pcall(ts.get_parser, 0, "scl")
  return ok and parser ~= nil
end

return M