Produce test data sets

"""
Sample a test data set from the distributions observed in production.

Rung N1. Standard library only, no file written, and still no network: the
distributions come from the caller, as the result of a `GROUP BY` would.

What this buys over N0. A schema says a quantity is between one and nine; it
does not say that most orders are for one item and almost none for nine. Data
drawn uniformly inside the bounds gives every rare case the same weight as the
common one, so a cache hit rate measured on it means nothing, a page laid out
for it looks nothing like the real one, and the slow query stays fast.

What is sampled here is the marginal distribution of each column, one column at
a time: the observed count per category, and the observed count per bucket of a
numeric column. Nothing is fitted and nothing is learnt — the observed table is
the model. Which also means the joint distribution is lost, and the tests next
to this file show what that costs.

Determinism works exactly as in N0, and for the same reason: a data set that is
only reproducible on average is not reproducible.
"""

FNV_OFFSET = 2166136261
FNV_PRIME = 16777619
MASK32 = 0xFFFFFFFF

UNIT = "\x1f"  # separates the parts of a cell key, and appears in none of them
ESCAPE = "\x1e"  # stands in for UNIT inside a part, and is escaped itself


def stable_hash(text: str) -> int:
    """FNV-1a on 32 bits, identical to the JavaScript version of this file."""
    digest = FNV_OFFSET
    for char in text:
        digest = ((digest ^ ord(char)) * FNV_PRIME) & MASK32
    return digest


def finalise(digest: int) -> int:
    """
    Scramble the low bits of a hash, the `fmix32` of MurmurHash3.

    Same reason as in rung N0: without it, the low bits of a cell follow the
    last digit of its row number, and two columns drawn from the same seed are
    not independent at all.
    """
    digest ^= digest >> 16
    digest = (digest * 0x85EBCA6B) & MASK32
    digest ^= digest >> 13
    digest = (digest * 0xC2B2AE35) & MASK32
    return digest ^ (digest >> 16)


def _part(text: str) -> str:
    """Keep UNIT out of a part, so that two different cells never share a key."""
    return str(text).replace(ESCAPE, ESCAPE + "0").replace(UNIT, ESCAPE + "1")


def draw(seed: str, field: str, row: int) -> int:
    """One 32-bit integer per cell, independent of the other cells."""
    return finalise(stable_hash(f"{_part(seed)}{UNIT}{_part(field)}{UNIT}{row}"))


def pick(counts: list[int], number: int) -> int:
    """
    Index of the bucket a drawn integer falls into, in proportion to `counts`.

    Cumulating integers rather than normalising to probabilities keeps the
    result exact, and identical in both languages: no floating point is
    involved anywhere in the decision.
    """
    total = sum(counts)
    if total <= 0:
        raise ValueError("a distribution needs at least one observation")
    target = number % total
    for index, count in enumerate(counts):
        if target < count:
            return index
        target -= count
    raise AssertionError("unreachable: the target is below the total")


def _sampler(spec: dict, seed: str, field: str):
    """Read one column's distribution once, and return what draws its value for a row."""
    kind = spec["type"]
    if kind == "categorical":
        # `counts` is a label to observed-count mapping, straight out of a
        # `GROUP BY`. A label seen zero times is never drawn, which is the
        # honest behaviour: it did not happen.
        #
        # The labels are sorted rather than taken in the order the query
        # returned them, so the same observed table always gives the same data.
        # It is also what keeps the two languages together: a JavaScript object
        # reorders its numeric-looking keys, and a postcode is one of those. The
        # key sorts by UTF-16 code unit, as JavaScript's `sort` does, or an emoji
        # and a full-width letter would come out in a different order. Sorting
        # once per column, not once per cell, keeps thousands of postcodes cheap.
        labels = sorted(spec["counts"], key=lambda label: label.encode("utf-16-be", "surrogatepass"))
        weights = [spec["counts"][label] for label in labels]
        return lambda row: labels[pick(weights, draw(seed, field, row))]
    if kind == "histogram":
        edges, counts = spec["edges"], spec["counts"]
        if len(edges) != len(counts) + 1:
            raise ValueError("a histogram needs one more edge than it has buckets")

        def within(row: int):
            bucket = pick(counts, draw(seed, field, row))
            low, high = edges[bucket], edges[bucket + 1]
            # A second, independent draw places the value inside its bucket. Within
            # a bucket the shape is unknown, so uniform is the only honest choice.
            return low + draw(seed, f"{field}{UNIT}within", row) % max(high - low, 1)

        return within
    raise ValueError(f"unknown distribution type {kind!r}")


