// HEAL is a mask and a fill now, not a patch pasted under each spot. Four
// things can go wrong and none of them throws: the circles of one stroke can
// stay separate shapes (which is what a row of round grey blobs on the photo
// is), the fill can fail to carry the light of the photo around it (a repair
// that reads as a mark of its own), the search can still borrow a patch from
// somewhere (the card three radii away, which is the brush the user complained
// about), or the composite can touch pixels it was never given (a smear across
// the whole frame). So the mask, the fill and the untouched ground are all
// measured here, on synthetic photos whose answer is known, with CanvasKit
// running the same two passes exportEngine runs.
//
// The fill now has a second, better author where one can be had — MI-GAN, over
// the mask (engine/inpaint.ts) — and this check deliberately measures the other
// one: the blur that stands in when the model cannot be had, which is the case
// on a machine without the runtime and the case this file can run offline. What
// is asserted about the model is only that the pass is still wired in, and with
// which polarity.
//
//   node scripts/heal-check.mjs
import assert from 'node:assert/strict';
import { mkdtempSync, readFileSync, writeFileSync } from 'node:fs';
import { tmpdir } from 'node:os';
import { join } from 'node:path';
import { fileURLToPath, pathToFileURL } from 'node:url';
import ts from 'typescript';
import CanvasKitInit from 'canvaskit-wasm/bin/full/canvaskit.js';

const transpile = (path) =>
  ts.transpileModule(readFileSync(new URL(path, import.meta.url), 'utf8'), {
    compilerOptions: { module: ts.ModuleKind.ESNext, target: ts.ScriptTarget.ES2022 },
  }).outputText;

const dir = mkdtempSync(join(tmpdir(), 'heal-check-'));
writeFileSync(join(dir, 'heal.mjs'), transpile('../shared/utils/heal.ts'));
const { readHeal, HEAL_MASK_STOPS, HEAL_FILL_BLUR, HEAL_FILL_MAX_LEVELS, HEAL_CUT_SKSL, HEAL_FILL_SKSL } = await import(
  pathToFileURL(join(dir, 'heal.mjs')).href
);

