tia-lsp.nvim/lua/tia_lsp/init.lua

local M = {}

local function setup_package_path()
  local plugin_path = debug.getinfo(1, "S").source:gsub("^@", ""):match("(.*/)") or ""
  plugin_path = plugin_path:gsub("/lua/tia_lsp", "")
  local lua_path = plugin_path .. "/lua"
  if not package.path:match(lua_path) then
    package.path = lua_path .. "/?.lua;" .. package.path
  end
end

function M.setup(opts)
  setup_package_path()
  opts = opts or {}

  local server_path = opts.server_path or vim.fn.expand("~/dev/tia-lsp/src/main.lua")
  
  local root_patterns = { ".git", "data_types", "plc.data.json" }
  
  local function start_lsp(bufnr)
    local fname = vim.api.nvim_buf_get_name(bufnr)
    if fname == "" then return end
    
    local root_dir = vim.fs.root(fname, root_patterns)
      or vim.fn.fnamemodify(fname, ":p:h")
    
    -- Prefer the mason-installed executable (vim.fn.exepath resolves on each
    -- buffer enter, so a :MasonInstall mid-session is picked up without restart).
    -- Fall back to running the server via `lua <server_path>` for manual installs.
    local tia_bin = vim.fn.exepath("tia-lsp")
    local cmd = tia_bin ~= "" and { "tia-lsp" } or { "lua", server_path }
    
    vim.lsp.start({
      name = "tia_lsp",
      cmd = cmd,
      root_dir = root_dir,
      on_attach = function(client, b)
        if opts.lsp and opts.lsp.on_attach then
          opts.lsp.on_attach(client, b)
        end
        
        if opts.lsp and opts.lsp.formatting then
          vim.api.nvim_buf_create_user_command(b, "SCLFormat", function()
            vim.lsp.buf.format({ name = "tia_lsp" })
          end, {})
        end
      end,
    })
  end
  
  -- Create FileType autocmd for future buffers
  vim.api.nvim_create_autocmd("FileType", {
    pattern = "scl",
    callback = function(args)
      start_lsp(args.buf)
      M.notify_parser_install_if_needed(args.buf)
      -- Start tree-sitter highlighting if the parser is registered
      pcall(vim.treesitter.start, args.buf, "scl")
    end,
  })
  
  -- Create global format command (works on current buffer)
  vim.api.nvim_create_user_command("SCLFormat", function()
    vim.lsp.buf.format({ name = "tia_lsp" })
  end, {})
  
  -- Start LSP for any existing scl buffers (handles lazy loading case)
  for _, bufnr in ipairs(vim.api.nvim_list_bufs()) do
    if vim.bo[bufnr].filetype == "scl" and vim.api.nvim_buf_is_loaded(bufnr) then
      start_lsp(bufnr)
      M.notify_parser_install_if_needed(bufnr)
      pcall(vim.treesitter.start, bufnr, "scl")
    end
  end

  M.setup_syntax(opts)
  M.setup_multiline_params()
  M.setup_attr_toggle()
  M.create_commands()

  if opts.auto_prefix ~= false then
    M.setup_auto_prefix()
  end
end

-- Find a compiled tree-sitter parser (.so) for SCL.
-- Looks in: opts.ts_parser_path, mason package dir, ~/dev/tia-lsp.
-- Returns the path to parser.so, or nil if not found.
function M.find_parser_so(opts)
  -- 1. Explicit path from opts
  if opts.ts_parser_path and vim.fn.filereadable(opts.ts_parser_path) == 1 then
    return opts.ts_parser_path
  end
  -- 2. Mason install: resolve exepath("tia-lsp") -> real bin -> package dir
  local tia_bin = vim.fn.exepath("tia-lsp")
  if tia_bin ~= "" then
    local real_bin = vim.fn.resolve(vim.fn.fnamemodify(tia_bin, ":p"))
    local pkg_dir = vim.fn.fnamemodify(real_bin, ":h:h")
    local so = pkg_dir .. "/parser.so"
    if vim.fn.filereadable(so) == 1 then
      return so
    end
  end
  -- 3. Local dev path
  local dev_so = vim.fn.expand("~/dev/tia-lsp/parser.so")
  if vim.fn.filereadable(dev_so) == 1 then
    return dev_so
  end
  return nil
end

-- Notify the user once per session if the tree-sitter parser isn't available.
local parser_notified = false

function M.notify_parser_install_if_needed(bufnr)
  if parser_notified then return end
  if vim.bo[bufnr].filetype ~= "scl" then return end
  local ok = pcall(vim.treesitter.get_parser, bufnr, "scl")
  if ok then return end
  parser_notified = true
  vim.schedule(function()
    local msg = "SCL: Tree-sitter parser not found. "
    if vim.fn.exepath("tia-lsp") ~= "" then
      msg = msg .. "Run :MasonInstall tia-lsp to install it."
    else
      msg = msg .. "Install tia-lsp via mason or compile parser.so manually."
    end
    vim.notify(msg, vim.log.levels.INFO)
  end)
