tia-lsp.nvim/lua/scl/auto_prefix.lua

-- Auto-prefixing module for SCL
-- Automatically adds # prefix to local variables when typed outside VAR sections
-- Replicates TIA Portal SCL behavior

local M = {}

local ts = vim.treesitter

-- Cache for local variables to avoid reparsing
local var_cache = {}
local cache_bufnr = nil
local cache_timestamp = nil

-- Clear cache when buffer changes
local function clear_cache()
  var_cache = {}
  cache_bufnr = nil
  cache_timestamp = nil
end

-- Keywords that should never be prefixed
local keywords = {
  "TRUE", "FALSE", "NOT", "AND", "OR", "XOR", "MOD",
  "IF", "THEN", "ELSE", "ELSIF", "END_IF",
  "FOR", "TO", "BY", "DO", "END_FOR",
  "WHILE", "END_WHILE",
  "REPEAT", "UNTIL", "END_REPEAT",
  "CASE", "OF", "END_CASE",
  "EXIT", "CONTINUE",
  "REGION", "END_REGION",
  "BEGIN", "VAR", "VAR_INPUT", "VAR_OUTPUT", "VAR_IN_OUT",
  "VAR_TEMP", "VAR_CONSTANT", "VAR_RETAIN", "VAR_NON_RETAIN", "END_VAR",
  "ORGANIZATION_BLOCK", "END_ORGANIZATION_BLOCK",
  "FUNCTION_BLOCK", "END_FUNCTION_BLOCK",
  "FUNCTION", "END_FUNCTION",
  "RETURN", "TYPE", "END_TYPE", "STRUCT", "END_STRUCT",
  "DATA_BLOCK", "END_DATA_BLOCK",
  "CONSTANT", "RETAIN", "NON_RETAIN",
}

-- Keywords that need semicolon after them (when at end of statement)
local end_keywords = {
  "END_IF", "END_FOR", "END_WHILE", "END_REPEAT", "END_CASE",
  "END_VAR", "END_FUNCTION", "END_FUNCTION_BLOCK", "END_ORGANIZATION_BLOCK",
  "END_TYPE", "END_STRUCT", "END_DATA_BLOCK", "END_REGION",
}

-- Map lowercase to uppercase for auto-conversion
local keyword_map = {}
for _, kw in ipairs(keywords) do
  keyword_map[kw:lower()] = kw
end

-- Check if word is a keyword
local function is_keyword(word)
  for _, kw in ipairs(keywords) do
    if word == kw then
      return true
    end
  end
  return false
end

-- Get list of local variable names (with caching)
local function get_local_variables()
  local bufnr = vim.api.nvim_get_current_buf()

  -- Return cached variables if available for current buffer
  if cache_bufnr == bufnr and cache_timestamp then
    local current_time = vim.loop.now()
    if current_time - cache_timestamp < 5000 then -- 5 second cache
      return var_cache
    end
  end

  local scl_vars = require("scl.variables")
  local vars = scl_vars.extract_variables(bufnr)

  var_cache = {}
  for _, var in ipairs(vars) do
    var_cache[var.name] = true
  end

  cache_bufnr = bufnr
  cache_timestamp = vim.loop.now()
  return var_cache
end

-- Check if we're inside a VAR declaration section
local function is_in_var_section()
  local bufnr = vim.api.nvim_get_current_buf()
  local cursor_row = vim.api.nvim_win_get_cursor(0)[1] - 1

  -- 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_declaration) @var
      ]])
      if query_ok then
        for _, match, _ in query:iter_matches(root, bufnr, 0, -1) do
          for _, nodes in pairs(match) do
            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 start_row, _, end_row, _ = node:range()
                if cursor_row >= start_row and cursor_row <= end_row then
                  return true
                end
              end
            end
          end
        end
        return false
      end
    end
  end

  -- Fallback: simple text-based detection
  local lines = vim.api.nvim_buf_get_lines(bufnr, 0, cursor_row + 1, false)
  local in_var = false
  for i = 1, cursor_row + 1 do
    local line = lines[i] or ""
    if line:match("^%s*VAR") then
      in_var = true
    elseif line:match("^%s*END_VAR") then
      in_var = false
    end
  end
  return in_var
