-- nlbwmon.lua (Corrected nlbw command for CSV output)

-- Helper function to read /tmp/dhcp.leases
local function get_mac_to_hostname_map()
  local mac_map = {}
  local leases_file = "/tmp/dhcp.leases"
  local file = io.open(leases_file, "r")

  if not file then
    io.stderr:write("[nlbwmon_collector] ERROR: Could not open DHCP leases file: " .. leases_file .. "\n")
    return mac_map
  end

  for line in file:lines() do
    local _, mac, _, hostname, _ = line:match("^(%d+)%s+([%x:]+)%s+([%d%.]+)%s+([^%s]*)%s+([^%s]*)$")
    if mac and hostname then
      local clean_hostname = (hostname == "*" or hostname == "" or hostname == nil) and "unknown_host" or hostname
      mac_map[string.upper(mac)] = clean_hostname
    elseif mac then
      mac_map[string.upper(mac)] = mac_map[string.upper(mac)] or "unknown_host"
    end
  end
  file:close()
  return mac_map
end

-- Main scrape function
local function scrape()
  -- Define Prometheus metrics local to the scrape function
  local metric_received_bytes = metric("nlbwmon_received_bytes_total", "counter")
  local metric_transmitted_bytes = metric("nlbwmon_transmitted_bytes_total", "counter")

  local mac_host_map = get_mac_to_hostname_map()

  -- Corrected nlbw command to explicitly set comma as separator
  local cmd = "nlbw -c csv -g mac -o mac -n -q -s,"
  local handle = io.popen(cmd)

  if not handle then
    io.stderr:write("[nlbwmon_collector] ERROR: Failed to execute nlbw command: " .. cmd .. "\n")
    return -- Exit scrape if command fails
  end

  local line_num = 0
  for line in handle:lines() do
    line_num = line_num + 1

    -- This header check expects actual commas if "-s," works
    if line_num == 1 and line == "mac,conns,rx_bytes,rx_pkts,tx_bytes,tx_pkts" then
      -- io.stderr:write("[nlbwmon_collector] Correct CSV header found and skipped.\n")
    else
      if line_num == 1 then -- It wasn't the expected comma-separated header
        io.stderr:write(
        "[nlbwmon_collector] WARNING: Unexpected first line from nlbw. Expected 'mac,conns,rx_bytes,rx_pkts,tx_bytes,tx_pkts'. Got: '" ..
        line .. "'\n")
      end

      -- This regex expects comma-separated values
      local mac_addr, conns, rx_bytes_str, rx_pkts, tx_bytes_str, tx_pkts_str = line:match(
      "^([%x:]+),([^,]*),([^,]*),([^,]*),([^,]*),([^%s]*)$")

      if mac_addr then -- At least MAC was parsed (implies commas were present if further fields are needed)
        local current_mac = string.upper(mac_addr)
        local received_val = tonumber(rx_bytes_str)
        local transmitted_val = tonumber(tx_bytes_str)

        if received_val ~= nil and transmitted_val ~= nil then
          local current_host
          if current_mac == "00:00:00:00:00:00" then
            current_host = "unknown_mac_bucket"
          else
            current_host = mac_host_map[current_mac] or "unknown_host"
          end

          local labels = {
            mac = current_mac,
            host = current_host
          }
          metric_received_bytes(labels, received_val)
          metric_transmitted_bytes(labels, transmitted_val)
        else
          io.stderr:write(string.format(
          "[nlbwmon_collector] WARNING: Could not parse byte values to number for MAC: %s. RX_STR: '%s', TX_STR: '%s'. Line: %s\n",
            current_mac, tostring(rx_bytes_str), tostring(tx_bytes_str), line))
        end
      else
        -- Only log this error if it wasn't the (potentially now correctly handled) header line
        if line_num > 1 or (line_num == 1 and line ~= "mac,conns,rx_bytes,rx_pkts,tx_bytes,tx_pkts") then
          io.stderr:write("[nlbwmon_collector] WARNING: Failed to parse nlbw data line structure (expected CSV): " ..
          line .. "\n")
        end
      end
    end
  end
  local close_status, close_err_msg, close_code = handle:close()
  if not close_status then
    io.stderr:write(string.format("[nlbwmon_collector] WARNING: nlbw command close error. Msg: %s, Code: %s\n",
      tostring(close_err_msg), tostring(close_code)))
  elseif type(close_status) == "number" and close_status ~= 0 then
    io.stderr:write(string.format("[nlbwmon_collector] WARNING: nlbw command exited with non-zero status: %s\n",
      tostring(close_status)))
  end
end

-- Return the scrape function
return {
  scrape = scrape
}