/**
 * Type-rule patches for numbl JIT builtins, applied to the in-process builtin
 * registry before lowering a script.
 *
 * numbl's JIT declines a handful of constructs its C/JS backends do not
 * implement yet (tensor `&`/`|`/`~`, two-arg `max`/`min` on tensors, matrix
 * `rand`/`randn`), and its comparison builtins deliberately return a
 * shape-unknown logical tensor because the lowerer has no static broadcast
 * helper. This backend executes the IR on WebGPU — numbl's own emitters never
 * run here — so all this sandbox needs from those builtins is a precise
 * *type*: everything else is supplied by the WGSL kernels.
 *
 * `registerBuiltin` overwrites by name (by design, for HMR), which makes it
 * the sanctioned way to swap a builtin's rules. The patches only ever widen:
 * every case the original transfer accepts is delegated to it unchanged, so
 * scalar behavior (including exact-value folding) is untouched.
 *
 * IMPORTANT: the patch is process-global to this module instance of numbl.
 * The CPU comparison runner executes scripts through numbl's normal engine in
 * a separate Web Worker, whose module instances are its own — it never sees
 * these patches.
 */
import {
  getBuiltin,
  registerBuiltin,
  type Builtin,
} from 'numbl-src/numbl-core/jit/builtins/index.ts';
import {
  isNumeric,
  isMultiElement,
  type NumericType,
  type Type,
} from 'numbl-src/numbl-core/jit/lowering/types.ts';
import { UnsupportedOnGpu } from './errors.ts';

const isRealTensor = (t: Type): t is NumericType =>
  isNumeric(t) && !t.isComplex && isMultiElement(t);
const isRealScalarish = (t: Type): t is NumericType =>
  isNumeric(t) && !t.isComplex && !isMultiElement(t);

const logicalTensor = (shape: number[]): NumericType => ({
  kind: 'Numeric',
  elem: 'logical',
  isComplex: false,
  dims: shape.map((n) => ({ kind: 'exact', value: n })),
  shape: [...shape],
  sign: 'nonneg',
});

const doubleTensor = (shape: number[], sign: NumericType['sign']): NumericType => ({
  kind: 'Numeric',
  elem: 'double',
  isComplex: false,
  dims: shape.map((n) => ({ kind: 'exact', value: n })),
  shape: [...shape],
  sign,
});

const sameShape = (a: number[], b: number[]): boolean =>
  a.length === b.length && a.every((d, i) => d === b[i]);

/**
 * The elementwise result shape for two operands, under the same rule the
 * kernels implement: operands are either scalars or tensors of one common
 * shape. Implicit expansion (`n x 1` with `1 x m`) is rejected here, matching
 * the emitter's refusal to broadcast.
 */
function elementwiseShape(name: string, argTypes: Type[]): number[] | null {
  let shape: number[] | null = null;
  for (const t of argTypes) {
    if (!isNumeric(t) || t.isComplex) {
      throw new UnsupportedOnGpu(`'${name}': operands must be real numeric`);
    }
    if (!isMultiElement(t)) continue;
    if (!t.shape) {
      throw new UnsupportedOnGpu(
        `'${name}': operand shape is not known at compile time`,
      );
    }
    if (shape && !sameShape(shape, t.shape)) {
      throw new UnsupportedOnGpu(
        `'${name}': operands are ${shape.join('x')} and ${t.shape.join('x')}; ` +
          `the GPU kernels do not broadcast — expand explicitly`,
      );
    }
    shape = t.shape;
  }
  return shape;
}

/** Wrap `transfer` so tensor operands get an elementwise result type instead
 *  of the original's decline (or, for comparisons, its shape-unknown type). */
function widenElementwise(
  name: string,
  makeResult: (shape: number[]) => NumericType,
  /** Arity the elementwise form requires. `max(x)` with one tensor arg is the
   *  reduction form and must stay with the original transfer. */
  arity?: number,
): void {
  const orig = getBuiltin(name);
  if (!orig) throw new Error(`patch: no builtin named '${name}'`);
  const origTransfer = orig.transfer.bind(orig);
  registerBuiltin({
    ...orig,
    transfer(argTypes: Type[], nargout: number): Type[] {
      const anyTensor = argTypes.some((t) => isNumeric(t) && isMultiElement(t));
      if (!anyTensor || (arity !== undefined && argTypes.length !== arity)) {
        return origTransfer(argTypes, nargout);
      }
      if (nargout > 1) {
        throw new UnsupportedOnGpu(`'${name}' returns one value`);
      }
      const shape = elementwiseShape(name, argTypes);
      // anyTensor guaranteed a tensor operand, so shape is non-null.
      return [makeResult(shape!)];
    },
  } as Builtin);
}

/** Comparisons: `lt` et al. accept tensors already but return unknown dims;
 *  replace the tensor case with the precise shape. */
const COMPARISONS = ['lt', 'le', 'gt', 'ge', 'eq', 'ne'];

/** Eager elementwise `&` / `|` / `xor` / `~`. (`&&`/`||` stay scalar-only —
 *  that is MATLAB semantics, not a backend gap.) */
const LOGICALS = ['and', 'or', 'xor', 'not'];

/** Two-arg elementwise forms of `max`/`min`. The one-arg reduction form is
 *  typed by the original transfer. */
const MINMAX = ['max', 'min'];

/** `rand(n)` / `rand(m, n)` / `randn(...)`: numbl's JS-JIT supports only the
 *  scalar form; here any exact-sized matrix form is a fill kernel. */
function widenRand(name: string, sign: NumericType['sign']): void {
  // `randn` has no JIT builtin at all in numbl yet; register it whole.
  const orig = getBuiltin(name);
  const origTransfer = orig
    ? orig.transfer.bind(orig)
    : (): Type[] => [{
        kind: 'Numeric', elem: 'double', isComplex: false,
        dims: [{ kind: 'exact', value: 1 }, { kind: 'exact', value: 1 }],
        shape: [1, 1], sign,
      } satisfies NumericType];
  registerBuiltin({
    ...(orig ?? { name }),
    name,
    transfer(argTypes: Type[], nargout: number): Type[] {
      if (argTypes.length === 0) return origTransfer(argTypes, nargout);
      if (nargout > 1) throw new UnsupportedOnGpu(`'${name}' returns one value`);
      if (argTypes.length > 2) {
        throw new UnsupportedOnGpu(
          `'${name}': only the 2-D forms ${name}(n) / ${name}(m, n) are supported`,
        );
      }
      const dims = argTypes.map((t) => {
        if (!isRealScalarish(t) || typeof t.exact !== 'number') {
          throw new UnsupportedOnGpu(
            `'${name}': array sizes must be known at compile time (assign the ` +
              `size from a literal, e.g. n = 1024)`,
          );
        }
        if (!Number.isInteger(t.exact) || t.exact < 0) {
          throw new UnsupportedOnGpu(`'${name}': sizes must be whole numbers`);
        }
        return t.exact;
      });
      const shape = dims.length === 1 ? [dims[0], dims[0]] : dims;
      return [doubleTensor(shape, sign)];
    },
  } as Builtin);
}

let applied = false;

/** Apply every patch, once per module instance. */
export function applyBuiltinPatches(): void {
  if (applied) return;
  applied = true;
  for (const name of COMPARISONS) widenElementwise(name, logicalTensor);
  for (const name of LOGICALS) widenElementwise(name, logicalTensor);
  for (const name of MINMAX) {
    widenElementwise(name, (shape) => doubleTensor(shape, 'unknown'), 2);
  }
  widenRand('rand', 'positive');
  widenRand('randn', 'unknown');
}