// 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.
//
//   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_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');

// --- 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)),
});

const W = 160;
const 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 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 pyramid down and up again, 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);

  const levels = [];
  let carried = cut.makeImageSnapshot();
  for (let k = 1; k <= HEAL_MAX_LEVELS; k++) {
    const lw = Math.max(1, W >> k);
    const lh = Math.max(1, H >> k);
    const level = surface(lw, lh);
    level
      .getCanvas()
      .drawImageRectOptions(
        carried,
        CanvasKit.XYWHRect(0, 0, carried.width(), carried.height()),
        CanvasKit.XYWHRect(0, 0, lw, lh),
        CanvasKit.FilterMode.Linear,
        CanvasKit.MipmapMode.None
      );
    levels.push({ surf: level, w: lw, h: lh });
    carried = level.makeImageSnapshot();
    if (lw === 1 && lh === 1) break;
  }
  let up = null;
  let fillImage = null;
  for (let k = levels.length - 1; k >= 0; k--) {
    const { surf, w: lw, h: lh } = levels[k];
    if (up) {
      surf
        .getCanvas()
        .drawImageRectOptions(
          up,
          CanvasKit.XYWHRect(0, 0, up.width(), up.height()),
          CanvasKit.XYWHRect(0, 0, lw, lh),
          CanvasKit.FilterMode.Linear,
          CanvasKit.MipmapMode.None
        );
    }
    const pulled = surf.makeImageSnapshot();
    if (k === 0) {
      fillImage = pulled;
      break;
    }
    up = pulled;
  }

  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), asChild(fillImage), asChild(mask)]));
  out.getCanvas().drawRect(CanvasKit.XYWHRect(0, 0, W, H), paint);
  return { pixels: pixelsOf(out.makeImageSnapshot()), mask: pixelsOf(mask) };
}

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 pyramid 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.
let outWorst = 0;
for (let p = 0; p < W * H; p++) {
  if (inMask(p % W, (p - (p % W)) / W)) continue;
  for (let c = 0; c < 3; c++) outWorst = Math.max(outWorst, Math.abs(made.pixels[p * 4 + c] - flat[p * 4 + c]));
}
assert.equal(outWorst, 0, `a pixel outside the mask 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`);

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