def sample_rows(distributions: dict, count: int, seed: str) -> list[dict]:
    """
    Return `count` rows, each column sampled from its own observed distribution.

    `distributions` maps a field name to either
    {"type": "categorical", "counts": {label: observed count}} or
    {"type": "histogram", "edges": [...], "counts": [...]}, where the edges are
    integers, the buckets are half-open and `edges` holds one more value than
    `counts`.
    """
    samplers = {field: _sampler(spec, seed, field) for field, spec in distributions.items()}
    return [{field: sample(row) for field, sample in samplers.items()} for row in range(count)]

/**
 * Sample a test data set from the distributions observed in production.
 *
 * Rung N1. No dependency, no file written, and still no network: the
 * distributions come from the caller, as the result of a `GROUP BY` would.
 *
 * What this buys over N0. A schema says a quantity is between one and nine; it
 * does not say that most orders are for one item and almost none for nine.
 * Data drawn uniformly inside the bounds gives every rare case the same weight
 * as the common one, so a cache hit rate measured on it means nothing, a page
 * laid out for it looks nothing like the real one, and the slow query stays
 * fast.
 *
 * What is sampled here is the marginal distribution of each column, one column
 * at a time: the observed count per category, and the observed count per
 * bucket of a numeric column. Nothing is fitted and nothing is learnt — the
 * observed table is the model. Which also means the joint distribution is
 * lost, and the tests next to this file show what that costs.
 *
 * Determinism works exactly as in N0, and for the same reason: a data set that
 * is only reproducible on average is not reproducible.
 */

const FNV_OFFSET = 2166136261;
const FNV_PRIME = 16777619;

// Unit separator: it joins the parts of a cell key, and appears in none of them.
const UNIT = '\u001f';
// Stands in for UNIT inside a part, and is escaped itself.
const ESCAPE = '\u001e';

/** FNV-1a on 32 bits, identical to the Python version of this file. */
export function stableHash(text) {
  let digest = FNV_OFFSET;
  for (const char of text) {
    digest = Math.imul(digest ^ char.codePointAt(0), FNV_PRIME) >>> 0;
  }
  return digest;
}

/**
 * Scramble the low bits of a hash, the `fmix32` of MurmurHash3.
 *
 * Same reason as in rung N0: without it, the low bits of a cell follow the last
 * digit of its row number, and two columns drawn from the same seed are not
 * independent at all.
 */
export function finalise(digest) {
  let h = digest >>> 0;
  h = (h ^ (h >>> 16)) >>> 0;
  h = Math.imul(h, 0x85ebca6b) >>> 0;
  h = (h ^ (h >>> 13)) >>> 0;
  h = Math.imul(h, 0xc2b2ae35) >>> 0;
  return (h ^ (h >>> 16)) >>> 0;
}

/** Keep UNIT out of a part, so that two different cells never share a key. */
function part(text) {
  return String(text).replaceAll(ESCAPE, `${ESCAPE}0`).replaceAll(UNIT, `${ESCAPE}1`);
}

/** One 32-bit integer per cell, independent of the other cells. */
export function draw(seed, field, row) {
  return finalise(stableHash(`${part(seed)}${UNIT}${part(field)}${UNIT}${row}`));
}

/**
 * Index of the bucket a drawn integer falls into, in proportion to `counts`.
 *
 * Cumulating integers rather than normalising to probabilities keeps the
 * result exact, and identical in both languages: no floating point is involved
 * anywhere in the decision.
 */
