/**
 * Pulpie WebUI - Block Chunker
 *
 * Implements token sequence packing matching Pulpie's exact chunker algorithm.
 * Groups blocks into sequences bounded by max_tokens (default 8,192).
 *
 * Sequence structure:
 *   [BOS] block_0 [<|sep|>] block_1 [<|sep|>] ... block_N [<|sep|>] [EOS]
 */

import { TOKEN_IDS } from './tokenizer.js';

export const DEFAULT_MAX_TOKENS = 1024;

/**
 * @typedef {Object} Chunk
 * @property {number[]} chunkIds - Array of token IDs including BOS, EOS, and SEP tokens
 * @property {number[]} blockIndices - Indices of the blocks included in this chunk
 */

/**
 * Pack block tokens into chunks conforming to EuroBERT's max_tokens limit.
 *
 * @param {number[][]} blockTokenIds - Array of token ID lists for each block
 * @param {Object} [options]
 * @param {number} [options.maxTokens=8192] - Maximum sequence length per forward pass
 * @param {number} [options.sepTokenId] - Token ID for <|sep|> (default 128256)
 * @param {number} [options.bosTokenId] - Token ID for <|begin_of_text|> (default 128000)
 * @param {number} [options.eosTokenId] - Token ID for <|end_of_text|> (default 128001)
 * @returns {Chunk[]}
 */
export function packChunks(
  blockTokenIds,
  {
    maxTokens = DEFAULT_MAX_TOKENS,
    sepTokenId = TOKEN_IDS.SEP,
    bosTokenId = TOKEN_IDS.BOS,
    eosTokenId = TOKEN_IDS.EOS,
  } = {}
) {
  if (!blockTokenIds || blockTokenIds.length === 0) {
    return [];
  }

  const chunks = [];
  let currentTokens = [];
  let currentBlockIndices = [];

  const overhead = (bosTokenId !== null ? 1 : 0) + (eosTokenId !== null ? 1 : 0);
  const budget = maxTokens - overhead;
  let currentLen = 0;

  for (let blockIdx = 0; blockIdx < blockTokenIds.length; blockIdx++) {
    const blockToks = blockTokenIds[blockIdx];
    const blockCost = blockToks.length + 1; // +1 for SEP token

    if (currentLen + blockCost > budget) {
      if (currentTokens.length > 0) {
        const chunkIds = [];
        if (bosTokenId !== null) chunkIds.push(bosTokenId);
        chunkIds.push(...currentTokens);
        if (eosTokenId !== null) chunkIds.push(eosTokenId);

        chunks.push({
          chunkIds,
          blockIndices: [...currentBlockIndices],
        });
      }

      currentTokens = [];
      currentBlockIndices = [];
      currentLen = 0;

      // Single block exceeds total budget on its own
      if (blockCost > budget) {
        const truncated = blockToks.slice(0, budget - 1);
        currentTokens.push(...truncated, sepTokenId);
        currentBlockIndices.push(blockIdx);
        currentLen = truncated.length + 1;
        continue;
      }
    }

    currentTokens.push(...blockToks, sepTokenId);
    currentBlockIndices.push(blockIdx);
    currentLen += blockCost;
  }

  if (currentTokens.length > 0) {
    const chunkIds = [];
    if (bosTokenId !== null) chunkIds.push(bosTokenId);
    chunkIds.push(...currentTokens);
    if (eosTokenId !== null) chunkIds.push(eosTokenId);

    chunks.push({
      chunkIds,
      blockIndices: [...currentBlockIndices],
    });
  }

  return chunks;
}