end

-- Check if we're inside an FB/Function call parameter list
local function is_in_fb_call_params()
  local cursor_col = vim.api.nvim_win_get_cursor(0)[2]
  local line = vim.api.nvim_get_current_line()

  -- If line ends with ");" or similar, we're at end of params
  if line:match("%s*%);?%s*$") then
    return false
  end

  -- Count parens before cursor position
  local open_count = 0
  local close_count = 0
  for i = 1, cursor_col do
    local c = line:sub(i, i)
    if c == "(" then
      open_count = open_count + 1
    elseif c == ")" then
      close_count = close_count + 1
    end
  end

  -- If there are unclosed parens before cursor, we're inside params
  if open_count > close_count then
    return true
  end

  return false
end

-- Check if we're after BEGIN keyword
local function is_after_begin()
  local bufnr = vim.api.nvim_get_current_buf()
  local cursor_row = vim.api.nvim_win_get_cursor(0)[1] - 1

  -- Search backwards from current line to find BEGIN
  for i = cursor_row, 0, -1 do
    local line = vim.api.nvim_buf_get_lines(bufnr, i, i + 1, false)[1] or ""
    if line:match("^%s*BEGIN%s*$") or line:match("^%s*BEGIN%s*;") then
      return true
    end
    -- If we find another block start, stop searching
    if line:match("^%s*ORGANIZATION_BLOCK%s+") or
       line:match("^%s*FUNCTION_BLOCK%s+") or
       line:match("^%s*FUNCTION%s+") then
      return false
    end
  end

  return false
end

-- Extract base variable from a word (handles member access like "equipmentData.cmCylinder")
local function get_base_variable(word)
  -- Remove leading # if present
  local clean_word = word:gsub("^#", "")

  -- Find first identifier before any dot
  local base = clean_word:match("^([%w_]+)")
  return base
end

-- Find the last identifier in the line (before cursor) that needs prefixing
-- Only prefixes BASE variables (first in member chain), not member names
local function find_last_variable_to_prefix(line)
  local vars = get_local_variables()

  -- Find all identifiers in the line
  local candidates = {}

  -- Find all word patterns and their positions using string.find
  local pos = 1
  while pos <= #line do
    local start_pos, end_pos, match = line:find("([%w_]+)", pos)
    if not start_pos then
      break
    end
    if not is_keyword(match) and vars[match] then
      -- Skip if this identifier is part of a member chain (preceded by ".")
      -- Skip if already prefixed with #
      local should_skip = false
      if start_pos > 1 then
        local char_before = line:sub(start_pos - 1, start_pos - 1)
        if char_before == "." or char_before == "#" then
          should_skip = true
        end
      end
      if not should_skip then
        table.insert(candidates, {
          name = match,
          start = start_pos,
          finish = end_pos,
        })
      end
    end
    pos = end_pos + 1
  end

  if #candidates == 0 then
    return nil, nil, nil
  end

  -- Sort by position (last one first)
  table.sort(candidates, function(a, b)
    return a.finish > b.finish
  end)

  -- Check each candidate from end to start
  -- We want the LAST one that should be prefixed
  for _, cand in ipairs(candidates) do
    local start_pos = cand.start
    local end_pos = cand.finish
    local name = cand.name

    -- Double-check: never prefix if preceded by . or #
    if start_pos > 1 then
      local char_before = line:sub(start_pos - 1, start_pos - 1)
      if char_before ~= "." and char_before ~= "#" then
        -- This is a base variable - return it
        return name, start_pos, end_pos
      end
    else
      -- At start of line - this is a base variable
      return name, start_pos, end_pos
    end
  end

  return nil, nil, nil
end