end

function M.setup_syntax(opts)
  -- Load the compiled tree-sitter parser (.so) via Neovim's built-in API.
  -- This bypasses nvim-treesitter's install system, which reloads its
  -- parsers module from source on :TSInstall and wipes runtime config.
  local parser_so = M.find_parser_so(opts)
  if parser_so then
    local ok, err = pcall(vim.treesitter.language.add, "scl", { path = parser_so })
    if not ok then
      vim.notify(
        "SCL: Failed to load tree-sitter parser from " .. parser_so .. ": " .. tostring(err),
        vim.log.levels.WARN
      )
    end
  end

  -- Also register with nvim-treesitter (best-effort, for :TSInstall support
  -- in sessions where the user wants to install from source).
  local has_ts, parsers = pcall(require, "nvim-treesitter.parsers")
  if has_ts and type(parsers) == "table" then
    local parser_config
    if type(parsers.get_parser_configs) == "function" then
      parser_config = parsers.get_parser_configs()
    else
      parser_config = parsers
    end
    if type(parser_config) == "table" then
      parser_config.scl = {
        install_info = {
          url = opts.ts_parser_url
            or "https://gitea.l-tech.rs/lazar/tia-lsp.git",
          files = { "src/parser.c" },
          generate_requires_npm = false,
          requires_generate_from_grammar = false,
        },
        filetype = "scl",
      }
    end
  end

  if opts.cmp then
    local function try_register()
      local cmp_ok, cmp = pcall(require, "blink.cmp")
      local blink_source_ok, blink_source = pcall(require, "tia.blink_cmp_source")
      if cmp_ok and cmp and cmp.add_source_provider and blink_source_ok then
        local source_config = {
          name = "scl",
          provider = blink_source,
          score_offset = 100,
        }
        pcall(cmp.add_source_provider, "scl", source_config)
        pcall(cmp.add_filetype_source, "scl", "scl")

        vim.notify("SCL: blink.cmp source registered", vim.log.levels.INFO)
        return true
      end
      return false
    end

    if not try_register() then
      vim.api.nvim_create_autocmd("BufEnter", {
        pattern = "*.scl",
        once = true,
        callback = try_register,
      })
    end
  end

  if opts.workspace_types ~= false then
    local ok, workspace_types = pcall(require, "tia.workspace_types")
    if ok then
      workspace_types.setup({
        project_root = opts.project_root,
        library_paths = opts.library_paths,
        debug = opts.debug,
      })
    end
  end
end

function M.setup_auto_prefix()
  local ok, auto_prefix = pcall(require, "tia.auto_prefix")
  if ok then
    auto_prefix.setup()
  end
end

function M.setup_multiline_params()
  local ok, mp = pcall(require, "tia.multiline_params")
  if not ok then
    vim.notify("SCL multiline_params module not found", vim.log.levels.WARN)
    return
  end

  vim.api.nvim_create_user_command("SCLMultilineParams", function()
    mp.fill_multiline_params()
  end, {})

  vim.keymap.set("i", "<Space>mp", function()
    mp.fill_multiline_params()
  end, { noremap = true, silent = true, desc = "SCL: Fill multiline parameters" })

  vim.keymap.set("i", "<Tab>", function()
    local ft = vim.bo.filetype
    if ft == "scl" or ft == "udt" then
      if mp.jump_to_next_param() then
        return ""
      end
    end
    return "<Tab>"
  end, { noremap = true, silent = true, desc = "SCL: Jump to next param or Tab" })
end

function M.prefix_current_word()
  local ok, auto_prefix = pcall(require, "tia.auto_prefix")
  if ok then
    auto_prefix.prefix_word()
  end
end

function M.show_variables()
  local ok, variables = pcall(require, "tia.variables")
  if not ok then
    vim.notify("SCL variables module not found", vim.log.levels.WARN)
    return
  end

  local vars = variables.extract_variables()
  if #vars == 0 then
    vim.notify("No local variables found", vim.log.levels.INFO)
    return
  end

  local lines = { "Local Variables:", "---" }
  for _, var in ipairs(vars) do
    local type_info = var.type and (" : " .. var.type) or ""
    table.insert(lines, "#" .. var.name .. type_info)
  end
  vim.notify(table.concat(lines, "\n"), vim.log.levels.INFO)
end

