#!/usr/bin/env lua

-- Copyright (c) Meta Platforms, Inc. and affiliates.
--
-- This source code is licensed under the MIT license found in the
-- LICENSE file in the root directory of this source tree.

--- Writes a custom vpp startup config file using node information.
-- @script update_vpp_startup_conf

local tg_utils = require "tg.utils"
local tg_net_utils = require "tg.net_utils"
local logger = require "tg.logger"
local argparse = require "argparse"
local dir = require "pl.dir"
local path = require "pl.path"
local Template = (require "pl.text").Template
local tablex = require "pl.tablex"
require("pl.stringx").import()

local C = {}

-- File paths
local CONFIG_FILE = "/data/cfg/node_config.json"
local NODE_INFO_FILE = "/var/run/node_info"
local VPP_STARTUP_CONF_FILE = "/var/run/vpp/startup.conf"
local VPP_STARTUP_TEMPLATE_FILE = "/etc/vpp/startup.conf.template"

-- Template strings
local DEVICE_FORMAT = Template(
  [[dev ${bus}:00.0 {
    workers ${workerId}
    hqos
    ${devargs}
  }]])
local ESMC_FORMAT = Template([[
esmc {
  on
  input ${input_interface}
  output ${output_interface}
  enable-tg-input
  enable-tg-output
}
]])
local PTPTC_FORMAT = Template([[
ptptc {
  interface ${interface}
  offset-ns ${offset_ns}
${port_list}
  rx-checksum-offload
  rx-ptp-classify-offload
}
]])
local PTPTC_PORT_FORMAT = Template("  port ${port}")
local INTERFACE_FORMAT = Template([[
  interface ${vpp_interface} {
    ${tap_interface}
    ipv6 slowpath off
    wired security on
  }
]])
local TGHQOS_FORMAT = Template([[
tghqos {
  fw-tx-ring-depth ${fw_tx_ring_depth}
}
]])

-- Constants
local DEFAULT_FW_TX_RING_DEPTH = 256
local DEFAULT_PORT_MAP = {
  -- Wigigs assigned to worker 1, 10G interface defaults to worker 0
  ["0000:01"] = {worker_id = 1},
  ["0001:01"] = {worker_id = 1},
  ["0002:03"] = {worker_id = 1},
  ["0002:04"] = {worker_id = 1},
}
local CONF_FILE_HEADER = [[
# NOTE: This file is automatically generated. DO NOT EDIT MANUALLY!
]]
local COREDUMP_ENABLED_CONF = [[
  coredump-size unlimited
  full-coredump
]]
local COREDUMP_DISABLED_CONF = [[
  # Disable coredumps
  coredump-size 1
]]