export function pick(counts, number) {
  const total = counts.reduce((sum, count) => sum + count, 0);
  if (total <= 0) throw new RangeError('a distribution needs at least one observation');
  let target = number % total;
  for (let index = 0; index < counts.length; index += 1) {
    if (target < counts[index]) return index;
    target -= counts[index];
  }
  throw new Error('unreachable: the target is below the total');
}

/** Read one column's distribution once, and return what draws its value for a row. */
function sampler(spec, seed, field) {
  switch (spec.type) {
    case 'categorical': {
      // `counts` maps a label to its observed count, straight out of a
      // `GROUP BY`. A label seen zero times is never drawn, which is the
      // honest behaviour: it did not happen.
      //
      // The labels are sorted rather than taken in the order the query
      // returned them, so the same observed table always gives the same data.
      // It is also what keeps the two languages together: an object reorders
      // its numeric-looking keys, and a postcode is one of those. `sort` orders
      // by UTF-16 code unit, and the Python version sorts the same way. Sorting
      // once per column, not once per cell, keeps thousands of postcodes cheap.
      const labels = Object.keys(spec.counts).sort();
      const weights = labels.map((label) => spec.counts[label]);
      return (row) => labels[pick(weights, draw(seed, field, row))];
    }
    case 'histogram': {
      const { edges, counts } = spec;
      if (edges.length !== counts.length + 1) {
        throw new RangeError('a histogram needs one more edge than it has buckets');
      }
      return (row) => {
        const bucket = pick(counts, draw(seed, field, row));
        const [low, high] = [edges[bucket], edges[bucket + 1]];
        // A second, independent draw places the value inside its bucket. Within
        // a bucket the shape is unknown, so uniform is the only honest choice.
        return low + (draw(seed, `${field}${UNIT}within`, row) % Math.max(high - low, 1));
      };
    }
    default:
      throw new RangeError(`unknown distribution type ${JSON.stringify(spec.type)}`);
  }
}

/**
 * Return `count` rows, each column sampled from its own observed distribution.
 *
 * `distributions` maps a field name to either
 * {type: 'categorical', counts: {label: observed count}} or
 * {type: 'histogram', edges, counts}, where the edges are integers, the
 * buckets are half-open and `edges` holds one more value than `counts`.
 */
export function sampleRows(distributions, count, seed) {
  const samplers = Object.entries(distributions).map(([field, spec]) => [field, sampler(spec, seed, field)]);
  const rows = [];
  for (let row = 0; row < count; row += 1) {
    rows.push(Object.fromEntries(samplers.map(([field, sample]) => [field, sample(row)])));
  }
  return rows;
}

"""
Write the text fields of a test data set by asking a general-purpose model.

Rung N3. This is the option people reach for first, and the only rung of this
entry that writes free text rather than picking from values the caller listed.

It also gives up the property the lower rungs were built on. There is no seed
here: the request carries a prompt and a temperature, and nothing that would make
a second call repeat the first. The data set has to be generated once and then
stored, like a fixture, not rebuilt on demand.

And a prompt is only a request: nothing in it holds the model to the keys you
asked for, the number of rows or the uniqueness of an identifier. Look at how
much of this file is checking rather than asking. That plumbing is the real cost
of the rung, and it is the part the tests can cover, because the model itself is
not testable.
"""

from __future__ import annotations

import json

# A cap chosen for this snippet, not a measured limit: the whole batch has to
# fit in one answer, and a model's output is capped in tokens.
MAX_ROWS = 50

# The provider named here is an example, not a recommendation: the reasoning
# holds for any general-purpose model API, and the client is swappable. Pass
# any object with a `complete(prompt=..., temperature=...)` method.
MODEL = "gpt-4.1-mini"  # an example id: check the parameters your model accepts


class GenerationUnavailable(Exception):
    """The provider could not be reached, or never returned a usable data set."""


class ProviderClient:
    """The one call this snippet makes, on top of the provider's SDK."""

    def __init__(self, sdk=None, model: str = MODEL):
        if sdk is None:  # pragma: no cover - needs a key and a network
            from openai import OpenAI

            sdk = OpenAI()
        self.sdk, self.model = sdk, model

    def complete(self, *, prompt: str, temperature: float) -> str:
        response = self.sdk.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": prompt}],
            temperature=temperature,
        )
        return response.choices[0].message.content


