import type { Environment } from '@marcbachmann/cel-js';

import {
  mulberry32,
  normalizeSeed,
  sampleWithoutReplacement,
  sampleWithReplacement,
  type RandomSource,
} from './prng';

/**
 * Random integer with exactly `digits` digits, e.g. random(5) -> 48574.
 * Ported verbatim from the backend's cel_evaluator for back-compat with
 * deployed workflows. Float precision degrades above 15 digits.
 */
function randomDigits(digits: bigint | number): bigint {
  const d = Number(digits);
  if (d <= 0) return 0n;
  const min = 10 ** (d - 1);
  const max = 10 ** d;
  return BigInt(Math.floor(Math.random() * (max - min)) + min);
}

function randomSampleImpl(
  list: unknown,
  count: bigint | number,
  seed: bigint | number | undefined,
  withReplacement: boolean,
): unknown[] {
  if (!Array.isArray(list)) return [];
  const rand: RandomSource = seed === undefined ? Math.random : mulberry32(normalizeSeed(seed));
  const k = Number(count);
  return withReplacement ? sampleWithReplacement(list, k, rand) : sampleWithoutReplacement(list, k, rand);
}

const RANDOM_SAMPLE_DESCRIPTION =
  'Sample items from a list. An int seed makes the result reproducible (same seed, same sample); ' +
  'withReplacement allows duplicates. Without replacement, count is clamped to the list size.';

const PARAM_LIST = { name: 'list', description: 'List to sample from' };
const PARAM_COUNT = { name: 'count', description: 'Number of items to sample' };
const PARAM_SEED = { name: 'seed', description: 'Seed for reproducible sampling' };
const PARAM_WITH_REPLACEMENT = { name: 'withReplacement', description: 'Allow duplicate picks' };

/** Register the shared base functions. Order and semantics are stable API. */
export function registerBaseFunctions(env: Environment): Environment {
  // Registered as dyn so cel-js doesn't reject null before reaching the handler.
  env.registerFunction('trim(dyn): string', (s: string) => (s == null ? '' : String(s).trim()), {
    description: 'Trim leading and trailing whitespace. Null-safe: null becomes "".',
    params: [{ name: 'value', description: 'Value to trim; null yields ""' }],
  });

  // Registered as dyn first arg so cel-js doesn't reject non-list inputs before the handler.
  env.registerFunction(
    'slice(dyn, int, int): list',
    (arr: unknown, start: bigint | number, end: bigint | number) => {
      if (!Array.isArray(arr)) return [];
      return arr.slice(Number(start), Number(end));
    },
    {
      description: 'Return list[start:end] with JS Array.prototype.slice semantics.',
      params: [
        { name: 'list', description: 'List to slice' },
        { name: 'start', description: 'Start index (inclusive)' },
        { name: 'end', description: 'End index (exclusive)' },
      ],
    },
  );

  env.registerFunction('random(int): int', (digits: bigint | number) => randomDigits(digits), {
    description: 'Random integer with exactly N digits, e.g. random(5) -> 48574. Unseeded.',
    params: [{ name: 'digits', description: 'Number of digits in the result' }],
  });

  env.registerFunction(
    'randomSample(list, int): list',
    (list: unknown, count: bigint | number) => randomSampleImpl(list, count, undefined, false),
    { description: RANDOM_SAMPLE_DESCRIPTION, params: [PARAM_LIST, PARAM_COUNT] },
  );

  env.registerFunction(
    'randomSample(list, int, int): list',
    (list: unknown, count: bigint | number, seed: bigint | number) =>
      randomSampleImpl(list, count, seed, false),
    { description: RANDOM_SAMPLE_DESCRIPTION, params: [PARAM_LIST, PARAM_COUNT, PARAM_SEED] },
  );

  env.registerFunction(
    'randomSample(list, int, bool): list',
    (list: unknown, count: bigint | number, withReplacement: boolean) =>
      randomSampleImpl(list, count, undefined, withReplacement),
    {
      description: RANDOM_SAMPLE_DESCRIPTION,
      params: [PARAM_LIST, PARAM_COUNT, PARAM_WITH_REPLACEMENT],
    },
  );

  env.registerFunction(
    'randomSample(list, int, int, bool): list',
    (list: unknown, count: bigint | number, seed: bigint | number, withReplacement: boolean) =>
      randomSampleImpl(list, count, seed, withReplacement),
    {
      description: RANDOM_SAMPLE_DESCRIPTION,
      params: [PARAM_LIST, PARAM_COUNT, PARAM_SEED, PARAM_WITH_REPLACEMENT],
    },
  );

  return env;
}