function C.buildDeviceConfigs(nodeInfo, nodeConfig)
  local useEepromMacs =
    tg_utils.get(nodeConfig, "envParams", "VPP_USE_EEPROM_MACS") == "1"
  local fwRecoveryEnabled =
    tg_utils.get(nodeConfig, "envParams", "VPP_FW_RECOVERY_ENABLED") == "1"

  local deviceConf = {}
  local num = tg_utils.get(nodeInfo, "NUM_WLAN_MACS")

  -- Build buslist
  local busList = {}
  if num == nil or tonumber(num) == 0 then
    busList = tablex.keys(DEFAULT_PORT_MAP)
    table.sort(busList)
    num = #busList
  else
    for i = 0, tonumber(num) - 1 do
      busList[#busList+1] = tg_utils.get(nodeInfo, "BUS_" .. i)
    end
  end

  for i = 0, tonumber(num) - 1 do
    local portParams = DEFAULT_PORT_MAP[busList[i + 1]]
    local workerId = portParams and portParams.worker_id or 1
    local bus = busList[i+1]
    local devargs = {}

    if useEepromMacs then
      local mac = tg_utils.get(nodeInfo, "MAC_" .. i)
      if mac and tg_net_utils.isMacAddr(mac) then
        devargs[#devargs+1] = "mac-address=" .. mac
      end
    end

    if not fwRecoveryEnabled then
      devargs[#devargs+1] = "no-fw-recovery=1"
    end

    devargs[#devargs+1] = "crash-on-fw-err=0"

    -- Enable non-commercial use flag for T-Bias
    devargs[#devargs+1] = "non-commercial-rf=1"

    -- Enable PMC mode firmware logging at startup to ensure early boot logs are
    -- fully captured
    devargs[#devargs+1] = "pmc-ext-host=1"

    local devargsLine
    if #devargs > 0 then
      devargsLine = "devargs " .. (","):join(devargs)
    else
      devargsLine = ""
    end

    deviceConf[#deviceConf+1] = DEVICE_FORMAT:substitute{
      bus = bus,
      devargs = devargsLine,
      workerId = workerId,
    }
  end
  return deviceConf
end

function C.buildVppStartupConfig(startupTemplate, nodeConfig, deviceConf)
  if startupTemplate == nil then
    return nil
  end

  -- Fill config parameters
  local coredumpConfig
  if tg_utils.get(nodeConfig, "envParams", "VPP_COREDUMP_ENABLED") == "1" then
    coredumpConfig = COREDUMP_ENABLED_CONF
  else
    coredumpConfig = COREDUMP_DISABLED_CONF
  end

  local ulaPrefix
  local cfg = tg_utils.get(nodeConfig, "envParams", "VPP_ULA_TEST_PREFIX")
  if cfg == nil or cfg:len() == 0 then
    ulaPrefix = ""
  else
    ulaPrefix = "ula-test-prefix " .. cfg
  end

  local wsecEnable
  cfg = tg_utils.get(nodeConfig, "radioParamsBase", "fwParams", "wsecEnable")
  if cfg ~= nil and cfg > 0 then
    wsecEnable = "wsec on"
  else
    wsecEnable = "wsec off"
  end

  local interfaceConfig = {}
  local cpeConfig = tg_utils.get(nodeConfig, "cpeConfig")
  for intf, ifCfg in pairs(cpeConfig or {}) do
    if ifCfg.wiredSecurityEnable == true then
      local tapInterface = "tap nic" .. string.match(intf, "%d+")

      interfaceConfig[#interfaceConfig+1] =
        INTERFACE_FORMAT:substitute{
          vpp_interface = intf,
          tap_interface = tapInterface
        }
    end
  end

  local esmcConfig
  if tg_utils.get(nodeConfig, "timingParams", "ESMC_ENABLED") == "1" then
    local port = tg_utils.get(nodeConfig, "timingParams", "PTP_VSC_PORT") or "0"
    local esmcIface = (port == "0")
      and "TenGigabitEthernet1"
      or "TenGigabitEthernet0"
    esmcConfig = ESMC_FORMAT:substitute{
      input_interface = esmcIface, output_interface = esmcIface
    }
  else
    esmcConfig = ""
  end

  local ptptcConfig
  local ptpIface =
    tg_utils.get(nodeConfig, "timingParams", "PTP_VPP_INTERFACE") or ""
  if ptpIface ~= "" then
    local ptpOffsetNs =
      tg_utils.get(nodeConfig, "timingParams", "PTP_VPP_OFFSET_NS") or "0"
    local ptpNxpPort =
      tg_utils.get(nodeConfig, "timingParams", "PTP_VPP_NXP_PORT") or ""
    local portList = {}
    if ptpNxpPort ~= "" then
      for _, port in ipairs(ptpNxpPort:split(",")) do
        portList[#portList+1] = PTPTC_PORT_FORMAT:substitute{port = port}
      end
    end
    ptptcConfig = PTPTC_FORMAT:substitute{
      interface = ptpIface,
      offset_ns = ptpOffsetNs,
      port_list = ("\n"):join(portList)
    }
  else
    ptptcConfig = ""
  end

  local tghqosFwTxRingDepth = tg_utils.get(
    nodeConfig, "envParams", "FW_TX_RING_DEPTH"
  ) or DEFAULT_FW_TX_RING_DEPTH
  local tghqosConfig = TGHQOS_FORMAT:substitute{
    fw_tx_ring_depth = tghqosFwTxRingDepth
  }

  local dpdkLogConfig = {}
  local dpdkLogStrs =
    tg_utils.get(nodeConfig, "envParams", "VPP_DPDK_LOG_STR") or ""
  for _, logStr in ipairs(dpdkLogStrs:split(",")) do
    dpdkLogConfig[#dpdkLogConfig+1] = "log-level " .. logStr
  end

  return Template(startupTemplate):substitute{
    coredump_config = coredumpConfig,
    devices = ("\n  "):join(deviceConf),
    dpdk_log_config = ("\n  "):join(dpdkLogConfig),
    esmc_config = esmcConfig,
    ptptc_config = ptptcConfig,
    tghqos_config = tghqosConfig,
    ula_test_prefix = ulaPrefix,
    wsec_enable = wsecEnable,
    interface_config = (""):join(interfaceConfig)
  }
end

function C.processTemplate(startupTemplate)
  if startupTemplate == nil then
    return nil
  end

  -- Replace header comment block
  local lines = startupTemplate:splitlines(true)
  local stripTo = 0
  for i, line in ipairs(lines) do
    if line:startswith("#") then
      stripTo = i
    else
      break
    end
  end
  for i = 1, stripTo do
    table.remove(lines, 1)
  end
  table.insert(lines, 1, CONF_FILE_HEADER)

  return table.concat(lines)
end

function C.main()
  local parser = argparse(
    "update_vpp_startup_conf",
    "Writes a custom vpp startup config file using node information."
  )
  parser:option(
    "-n --node_config_file", "Path to the node configuration file", CONFIG_FILE
  )
  parser:option(
    "-i --node_info_file", "Path to the node info file", NODE_INFO_FILE
  )
  parser:option(
    "-t --vpp_startup_conf_template",
    "Path to the vpp startup configuration template", VPP_STARTUP_TEMPLATE_FILE
  )
  parser:option(
    "-o --vpp_startup_conf_file", "Output vpp startup configuration file path",
    VPP_STARTUP_CONF_FILE
  )
  parser:flag("-v --verbose", "Enable logging", false)
  local args = parser:parse()

  -- Set up logging
  if not args.verbose then
    logger.level = "critical"
  end
  logger.startSyslog()
  local exit = function(code)
    logger.stopSyslog()
    os.exit(code)
  end

  -- Read input files
  local startupTemplate = tg_utils.readFile(args.vpp_startup_conf_template)
  if startupTemplate == nil then
    logger.error(
      "Failed to read template file %s", args.vpp_startup_conf_template
    )
    exit(1)
  end
  local nodeConfig = tg_utils.readJsonFile(args.node_config_file)
  if nodeConfig == nil then
    logger.error("Failed to read node config file %s", args.node_config_file)
    nodeConfig = {}
  end
  local nodeInfo = tg_utils.readEnvFile(args.node_info_file)
  if nodeInfo == nil then
    logger.error("Failed to read node info file %s", args.node_info_file)
    nodeInfo = {}
  end

  -- Write VPP config
  local newVppConfig = C.buildVppStartupConfig(
    C.processTemplate(startupTemplate),
    nodeConfig,
    C.buildDeviceConfigs(nodeInfo, nodeConfig)
  )
  logger.info(
    "Writing vpp startup configuration to file %s",
    args.vpp_startup_conf_file
  )
  dir.makepath(path.dirname(args.vpp_startup_conf_file))
  if not tg_utils.writeFile(args.vpp_startup_conf_file, newVppConfig) then
    exit(1)
  end

  logger.stopSyslog()
end

if tg_utils.isMain() then
  C.main()
else
  return C
end
