--[[
Static embedding provider for neural feature fusion.

The cheap, multilingual successor to fasttext_embed: words produced by
rspamd's regular tokenization pipeline are re-tokenized into WordPiece
subword tokens and embedded by mean-pooling rows of a static
token-embedding matrix (Model2Vec style). No neural forward pass is
involved; all heavy lifting (tokenizer, mmap-ed float32 matrix, pooling)
lives in the rspamd_static_embed C module, so the model is shared between
workers and never copied into the Lua heap.

The WordPiece tokenizer is internal to the vectorizer: the global
word-breaking / statistics tokenization is not affected, so Bayes tokens
and fuzzy hashes stay exactly as they were.

A model directory must contain config.json, vocab.txt and the matrix file
(see the rspamd_static_embed module docs); models are shipped separately
as data, like FastText models. Any deviation from the supported spec
disables the provider with an explicit error - there is no silent
fallback.

Note: the rspamd_static_embed module also exposes per-token sequence
access (model:get_token_vectors) for external consumers such as offline
trainers exporting order-aware text features. The provider path is not
affected: fusion vectors must stay fixed-dim, so only the pooled
get_sentence_vector is used here.

Configuration example in neural.conf:
  providers = [
    {
      type = "static_embed";
      model = "/path/to/model_dir";
      weight = 1.0;
    }
  ];
]] --

local rspamd_logger = require "rspamd_logger"
local lua_mime = require "lua_mime"

local N = "neural.static_embed"

local exports = {}

-- May be nil on incomplete builds; checked in load_model
local se_ok, rspamd_static_embed = pcall(require, "rspamd_static_embed")

-- Cache of loaded models: dir -> model; load errors are cached too so that
-- a broken config is reported once instead of on every scanned message
local loaded_models = {}
local failed_models = {}

exports.load_model = function(dir)
  if loaded_models[dir] then
    return loaded_models[dir]
  end
  if failed_models[dir] then
    return nil, failed_models[dir]
  end

  local model, err
  if not se_ok then
    err = 'rspamd_static_embed module is not available'
  else
    model, err = rspamd_static_embed.load(dir)
  end

  if not model then
    failed_models[dir] = err or 'unknown error'
    return nil, failed_models[dir]
  end

  loaded_models[dir] = model
  return model
end

-- Extract words exactly like fasttext_embed does
-- Append the message's words to `words`, up to opts.max_words TOTAL entries
-- (0/absent = unbounded). The cap bounds the Lua-side materialisation and the
-- embedding cost: without it a multi-megabyte text part turns into a
-- million-entry Lua string table plus a mean-pool over every token — observed
-- at hundreds of ms per extra megabyte inside a synchronous rule. A mean over
-- the first few hundred words is statistically indistinguishable for
-- classification, but pick the cap consciously: it changes the vectors of
-- long texts, so train and inference must agree on it.
local function extract_words(task, opts, words)
  words = words or {}
  local how = opts.word_form or 'norm'
  local cap = tonumber(opts.max_words) or 0

  local parts
  if opts.all_parts then
    parts = task:get_text_parts()
  else
    local sel = lua_mime.get_displayed_text_part(task)
    if sel then
      parts = { sel }
    else
      parts = task:get_text_parts()
    end
  end

  if not parts then
    return words
  end

  for _, part in ipairs(parts) do
    local pw = part:get_words(how)
    if pw then
      for _, w in ipairs(pw) do
        if type(w) == 'string' and #w > 0 then
          words[#words + 1] = w
          if cap > 0 and #words >= cap then
            return words
          end
        end
      end
    end
  end

  return words
end

-- Provider registration is skipped when neural is not loadable (e.g. in
-- unit tests); the exported helpers are still usable in that case
local neural_ok, neural_common = pcall(require, "plugins/neural")

if neural_ok then
  neural_common.register_provider('static_embed', {
    init = function(pcfg)
      if not pcfg.model then
        rspamd_logger.errx(rspamd_config, '%s: no model directory specified', N)
        return
      end

      local model, err = exports.load_model(pcfg.model)
      if model then
        rspamd_logger.infox(rspamd_config, '%s: loaded model from %s: %s tokens, dim=%s',
          N, pcfg.model, model:get_vocab_size(), model:get_dimension())
      else
        rspamd_logger.errx(rspamd_config, '%s: cannot load model from %s: %s; provider disabled',
          N, pcfg.model, err)
      end
    end,
    collect_async = function(task, ctx, cont)
      local pcfg = ctx.config or {}

      local model = pcfg.model and exports.load_model(pcfg.model) or nil
      if not model then
        rspamd_logger.debugm(N, task, 'static_embed: no model available; skip')
        cont(nil)
        return
      end

      -- Subject words go FIRST (kept from the original prepend semantics),
      -- so seed the table with them and let extract_words append the body
      -- up to the cap. The old `table.insert(words, 1, w)` prepend was
      -- O(#words) per subject word — a ten-word subject over a large
      -- message's word table meant millions of element shifts.
      local words = {}
      if pcfg.include_subject ~= false then
        local subj = task:get_subject()
        if subj and #subj > 0 then
          for w in subj:gmatch('%S+') do
            words[#words + 1] = w
          end
        end
      end
      extract_words(task, {
        word_form = pcfg.word_form or 'norm',
        all_parts = pcfg.all_parts,
        max_words = pcfg.max_words,
      }, words)

      -- Empty input produces a zero vector: dimensionality stays stable
      local vec, ntokens = model:get_sentence_vector(words)

      local meta = {
        name = pcfg.name or 'static_embed',
        type = 'static_embed',
        dim = model:get_dimension(),
        weight = ctx.weight or 1.0,
        tokens = ntokens,
      }

      rspamd_logger.debugm(N, task, 'static_embed: produced %s-dim vector from %s tokens (%s words)',
        meta.dim, ntokens, #words)
      cont(vec, meta)
    end,
  })
end

return exports
