/**
 * Watermark PRFs (pseudo-random functions).
 *
 * Constants follow the reference implementations:
 * - HASH_KEY = 15485863 (the 1,000,000th prime) is the multiplicative salt
 *   used by both jwkirchenbauer/lm-watermarking and bohanhou14/SemStamp.
 */

import { hashCombine, mix64, toUnitFloat } from './rng';

export const HASH_KEY = 15485863n;
const MOD64 = (1n << 64n) - 1n;

/**
 * Kirchenbauer "lefthash" (simple_1) seed: salt * prev_token_id, mod 2^64-1.
 * Matches the reference `additive_prf` with context_width=1.
 */
export function kirchenbauerSeed(prevTokenId: number, hashKey: bigint = HASH_KEY): bigint {
  return (hashKey * BigInt(prevTokenId)) % MOD64;
}

/**
 * Green-list membership test. Instead of permuting the whole vocabulary
 * (torch.randperm in the reference), we use a per-token PRF judgment:
 * token is green iff U(seed, key, tokenId) < gamma.
 * The green-list size is then Binomial(|V|, gamma) instead of exactly
 * gamma*|V|; negligible for the z-test (documented in README).
 *
 * `key` is the user-chosen secret. It is mixed in here rather than into the
 * seed formula so the displayed derivation stays literally
 * `15485863 * prev_token mod 2^64-1`, as in the reference implementation.
 */
export function isGreenToken(
  seed: bigint,
  tokenId: number,
  gamma: number,
  key: bigint = 0n,
): boolean {
  return greenUniform(seed, tokenId, key) < gamma;
}

/** The underlying uniform used by isGreenToken (exposed for UI/inspector). */
export function greenUniform(seed: bigint, tokenId: number, key: bigint = 0n): number {
  return toUnitFloat(hashCombine(seed, key, BigInt(tokenId)));
}

/**
 * TextSeal-style PRF: R_v = U(key, context window, candidate token v) in [0,1).
 * The watermark context window is the previous k tokens (paper default k=3).
 *
 * NOTE: the official facebookresearch/textseal repo derives R via a
 * cryptographic hash (appendix A.1). We use splitmix64-based mixing for
 * per-token speed in the browser; this is an educational deviation that
 * preserves the statistical construction (uniform R per (context, v, key)).
 */
export function textsealR(
  key: bigint,
  contextTokenIds: number[],
  candidateTokenId: number,
): number {
  const parts: bigint[] = [key];
  for (const t of contextTokenIds) parts.push(BigInt(t));
  parts.push(0x746578747365616cn); // "textseal" domain separator
  parts.push(BigInt(candidateTokenId));
  let u = toUnitFloat(hashCombine(...parts));
  // Guard against exact 0/1 for downstream log computations.
  if (u <= 0) u = Number.MIN_VALUE;
  if (u >= 1) u = 1 - Number.EPSILON;
  return u;
}

/** Derive a stable 64-bit key from a human-readable secret string. */
export function keyFromString(secret: string): bigint {
  let h = 0xcbf29ce484222325n; // FNV offset basis
  for (let i = 0; i < secret.length; i++) {
    h = mix64((h ^ BigInt(secret.charCodeAt(i))) & MOD64);
  }
  return h;
}

/** k-SemStamp seed: prev sentence's cluster id * prime salt, mixed with the key. */
export function ksemstampSeed(
  clusterId: number,
  key: bigint = 0n,
  hashKey: bigint = HASH_KEY,
): bigint {
  const base = (hashKey * BigInt(clusterId + 1)) % MOD64; // +1 so cluster 0 doesn't zero out
  return key === 0n ? base : hashCombine(base, key);
}

/** Derive a 64-bit key from a numeric seed entered by the user. */
export function keyFromSeed(seed: number): bigint {
  return mix64(BigInt(Math.trunc(seed)) & MOD64) | 1n; // never zero
}