// The other half of the check: the passes below are these two shaders, so the
// engine has to be compiling them — and it must not still be building the
// per-spot shader this replaced.
const engine = readFileSync(new URL('../src/engine/exportEngine.ts', import.meta.url), 'utf8');
assert.match(engine, /healCutEffect = Skia\.RuntimeEffect\.Make\(HEAL_CUT_SKSL\)/, 'exportEngine does not compile HEAL_CUT_SKSL');
assert.match(engine, /healFillEffect = Skia\.RuntimeEffect\.Make\(HEAL_FILL_SKSL\)/, 'exportEngine does not compile HEAL_FILL_SKSL');
assert.doesNotMatch(engine, /healSkSL\(/, 'the per-spot heal shader is still in exportEngine');
assert.doesNotMatch(engine, /healUniforms\(/, 'the per-spot uniform block is still in exportEngine');
// ...and the fill's level count is the one this check builds, not a constant.
assert.match(engine, /HEAL_FILL_MAX_LEVELS/, 'exportEngine does not cap the fill levels at HEAL_FILL_MAX_LEVELS');
assert.match(engine, /1 \+ Math\.ceil\(Math\.log2\(2 \+ span \/ 2 \/ widest\)\)/, 'exportEngine no longer scales the fill levels with the depth of the hole');
// ...and the fill is built on the brush's box, read back at the frame's origin.
assert.match(engine, /imageShaderChildAt\(fillImage, boxLeft, boxTop\)/, 'the fill child is read at its own origin instead of the frame\'s');
// ...and the box the fill is built on is drawn at the surface's own corner, which
// is where that child reads it. Drawn at the frame's coordinates instead, the two
// disagree by the box's corner, and the repair comes back a patch of the frame
// from up and to the left of the spot (the case at the foot of this file).
assert.match(engine, /drawImageRect\(cutImage, box, onto, /, 'the fill is drawn at the box\'s frame coordinates, not at the surface\'s own corner');

// The model half is wired into the same pass: the masked crop is asked for, the
// answer is written into the fill surface over the box, and the mask it is
// handed is built from the brush's own circles. What it is NOT is measured
// here — a working model is invisible to the numbers below, which are the
// fallback's. The polarity is MI-GAN's, 255 KEEP and 0 HOLE, and it is the one
// thing that, got backwards, repaints the frame and leaves the hole (measured in
// engine/inpaint.ts; the probe is gone).
const inpaint = readFileSync(new URL('../src/engine/inpaint.ts', import.meta.url), 'utf8');
assert.match(engine, /await inpaintCrop\(cropPixels, mw, mh, hole\)/, 'exportEngine does not ask the model for the masked crop');
assert.match(engine, /patchSurface\.getCanvas\(\)\.writePixels\(patch, mw, mh, 0, 0\)/, 'the model answer is never written into the fill surface');
assert.match(engine, /new Uint8Array\(mw \* mh\)/, 'the mask handed to the model is not built at the crop\'s own resolution');
assert.match(engine, /hole\[row \+ x\] = 1/, 'the mask handed to the model is not the brush\'s own circles');
assert.match(inpaint, /fill\(255\)/, 'the model mask no longer starts at KEEP');
assert.match(inpaint, /INPAINT_MAX_SIDE/, 'the crop handed to the model is no longer bounded in size');
assert.match(inpaint, /inpaintCropSize/, 'nothing tells the renderer the resolution the model will run at');

// --- the stored spots --------------------------------------------------------
// A recipe written by the older brush carries sx/sy. It still loads: those two
// numbers are no longer read, and nothing else about the spot moves.
assert.deepEqual(readHeal([{ x: 0.5, y: 0.5, r: 0.02, sx: 0.9, sy: 0.1 }]), [{ x: 0.5, y: 0.5, r: 0.02 }]);
// Out of frame, non-numeric, or a radius of nothing: a spot the brush and the
// renderer would disagree about is dropped either way.
assert.deepEqual(readHeal([{ x: 2, y: -1, r: 0.05 }, { x: 'a', y: 0, r: 0.05 }, { x: 0.5, y: 0.5, r: 0 }]), [
  { x: 1, y: 0, r: 0.05 },
  { x: 0, y: 0, r: 0.05 },
]);

const CanvasKit = await CanvasKitInit({
  locateFile: () => fileURLToPath(new URL('../node_modules/canvaskit-wasm/bin/full/canvaskit.wasm', import.meta.url)),
});

// The frame every pass below is built on. One size for most of the check; the
// last case needs a bigger one, which is the whole point of it (see there), so
// these two are the frame and the functions here read them.
let W = 160;
let H = 160;
const GREY = 118;
const DUST = 38;
const R = 0.05; // the brush: an 8px radius on this frame
// A stroke of nine spots, one every 0.6 of a radius, running down the frame.
const spots = Array.from({ length: 9 }, (_, i) => ({ x: 0.5, y: (48 + i * R * 0.6 * W) / H, r: R }));
const cx = 0.5 * W;
const cy = 48 + 4 * R * 0.6 * W;
// Three radii to the right of the stroke: where the old search looked.
const CARD_X = Math.round(cx + 3.2 * R * W);

const shadeFlat = (x, y) => {
  if (Math.abs(x - CARD_X) < 6 && Math.abs(y - cy) < 6) return [220, 30, 30];
  if (Math.hypot(x - cx, y - cy) < 6) return [DUST, DUST, DUST];
  return [GREY, GREY, GREY];
};
const shadeRamp = (x) => {
  const v = 40 + (120 * x) / (W - 1);
  return [v, v, v];
};

function bytesOf(shade) {
  const rgba = new Uint8Array(W * H * 4);
  for (let y = 0; y < H; y++) {
    for (let x = 0; x < W; x++) {
      const [r, g, b] = shade(x, y);
      const i = (y * W + x) * 4;
      rgba[i] = r;
      rgba[i + 1] = g;
      rgba[i + 2] = b;
      rgba[i + 3] = 255;
    }
  }
  return rgba;
}
const imageOf = (rgba) =>
  CanvasKit.MakeImage(
    {
      width: W,
      height: H,
      colorType: CanvasKit.ColorType.RGBA_8888,
      alphaType: CanvasKit.AlphaType.Unpremul,
      colorSpace: CanvasKit.ColorSpace.SRGB,
    },
    rgba,
    W * 4
  );
const surface = (w, h) => {
  const s = CanvasKit.MakeSurface(w, h);
  s.getCanvas().clear(CanvasKit.TRANSPARENT);
  return s;
};
const pixelsOf = (image) => {
  const s = surface(image.width(), image.height());
  s.getCanvas().drawImage(image, 0, 0);
  return s.getCanvas().readPixels(0, 0, {
    width: image.width(),
    height: image.height(),
    colorType: CanvasKit.ColorType.RGBA_8888,
    alphaType: CanvasKit.AlphaType.Unpremul,
    colorSpace: CanvasKit.ColorSpace.SRGB,
  });
};
const asChild = (image) =>
  image.makeShaderOptions(CanvasKit.TileMode.Clamp, CanvasKit.TileMode.Clamp, CanvasKit.FilterMode.Linear, CanvasKit.MipmapMode.None);
// The fill's own child, read in the frame's coordinates: the fill is built on
// the brush's box (below), and a whole-pixel shift is what the engine does.
const asChildAt = (image, x, y) =>
  image.makeShaderOptions(
    CanvasKit.TileMode.Clamp,
    CanvasKit.TileMode.Clamp,
    CanvasKit.FilterMode.Linear,
    CanvasKit.MipmapMode.None,
    [1, 0, x, 0, 1, y, 0, 0, 1]
  );

// The brush's circles as ONE mask: a radial gradient per spot, solid out to
// HEAL_FEATHER of its radius and clear at its edge, each laid over the last.
function maskOf(spotsIn) {
  const surf = surface(W, H);
  const canvas = surf.getCanvas();
  for (const s of spotsIn) {
    const rad = s.r * W;
    const shader = CanvasKit.Shader.MakeRadialGradient(
      [s.x * W, s.y * H],
      rad,
      new Float32Array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0]),
      HEAL_MASK_STOPS,
      CanvasKit.TileMode.Clamp
    );
    const paint = new CanvasKit.Paint();
    paint.setAntiAlias(true);
    paint.setShader(shader);
    canvas.drawCircle(s.x * W, s.y * H, rad, paint);
  }
  return surf.makeImageSnapshot();
}

// exportEngine's 6e, pass for pass: cut, the blurred fill, and the fill laid
// back under the brush.
function healPass(rgba, spotsIn) {
  const src = imageOf(rgba);
  const mask = maskOf(spotsIn);
  const cutEffect = CanvasKit.RuntimeEffect.Make(HEAL_CUT_SKSL);
  assert.ok(cutEffect, 'HEAL_CUT_SKSL does not compile');
  const cut = surface(W, H);
  const cutPaint = new CanvasKit.Paint();
  cutPaint.setShader(cutEffect.makeShaderWithChildren([], [asChild(src), asChild(mask)]));
  cut.getCanvas().drawRect(CanvasKit.XYWHRect(0, 0, W, H), cutPaint);

  // The fill: the cut frame blurred at a few widths, widest first, each finer
  // one drawn over it, the untouched cut on top. A blur is a weighted mean of
  // the premultiplied pixels around each pixel, and the hole weighs nothing, so
  // the hole comes back as the light of its rim. Widest first means the widest
  // blur only has a say where every finer one has nothing — the middle of a hole
  // deeper than a brush — and the untouched cut wins wherever the photo is still
  // there, so the rim is the rim exactly.
  //
  // Blurs and not the pyramid of halvings the engine used to build: this Skia
  // gets the premultiply wrong on a snapshot drawn with a SCALE — its rgb comes
  // back with the alpha multiplied in a second time — and a pyramid is made of
  // exactly those halvings, so a partly transparent rim came back a wash of
  // grey. Every draw here is 1:1, which filters correctly.
  const cutImage = cut.makeImageSnapshot();
  // The count and the widths are exportEngine's own, to the letter: a check
  // that picks its own levels is checking a renderer nobody ships. So is the
  // surface the fill is built on — the brush's own box, padded by three sigma of
  // the widest level, rather than the frame.
  const widest = Math.max(...spotsIn.map((s) => s.r * W));
  const xs = spotsIn.map((s) => s.x * W);
  const ys = spotsIn.map((s) => s.y * H);
  const span = Math.hypot(Math.max(...xs) - Math.min(...xs), Math.max(...ys) - Math.min(...ys));
  const levelCount = Math.min(
    HEAL_FILL_MAX_LEVELS,
    Math.max(2, 1 + Math.ceil(Math.log2(2 + span / 2 / widest)))
  );
  const pad = 3 * HEAL_FILL_BLUR * widest * 2 ** (levelCount - 1);
  const boxLeft = Math.max(0, Math.floor(Math.min(...xs) - widest - pad));
  const boxTop = Math.max(0, Math.floor(Math.min(...ys) - widest - pad));
  const boxRight = Math.min(W, Math.ceil(Math.max(...xs) + widest + pad));
  const boxBottom = Math.min(H, Math.ceil(Math.max(...ys) + widest + pad));
  const box = CanvasKit.XYWHRect(boxLeft, boxTop, boxRight - boxLeft, boxBottom - boxTop);
  // ...and the box's region is drawn at the surface's OWN corner, which is what
  // the child below reads it back with. Both rects the `box` — the frame's own
  // coordinates on a surface the size of the box — is the draw that leaves the
  // surface's left and top strips empty, clips its far side, and hands the
  // composite a patch of the frame from boxLeft, boxTop up and to the left of the
  // spot: a repair the wrong shade, on a ramp a far-off one. `box, onto`, and the
  // case at the foot of this file is what keeps it that way.
  const onto = CanvasKit.XYWHRect(0, 0, boxRight - boxLeft, boxBottom - boxTop);
  const fillBox = surface(boxRight - boxLeft, boxBottom - boxTop);
  for (let k = levelCount; k >= 1; k--) {
    const sigma = HEAL_FILL_BLUR * widest * 2 ** (k - 1);
    const blurPaint = new CanvasKit.Paint();
    blurPaint.setImageFilter(CanvasKit.ImageFilter.MakeBlur(sigma, sigma, CanvasKit.TileMode.Clamp, null));
    fillBox.getCanvas().drawImageRect(cutImage, box, onto, blurPaint);
  }
  fillBox.getCanvas().drawImageRect(cutImage, box, onto, new CanvasKit.Paint());
  const fillImage = fillBox.makeImageSnapshot();

  const fillEffect = CanvasKit.RuntimeEffect.Make(HEAL_FILL_SKSL);
  assert.ok(fillEffect, 'HEAL_FILL_SKSL does not compile');
  const out = surface(W, H);
  const paint = new CanvasKit.Paint();
  paint.setShader(fillEffect.makeShaderWithChildren([], [asChild(src), asChildAt(fillImage, boxLeft, boxTop), asChild(mask)]));
  out.getCanvas().drawRect(CanvasKit.XYWHRect(0, 0, W, H), paint);
  return { pixels: pixelsOf(out.makeImageSnapshot()), mask: pixelsOf(mask), boxLeft, boxTop };
}

const flat = bytesOf(shadeFlat);
const made = healPass(flat, spots);
const alphaAt = (x, y) => made.mask[(y * W + x) * 4 + 3];
const inMask = (x, y) => alphaAt(x, y) > 200;
const rgbAt = (x, y) => {
  const i = (y * W + x) * 4;
  return [made.pixels[i], made.pixels[i + 1], made.pixels[i + 2]];
};

// 1. The mask is one shape. Painting a stroke as nine overlapping spots has to
//    leave one painted region, not nine: every painted pixel is reachable from
//    the first spot's centre without ever stepping off the mask. This is what
//    "mask liền mạch" is, and it is the thing the row of circles failed.
const painted = [];
for (let y = 0; y < H; y++) for (let x = 0; x < W; x++) if (inMask(x, y)) painted.push(y * W + x);
const seen = new Set();
const queue = [painted[0]];
seen.add(painted[0]);
while (queue.length) {
  const p = queue.pop();
  const x = p % W;
  const y = (p - x) / W;
  for (const [dx, dy] of [
    [1, 0],
    [-1, 0],
    [0, 1],
    [0, -1],
  ]) {
    const nx = x + dx;
    const ny = y + dy;
    if (nx < 0 || ny < 0 || nx >= W || ny >= H) continue;
    const q = ny * W + nx;
    if (seen.has(q) || !inMask(nx, ny)) continue;
    seen.add(q);
    queue.push(q);
  }
}
assert.equal(seen.size, painted.length, `${painted.length - seen.size} painted pixels are not reachable from the rest of the stroke`);
// ...and the mask is solid in the middle and clear well outside it.
assert.equal(alphaAt(Math.round(cx), Math.round(cy)), 255, 'the middle of the stroke is not fully painted');
assert.equal(alphaAt(2, 2), 0, 'the mask leaked into the corner of the frame');

// 2. The dust is gone and nothing was borrowed to do it: every painted pixel
//    ends on the ground's own value, 3.2 radii from a red card that a patch
//    search would have pasted in.
let worstIn = 0;
for (const p of painted) {
  const i = p * 4;
  for (let c = 0; c < 3; c++) worstIn = Math.max(worstIn, Math.abs(made.pixels[i + c] - GREY));
}
assert.ok(worstIn <= 4, `a painted pixel is off the ground it sits in by ${worstIn} levels (dust left, or a patch borrowed)`);

// 3. Nothing survives INSIDE the stroke either: neighbouring painted pixels
//    differ by a level or two, the way the ground does, where a spot solved on
//    its own would step at every seam.
let worstStep = 0;
for (let y = 0; y + 1 < H; y++) {
  for (let x = 0; x + 1 < W; x++) {
    if (!inMask(x, y) || !inMask(x + 1, y)) continue;
    worstStep = Math.max(worstStep, Math.abs(rgbAt(x, y)[0] - rgbAt(x + 1, y)[0]));
  }
}
assert.ok(worstStep <= 4, `the fill steps by ${worstStep} levels between neighbours inside the mask`);

// 4. The repair carries the LIGHT of where it landed. On a photo running from
//    dark to light across the frame, the painted pixels come back on that ramp
//    rather than at one flat value borrowed from somewhere else. The bound is
//    the ramp's own steepness, not zero: the fill reads the rim of the hole, and
//    where the hole takes weight off one side of a pixel its average slides that
//    way by about a pixel's worth of the ramp — 0.75 levels a pixel across this
//    frame. On the flat ground of 2 and 3 it is nothing at all.
const ramped = healPass(bytesOf(shadeRamp), spots);
let worstRamp = 0;
for (let y = 0; y < H; y++) {
  for (let x = 0; x < W; x++) {
    if (ramped.mask[(y * W + x) * 4 + 3] < 250) continue;
    worstRamp = Math.max(worstRamp, Math.abs(ramped.pixels[(y * W + x) * 4] - shadeRamp(x)[0]));
  }
}
assert.ok(worstRamp <= 6, `the fill is off the ramp it sits in by ${worstRamp} levels`);

// 5. Outside the mask, nothing moved. The pass is asked for the painted pixels
//    and no others — and "painted" is the whole brush, feather included: the last
//    15% of a spot is where the repair is mixed back in, so it is part of the
//    repair. Past the brush's own edge the photo is the photo.
let outWorst = 0;
let outAt = 0;
for (let p = 0; p < W * H; p++) {
  const x = p % W;
  const y = (p - x) / W;
  if (alphaAt(x, y) > 0) continue;
  for (let c = 0; c < 3; c++) {
    const off = Math.abs(made.pixels[p * 4 + c] - flat[p * 4 + c]);
    if (off > outWorst) {
      outWorst = off;
      outAt = x;
    }
  }
}
assert.equal(outWorst, 0, `a pixel outside the mask (x=${outAt}) changed by ${outWorst} levels`);

// 6. The same photo and the same spots give the same photo, byte for byte.
assert.deepEqual(
  Buffer.from(healPass(flat, spots).pixels),
  Buffer.from(made.pixels),
  'the pass is not deterministic'
);

// 7. A mask with no rim to read from — one spot over the whole frame — leaves
//    the photo alone. Nothing is painted black, because nothing was found to
//    paint it with, and a frame is still a frame.
const all = healPass(flat, [{ x: 0.5, y: 0.5, r: 2 }]);
let worstAll = 0;
for (let p = 0; p < W * H; p++) {
  for (let c = 0; c < 3; c++) worstAll = Math.max(worstAll, Math.abs(all.pixels[p * 4 + c] - flat[p * 4 + c]));
}
assert.equal(worstAll, 0, `a frame-wide mask changed the photo by ${worstAll} levels`);

// 8. The shape the brush is actually used in: one long thin run right across the
//    frame, on a photo that runs dark to light along it. Its rim is the pixels
//    just above and below it wherever it passes, so the repair comes back on the
//    ramp there — not at the frame's average colour. A hole this long needs the
//    level count to grow with it: too few levels and the middle of the run is
//    farther from any rim than the widest blur reaches, so the fill comes back
//    the average of a blur too wide to have read the rim — one flat mark.
const band = [];
for (let i = 0; i <= 60; i++) band.push({ x: 0.12 + (0.76 * i) / 60, y: 0.5, r: 0.02 });
const across = healPass(bytesOf(shadeRamp), band);
let worstBand = 0;
let bandAt = 0;
for (let y = 0; y < H; y++) {
  for (let x = 0; x < W; x++) {
    if (across.mask[(y * W + x) * 4 + 3] < 250) continue;
    const off = Math.abs(across.pixels[(y * W + x) * 4] - shadeRamp(x)[0]);
    if (off > worstBand) {
      worstBand = off;
      bandAt = x;
    }
  }
}
assert.ok(worstBand <= 8, `a painted pixel at x=${bandAt} is ${worstBand} levels off the rim it sits in`);

// 9. The one case every size above hides: a frame big enough that the fill's box
//    sits INSIDE it. On a 160px frame with a 8px brush the padding is wider than
//    the frame, so the box is the frame and its corner is (0,0) — a fill drawn at
//    the box's coordinates and one drawn at the surface's own are then the same
//    draw, and a repair off by the box's corner is off by nothing. On a real photo
//    it is not nothing: a brush at the sensor-dust end of the scale, a frame a few
//    thousand pixels across, and the corner is hundreds of pixels up and to the
//    left of the spot. That is the shape the reader heals with, and on a ramp the
//    patch comes back at the ramp's value hundreds of pixels away — the healed
//    area is the wrong shade, which is what this case is here to catch.
{
  const keptW = W;
  const keptH = H;
  // Tall enough that the box clears BOTH edges: the padding is three blur
  // radii of the widest brush, which on this frame is 192px, so a shorter
  // frame would push boxTop to 16 and the corner would nearly be the edge.
  W = 640;
  H = 720;
  const shadeWide = (x) => {
    const v = 40 + (120 * x) / (W - 1);
    return [v, v, v];
  };
  const spot = [{ x: 0.5, y: 0.5, r: 0.05 }]; // a 32px brush on this frame
  const wide = healPass(bytesOf(shadeWide), spot);
  assert.ok(
    wide.boxLeft > 16 && wide.boxTop > 16,
    `the fill's box is not inset on this frame (boxLeft=${wide.boxLeft}, boxTop=${wide.boxTop}) — this case is no longer testing the bug it exists for`
  );
  let worstWide = 0;
  let wideAt = 0;
  for (let y = 0; y < H; y++) {
    for (let x = 0; x < W; x++) {
      if (wide.mask[(y * W + x) * 4 + 3] < 250) continue;
      const off = Math.abs(wide.pixels[(y * W + x) * 4] - shadeWide(x)[0]);
      if (off > worstWide) {
        worstWide = off;
        wideAt = x;
      }
    }
  }
  assert.ok(
    worstWide <= 4,
    `a spot healed at the middle of a ${W}x${H} frame came back ${worstWide} levels off at x=${wideAt} — the fill is read at the box's corner, not at the spot (box is at ${wide.boxLeft},${wide.boxTop})`
  );
  W = keptW;
  H = keptH;
}

console.log('ok  the stroke is one mask, the fill carries the light, and nothing outside it moved');