def build_prompt(fields: list[str], count: int, unique_field: str | None) -> str:
    """The instructions, kept next to the checks that verify they were followed."""
    lines = [
        f"Write {count} rows of test data for a fictional application.",
        "Each row is a JSON object with exactly these keys, and string values:",
        ", ".join(fields) + ".",
        "Invent every value: it must match no real person, company or address.",
        f"Answer with JSON only: a list of {count} objects, and nothing else.",
    ]
    if unique_field:
        lines.insert(3, f"Every value of `{unique_field}` must differ from the others.")
    return "\n".join(lines)


def check(rows, fields: list[str], count: int, unique_field: str | None) -> None:
    """
    Verify what the model was asked for: the row count, exactly the keys,
    non-empty strings, and the uniqueness requested.
    """
    if not isinstance(rows, list) or len(rows) != count:
        raise ValueError(f"expected a list of {count} rows")
    for row in rows:
        if not isinstance(row, dict) or set(row) != set(fields):
            raise ValueError(f"a row does not carry exactly the keys {fields}")
        if not all(isinstance(value, str) and value.strip() for value in row.values()):
            raise ValueError("a value is empty, or is not a string")
    if unique_field:
        values = [row[unique_field] for row in rows]
        if len(set(values)) != len(values):
            raise ValueError(f"the model repeated a value of {unique_field!r}")


def write_rows(
    fields: list[str],
    count: int,
    *,
    unique_field: str | None = None,
    client=None,
    attempts: int = 3,
    temperature: float = 1.0,
) -> list[dict]:
    """
    Return `count` rows of invented text, or raise rather than return junk.

    `client` is injected so this function can be tested without a network call.
    In production it defaults to a real provider client.

    The temperature is high on purpose: varied prose is the only reason to be
    on this rung at all. It is also why the answer has to be checked, and why
    the same call twice gives two different data sets.
    """
    # A model charges by the token. Refusing a batch no answer can satisfy
    # before calling is not an optimisation, it is a cost control.
    if not isinstance(count, int) or not 0 < count <= MAX_ROWS:
        raise ValueError(f"ask for a whole number of rows, between one and {MAX_ROWS}")
    if unique_field is not None and unique_field not in fields:
        raise ValueError(f"the unique field {unique_field!r} is not one of the fields")
    client = client or ProviderClient()

    prompt = build_prompt(fields, count, unique_field)
    last_error: Exception | None = None
    for _ in range(attempts):
        try:
            answer = client.complete(prompt=prompt, temperature=temperature)
            if not isinstance(answer, str):  # a refusal comes back with no content
                raise ValueError("the model returned no text")
            rows = json.loads(answer)
            check(rows, fields, count, unique_field)
            return rows
        except Exception as error:  # noqa: BLE001 - a bad answer is retried like a failure
            last_error = error
    raise GenerationUnavailable(str(last_error))

/**
 * Write the text fields of a test data set by asking a general-purpose model.
 *
 * Rung N3. This is the option people reach for first, and the only rung of
 * this entry that writes free text rather than picking from values the caller
 * listed.
 *
 * It also gives up the property the lower rungs were built on. There is no
 * seed here: the request carries a prompt and a temperature, and nothing that
 * would make a second call repeat the first. The data set has to be generated
 * once and then stored, like a fixture, not rebuilt on demand.
 *
 * And a prompt is only a request: nothing in it holds the model to the keys
 * you asked for, the number of rows or the uniqueness of an identifier. Look
 * at how much of this file is checking rather than asking. That plumbing is the
 * real cost of the rung, and it is the part the tests can cover, because the
 * model itself is not testable.
 */

// A cap chosen for this snippet, not a measured limit: the whole batch has to
// fit in one answer, and a model's output is capped in tokens.
export const MAX_ROWS = 50;

