/**
 * The host-provided operations a model .m can call — the only places where a
 * point reads anything but itself.
 *
 * Everything else a model does is element-wise, so these five are where the
 * physics that couples neighbours (and couples the two grids) lives:
 *
 *   dxx(u)        second difference along the string        (string -> string)
 *   dxxxx(u)      fourth difference along the string        (string -> string)
 *   lapw(p, w)    wall-masked 7-point Laplacian in the air  (air, air -> air)
 *   spread(a)     sample a string field at each air point's x  (string -> air)
 *   bridge(u)     du/dx at the string's bridge end, broadcast  (string -> air)
 *
 * numbl needs only their *type rule* in order to lower a call site. It gets
 * that from a `.mtoc2.js` workspace file — numbl's sanctioned extension point
 * for a JS-defined builtin (see `mtoc2UserFunctionsByName` in numbl's
 * LoweringContext). The file is evaluated in a bare CommonJS sandbox with no
 * imports available, so `transfer` builds numbl `Type` objects as plain
 * literals, and the grid sizes are baked in by the generator below (a grid
 * change recompiles anyway).
 *
 * The `emit`/`cBody` exports exist only because the loader's contract requires
 * them; we never emit C. The actual implementation is supplied by the WGSL
 * backend (src/mgpu/ops.ts).
 */

export interface GridSizes {
  /** Air grid points, nx*ny*nz. Air fields are npts x 1 column vectors. */
  npts: number;
  /** String nodes. String fields are ns x 1 column vectors. */
  ns: number;
}

/** One external operation's shape contract. */
export interface ExternalOpSpec {
  name: string;
  /** Element counts of the arguments, in order. */
  args: ('air' | 'string')[];
  out: 'air' | 'string';
}

export const EXTERNAL_OP_SPECS: ExternalOpSpec[] = [
  { name: 'dxx', args: ['string'], out: 'string' },
  { name: 'dxxxx', args: ['string'], out: 'string' },
  { name: 'lapw', args: ['air', 'air'], out: 'air' },
  { name: 'spread', args: ['string'], out: 'air' },
  { name: 'bridge', args: ['string'], out: 'air' },
];

/** Names the WGSL backend must implement as dispatches rather than
 *  element-wise kernels, with the argument count of each. */
export const EXTERNAL_OPS = new Map(EXTERNAL_OP_SPECS.map((s) => [s.name, s]));

const numericType = (rows: number): string =>
  `{ kind: "Numeric", elem: "double", isComplex: false, ` +
  `dims: [{ kind: "exact", value: ${rows} }, { kind: "exact", value: 1 }], ` +
  `shape: [${rows}, 1], sign: "unknown" }`;

/** Source for one op's `.mtoc2.js`: fields in, field out, shapes checked. */
function opSource(spec: ExternalOpSpec, g: GridSizes): string {
  const size = (k: 'air' | 'string'): number => (k === 'air' ? g.npts : g.ns);
  const checks = spec.args
    .map((kind, i) => {
      const n = size(kind);
      const what = kind === 'air' ? 'an air field' : 'a string field';
      return `
  var a${i} = argTypes[${i}];
  if (!a${i} || a${i}.kind !== "Numeric" || a${i}.isComplex) {
    throw new Error("${spec.name}: argument ${i + 1} must be a real numeric array");
  }
  var s${i} = a${i}.shape;
  if (!s${i} || s${i}.length !== 2 || s${i}[0] !== ${n} || s${i}[1] !== 1) {
    throw new Error(
      "${spec.name}: argument ${i + 1} must be ${what} (${n}x1), not " +
        (s${i} ? s${i}.join("x") : "unknown shape")
    );
  }`;
    })
    .join('\n');
  return `
exports.name = ${JSON.stringify(spec.name)};

exports.transfer = function (argTypes, nargout) {
  if (argTypes.length !== ${spec.args.length}) {
    throw new Error("${spec.name} takes ${spec.args.length} argument(s), got " + argTypes.length);
  }
  if (nargout > 1) {
    throw new Error("${spec.name} returns one value, but " + nargout + " were requested");
  }
${checks}
  return [${numericType(size(spec.out))}];
};

// Never called: this project executes the IR on WebGPU and emits no C.
exports.emit = function () {
  throw new Error("${spec.name}: no C backend (this op runs on WebGPU)");
};
exports.cBody = function () {
  return "";
};
`;
}

/** Workspace files that make the ops resolvable during lowering. */
export function externalOpFiles(g: GridSizes): { name: string; source: string }[] {
  return EXTERNAL_OP_SPECS.map((spec) => ({
    name: `${spec.name}.mtoc2.js`,
    source: opSource(spec, g),
  }));
}