-- Find base variable that ends at a specific column
-- Used when typing dot - finds variable before the dot
local function find_variable_before_dot(line, dot_col)
  local vars = get_local_variables()

  -- Go backwards from dot position to find word boundary
  local word_end = dot_col - 1
  while word_end > 0 and line:sub(word_end, word_end):match("[%w_]") do
    word_end = word_end - 1
  end
  word_end = word_end + 1

  local word = line:sub(word_end, dot_col - 1)
  if word == "" then
    return nil, nil, nil
  end

  -- Check if already prefixed
  if word:sub(1, 1) == "#" then
    return nil, nil, nil
  end

  -- Get base variable (before any member access)
  local base_var = get_base_variable(word)
  if not base_var then
    return nil, nil, nil
  end

  -- Check if it's a known local variable
  if not vars[base_var] then
    return nil, nil, nil
  end

  -- Check if already preceded by # (shouldn't happen due to check above, but be safe)
  if word_end > 1 then
    local prev_char = line:sub(word_end - 1, word_end - 1)
    if prev_char == "#" then
      return nil, nil, nil
    end
  end

  return base_var, word_end, word_end + #word - 1
end

-- Add # prefix to a variable at given position
local function add_prefix(bufnr, row, start_col, var_name)
  local line = vim.api.nvim_buf_get_lines(bufnr, row - 1, row, false)[1]
  local current_var = line:sub(start_col, start_col + #var_name - 1)

  -- Check if already prefixed
  if current_var:sub(1, 1) == "#" then
    return
  end

  -- Extra safety: check if the character before is . or #
  if start_col > 1 then
    local char_before = line:sub(start_col - 1, start_col - 1)
    if char_before == "." or char_before == "#" then
      return
    end
  end

  -- Insert # before the variable
  vim.api.nvim_buf_set_text(bufnr, row - 1, start_col - 1, row - 1, start_col - 1, {"#"})
end

-- Main auto-prefix function for space/semicolon triggers
local function auto_prefix_current_line()
  -- Check if auto-prefixing is enabled
  if not (vim.g.scl_auto_prefix or false) then
    return
  end

  -- Only in SCL files
  if vim.bo.filetype ~= "scl" then
    return
  end

  -- Don't prefix if inside VAR section
  if is_in_var_section() then
    return
  end

  -- Don't prefix if inside FB/Function call parameters
  if is_in_fb_call_params() then
    return
  end

  -- Only prefix after BEGIN
  if not is_after_begin() then
    return
  end

  local bufnr = vim.api.nvim_get_current_buf()
  local row = vim.api.nvim_win_get_cursor(0)[1]
  local line = vim.api.nvim_buf_get_lines(bufnr, row - 1, row, false)[1]

  -- Find the last variable that needs prefixing
  local var_name, start_col, end_col = find_last_variable_to_prefix(line)

  if var_name then
    add_prefix(bufnr, row, start_col, var_name)
  end
end

-- Auto-prefix triggered on dot character
local function on_dot_typed()
  -- Check if auto-prefixing is enabled
  if not (vim.g.scl_auto_prefix or false) then
    return
  end

  -- Only in SCL files
  if vim.bo.filetype ~= "scl" then
    return
  end

  -- Don't prefix if inside VAR section
  if is_in_var_section() then
    return
  end

  -- Don't prefix if inside FB/Function call parameters
  if is_in_fb_call_params() then
    return
  end

  -- Only prefix after BEGIN
  if not is_after_begin() then
    return
  end

  local bufnr = vim.api.nvim_get_current_buf()
  local row = vim.api.nvim_win_get_cursor(0)[1]
  local col = vim.api.nvim_win_get_cursor(0)[2]
  local line = vim.api.nvim_buf_get_lines(bufnr, row - 1, row, false)[1]

  -- Find variable before the dot (cursor is on the dot)
  local var_name, start_col, _ = find_variable_before_dot(line, col)

  if var_name then
    add_prefix(bufnr, row, start_col, var_name)
  end
end

-- Auto-uppercase keywords when typing space
local function auto_uppercase_keyword()
  -- Only in SCL files
  local ft = vim.bo.filetype
  if ft ~= "scl" and ft ~= "udt" and ft ~= "db" then
    return
  end

  local bufnr = vim.api.nvim_get_current_buf()
  local row = vim.api.nvim_win_get_cursor(0)[1]
  local col = vim.api.nvim_win_get_cursor(0)[2]
  local line = vim.api.nvim_buf_get_lines(bufnr, row - 1, row, false)[1]

  -- Check if we just typed a space
  if col < 1 or line:sub(col, col) ~= " " then
    return
  end

  -- Find word before the space
  local word_end = col - 1
  local word_start = word_end
  while word_start > 0 and line:sub(word_start, word_start):match("[%w_]") do
    word_start = word_start - 1
  end
  word_start = word_start + 1

  if word_start > word_end then
    return
  end

  local word = line:sub(word_start, word_end)
  local lower_word = word:lower()

  -- Check if it's a keyword that needs uppercase
  if keyword_map[lower_word] and word ~= keyword_map[lower_word] then
    local upper_word = keyword_map[lower_word]
    vim.api.nvim_buf_set_text(bufnr, row - 1, word_start - 1, row - 1, word_end, { upper_word })
  end
end

-- Auto-add semicolon after END_* keywords
local function auto_add_semicolon()
  -- Only in SCL files
  local ft = vim.bo.filetype
  if ft ~= "scl" and ft ~= "udt" and ft ~= "db" then
    return
  end

  local bufnr = vim.api.nvim_get_current_buf()
  local row = vim.api.nvim_win_get_cursor(0)[1]
  local line = vim.api.nvim_buf_get_lines(bufnr, row - 1, row, false)[1]
  local trimmed = line:match("^%s*(.-)%s*$")

  -- Check if line ends with an END_* keyword (without semicolon)
  for _, end_kw in ipairs(end_keywords) do
    if trimmed:upper() == end_kw then
      -- Check if already has semicolon
      if not line:match(";%s*$") then
        local new_line = line .. ";"
        vim.api.nvim_buf_set_lines(bufnr, row - 1, row, false, { new_line })
        -- Move cursor to end of line
        vim.api.nvim_win_set_cursor(0, { row, #new_line })
      end
      break
    end
  end
end

-- Setup auto-prefixing on specific triggers
function M.setup()
  local augroup = vim.api.nvim_create_augroup("SCLAutoPrefix", { clear = true })

  -- Enable by default
  vim.g.scl_auto_prefix = true

  -- Clear cache on text changes
  vim.api.nvim_create_autocmd({ "TextChanged", "TextChangedI", "BufEnter" }, {
    group = augroup,
    callback = function()
      if vim.bo.filetype == "scl" then
        clear_cache()
      end
    end,
  })

  -- Trigger on InsertLeave
  vim.api.nvim_create_autocmd("InsertLeave", {
    group = augroup,
    callback = function()
      local ft = vim.bo.filetype
      if ft == "scl" or ft == "udt" or ft == "db" then
        -- Auto-add semicolon after END_* keywords
        auto_add_semicolon()
        -- Variable prefixing (only for scl)
        if ft == "scl" then
          auto_prefix_current_line()
        end
      end
    end,
  })

  -- Trigger on space character in insert mode
  vim.api.nvim_create_autocmd("TextChangedI", {
    group = augroup,
    callback = function()
      local ft = vim.bo.filetype
      if ft == "scl" or ft == "udt" or ft == "db" then
        local line = vim.api.nvim_get_current_line()
        local cursor_col = vim.api.nvim_win_get_cursor(0)[2]

        if cursor_col > 0 and line:sub(cursor_col, cursor_col) == " " then
          -- Auto-uppercase keywords
          auto_uppercase_keyword()
          -- Variable prefixing (only for scl)
          if ft == "scl" then
            auto_prefix_current_line()
          end
        end
      end
    end,
  })

  -- Trigger on semicolon character in insert mode
  vim.api.nvim_create_autocmd("TextChangedI", {
    group = augroup,
    callback = function()
      if vim.bo.filetype == "scl" then
        vim.schedule(function()
          local line = vim.api.nvim_get_current_line()
          local cursor_col = vim.api.nvim_win_get_cursor(0)[2]
          local line_len = #line

          if cursor_col > 0 and cursor_col <= line_len and line:sub(cursor_col, cursor_col) == ";" then
            -- UDT type auto-quoting in VAR section
            local in_var = is_in_var_section()
            if in_var then
              local before_semicolon = line:sub(1, cursor_col - 1)
              local type_name = before_semicolon:match(":%s*([%w_]+)%s*$")
              if type_name then
                local udt_parser = require("scl.udt_parser")
                local is_udt = udt_parser.is_udt_type(type_name)
                if is_udt and not before_semicolon:match(":%s*\"" .. type_name .. "\"") then
                  local new_before = before_semicolon:gsub(":%s*" .. type_name .. "%s*$", ':"' .. type_name .. '"')
                  vim.api.nvim_set_current_line(new_before .. ";")
                end
              end
            end

            -- Variable prefixing
            auto_prefix_current_line()
          end
        end)
      end
    end,
  })

  -- Trigger on dot character in insert mode
  vim.api.nvim_create_autocmd("TextChangedI", {
    group = augroup,
    callback = function()
      if vim.bo.filetype == "scl" then
        local line = vim.api.nvim_get_current_line()
        local cursor_col = vim.api.nvim_win_get_cursor(0)[2]

        -- Check if the character at cursor is a dot
        if cursor_col > 0 and line:sub(cursor_col, cursor_col) == "." then
          on_dot_typed()
        end
      end
    end,
  })

  -- Auto-add semicolon when pressing Enter after END_* keyword
  vim.api.nvim_create_autocmd("TextChangedI", {
    group = augroup,
    callback = function()
      local ft = vim.bo.filetype
      if ft ~= "scl" and ft ~= "udt" and ft ~= "db" then
        return
      end

      local bufnr = vim.api.nvim_get_current_buf()
      local row = vim.api.nvim_win_get_cursor(0)[1]
      local line = vim.api.nvim_buf_get_lines(bufnr, row - 1, row, false)[1]

      -- Check if we're on a new/empty line (Enter was pressed)
      if line and line:match("^%s*$") and row > 1 then
        local prev_line = vim.api.nvim_buf_get_lines(bufnr, row - 2, row - 1, false)[1]
        if prev_line then
          local trimmed = prev_line:match("^%s*(.-)%s*$")
          -- Check if previous line is an END_* keyword without semicolon
          for _, end_kw in ipairs(end_keywords) do
            if trimmed:upper() == end_kw and not prev_line:match(";%s*$") then
              -- Add semicolon to previous line
              local new_prev = prev_line .. ";"
              vim.api.nvim_buf_set_lines(bufnr, row - 2, row - 1, false, { new_prev })
              break
            end
          end
        end
      end
    end,
  })
end

-- Manual trigger function
function M.prefix_word()
  local bufnr = vim.api.nvim_get_current_buf()
  local row = vim.api.nvim_win_get_cursor(0)[1]
  local col = vim.api.nvim_win_get_cursor(0)[2] + 1
  local line = vim.api.nvim_buf_get_lines(bufnr, row - 1, row, false)[1]

  -- Find word under cursor
  local start_col = col
  while start_col > 1 and line:sub(start_col - 1, start_col - 1):match("[%w_]") do
    start_col = start_col - 1
  end

  local end_col = col - 1
  while end_col <= #line and line:sub(end_col, end_col):match("[%w_]") do
    end_col = end_col + 1
  end
  end_col = end_col - 1

  local word = line:sub(start_col, end_col)
  if word == "" or word:sub(1, 1) == "#" then
    return
  end

  local base_var = get_base_variable(word)
  if not base_var then
    return
  end

  add_prefix(bufnr, row, start_col, word)
end

-- Get prefix enabled state
function M.is_enabled()
  return vim.g.scl_auto_prefix or false
end

-- Toggle prefix enabled state
function M.toggle()
  vim.g.scl_auto_prefix = not (vim.g.scl_auto_prefix or false)
  local status = vim.g.scl_auto_prefix and "enabled" or "disabled"
  vim.notify("SCL auto-prefix " .. status, vim.log.levels.INFO)
end

return M