// The provider named here is an example, not a recommendation: the reasoning
// holds for any general-purpose model API, and the client is swappable. Pass
// any object with a `complete({ prompt, temperature })` method.
export const MODEL = 'gpt-4.1-mini'; // an example id: check the parameters your model accepts

export class GenerationUnavailable extends Error {}

export async function providerClient(sdk, model = MODEL) {
  if (!sdk) {
    const { OpenAI } = await import('openai');
    sdk = new OpenAI();
  }
  return {
    async complete({ prompt, temperature }) {
      const response = await sdk.chat.completions.create({
        model,
        messages: [{ role: 'user', content: prompt }],
        temperature,
      });
      return response.choices[0].message.content;
    },
  };
}

/** The instructions, kept next to the checks that verify they were followed. */
export function buildPrompt(fields, count, uniqueField) {
  const lines = [
    `Write ${count} rows of test data for a fictional application.`,
    'Each row is a JSON object with exactly these keys, and string values:',
    `${fields.join(', ')}.`,
    'Invent every value: it must match no real person, company or address.',
    `Answer with JSON only: a list of ${count} objects, and nothing else.`,
  ];
  if (uniqueField) {
    lines.splice(3, 0, `Every value of \`${uniqueField}\` must differ from the others.`);
  }
  return lines.join('\n');
}

/**
 * Verify what the model was asked for: the row count, exactly the keys,
 * non-empty strings, and the uniqueness requested.
 */
export function check(rows, fields, count, uniqueField) {
  if (!Array.isArray(rows) || rows.length !== count) {
    throw new TypeError(`expected a list of ${count} rows`);
  }
  for (const row of rows) {
    const keys = row && typeof row === 'object' ? Object.keys(row) : [];
    if (keys.length !== fields.length || !fields.every((field) => keys.includes(field))) {
      throw new TypeError(`a row does not carry exactly the keys ${fields.join(', ')}`);
    }
    if (!Object.values(row).every((value) => typeof value === 'string' && value.trim())) {
      throw new TypeError('a value is empty, or is not a string');
    }
  }
  if (uniqueField) {
    const values = rows.map((row) => row[uniqueField]);
    if (new Set(values).size !== values.length) {
      throw new TypeError(`the model repeated a value of ${uniqueField}`);
    }
  }
}

/**
 * Return `count` rows of invented text, or throw rather than return junk.
 *
 * The temperature is high on purpose: varied prose is the only reason to be on
 * this rung at all. It is also why the answer has to be checked, and why the
 * same call twice gives two different data sets.
 *
 * @param {string[]} fields
 * @param {number} count
 * @param {object} options
 * @param {string} [options.uniqueField] field whose values must not repeat
 * @param {{complete: Function}} [options.client] injected so this can be
 *   tested without a network call; defaults to a real provider client
 * @param {number} [options.attempts]
 * @param {number} [options.temperature]
 */
export async function writeRows(fields, count, options = {}) {
  const { uniqueField, attempts = 3, temperature = 1 } = options;
  let { client } = options;

  // A model charges by the token. Refusing a batch no answer can satisfy
  // before calling is not an optimisation, it is a cost control.
  if (!Number.isInteger(count) || !(count > 0 && count <= MAX_ROWS)) {
    throw new RangeError(`ask for a whole number of rows, between one and ${MAX_ROWS}`);
  }
  if (uniqueField !== undefined && !fields.includes(uniqueField)) {
    throw new RangeError(`the unique field ${uniqueField} is not one of the fields`);
  }
  // Needs a key and a network, so it is never reached in the tests.
  client ??= await providerClient();

  const prompt = buildPrompt(fields, count, uniqueField);
  let lastError;
  for (let attempt = 0; attempt < attempts; attempt += 1) {
    try {
      const answer = await client.complete({ prompt, temperature });
      // A refusal comes back with no content.
      if (typeof answer !== 'string') throw new TypeError('the model returned no text');
      const rows = JSON.parse(answer);
      check(rows, fields, count, uniqueField);
      return rows;
    } catch (error) {
      // A bad answer is retried exactly like a provider failure.
      lastError = error;
    }
  }
  throw new GenerationUnavailable(String(lastError));
}