function M.show_workspace_types()
  local ok_ws, ws = pcall(require, "tia.workspace_types")
  local ok_udt, udt = pcall(require, "tia.udt_parser")
  local ok_fb, fb = pcall(require, "tia.fb_parser")

  if not ok_ws or not ok_udt or not ok_fb then
    vim.notify("SCL workspace modules not found", vim.log.levels.WARN)
    return
  end

  local root = ws.get_project_root()
  local udt_names = udt.get_all_udt_names()
  local fb_names = fb.get_all_fb_names()

  if not root then
    vim.notify("No project root detected", vim.log.levels.WARN)
    return
  end

  vim.notify("Project: " .. root .. "\nUDTs: " .. #udt_names .. "\nFBs/Functions: " .. #fb_names, vim.log.levels.INFO)
end

function M.rescan_workspace_types()
  local ok_ws, ws = pcall(require, "tia.workspace_types")
  local ok_udt, udt = pcall(require, "tia.udt_parser")
  local ok_fb, fb = pcall(require, "tia.fb_parser")

  if not ok_ws or not ok_udt or not ok_fb then
    vim.notify("SCL workspace modules not found", vim.log.levels.WARN)
    return
  end

  udt.clear_cache()
  fb.clear_cache()
  ws.clear_root_cache()

  local udt_result = ws.scan_project_udts()
  local fb_result = ws.scan_project_fbs()

  vim.notify("Scanned UDTs: " .. udt_result.parsed_count .. "/" .. udt_result.total_files ..
            ", FBs: " .. fb_result.parsed_count .. "/" .. fb_result.total_files, vim.log.levels.INFO)
end

function M.create_commands()
  vim.api.nvim_create_user_command("SCLShowVariables", function() M.show_variables() end, {})
  vim.api.nvim_create_user_command("SCLShowWorkspaceTypes", function() M.show_workspace_types() end, {})
  vim.api.nvim_create_user_command("SCLRescanWorkspaceTypes", function() M.rescan_workspace_types() end, {})
  vim.api.nvim_create_user_command("SCLPrefixWord", function() M.prefix_current_word() end, {})
  
  -- Attribute toggle commands
  vim.api.nvim_create_user_command("SCLToggleAttrBlock", function()
    local ok, at = pcall(require, "tia.attr_toggle")
    if ok then
      at.toggle_attr_block()
    else
      vim.notify("SCL: Failed to load attr_toggle module: " .. tostring(at), vim.log.levels.ERROR)
    end
  end, {})
  
  vim.api.nvim_create_user_command("SCLExpandAllAttrBlocks", function()
    local ok, at = pcall(require, "tia.attr_toggle")
    if ok then
      at.expand_all_attr_blocks()
    else
      vim.notify("SCL: Failed to load attr_toggle module: " .. tostring(at), vim.log.levels.ERROR)
    end
  end, {})
  
  vim.api.nvim_create_user_command("SCLCollapseAllAttrBlocks", function()
    local ok, at = pcall(require, "tia.attr_toggle")
    if ok then
      at.collapse_all_attr_blocks()
    else
      vim.notify("SCL: Failed to load attr_toggle module: " .. tostring(at), vim.log.levels.ERROR)
    end
  end, {})
end

function M.setup_attr_toggle()
  local ok, at = pcall(require, "tia.attr_toggle")
  if not ok then
    vim.notify("SCL: attr_toggle module not found", vim.log.levels.WARN)
    return
  end
  
  at.setup()

  local function setup_buffer_keymaps(bufnr)
    local opts = { buffer = bufnr, silent = true }
    local ft = vim.bo[bufnr].filetype
    
    if ft ~= "scl" and ft ~= "udt" then
      return
    end

    -- Toggle individual attribute block under cursor
    vim.keymap.set({ "n", "i" }, "<Leader>xa", function()
      at.toggle_attr_block()
    end, vim.tbl_extend("force", opts, { desc = "SCL: Toggle attribute block" }))

    -- Expand all collapsed blocks
    vim.keymap.set("n", "<Leader>xae", function()
      at.expand_all_attr_blocks()
    end, vim.tbl_extend("force", opts, { desc = "SCL: Expand all attribute blocks" }))

    -- Collapse all attribute blocks
    vim.keymap.set("n", "<Leader>xac", function()
      at.collapse_all_attr_blocks()
    end, vim.tbl_extend("force", opts, { desc = "SCL: Collapse all attribute blocks" }))
  end

  -- Set up keybindings for SCL files only
  vim.api.nvim_create_autocmd("FileType", {
    pattern = { "scl", "udt" },
    callback = function(args)
      setup_buffer_keymaps(args.buf)
    end,
  })
  
  -- Also check if there are already SCL buffers open
  for _, bufnr in ipairs(vim.api.nvim_list_bufs()) do
    if vim.api.nvim_buf_is_loaded(bufnr) then
      setup_buffer_keymaps(bufnr)
    end
  end
end

return M