import type { HealSpot } from '../types';

// FX tab > HEAL — the dust brush, and the patch search behind it.
//
// A spot is a circle on the rendered photo plus the patch it borrows: the
// renderer copies the pixels at (sx, sy) onto (x, y), feathers the edge, and
// shifts what it pasted onto the light of the place it landed in, so a repair
// is a draw of the same picture rather than a blur over the dust. All five
// numbers are fractions of the render — x/y/sx/sy of its width and height, r of
// its width — which is what makes one set of spots survive the preview and the
// export rendering the same photo at two sizes, and keeps the circle round
// whatever the photo's shape.
//
// There is no ceiling on the list. The shader is built to carry exactly the
// spots the recipe holds (healSkSL), so a new repair can never push an old one
// out: the dust you healed first is the dust that stays healed.
//
// The source is SEARCHED for rather than asked for. Lightroom picks the patch
// and lets you drag the second circle afterwards; the search below is the same
// idea without the second circle, and it is a pure function of a sampler so a
// synthetic picture can hold it to account.
//
// Feather, as a fraction of the radius: inside it the patch is copied, outward
// it fades to nothing, so the circle leaves no rim of its own. It is the outer
// 15% and no more, because that band is the only place the dust being repaired
// is mixed back into the patch — a wide fade keeps the speck's own edge alive
// as a faint ring inside the circle, which is a blur of the dust rather than a
// repair. sub-pixel at the default brush, still a soft edge at a big one.
export const HEAL_FEATHER = 0.85;
// The brush's radius, as a fraction of the photo's width — where this tool
// starts, at the sensor-dust end, where a speck is a few thousandths of the
// frame. The bounds and the wheel that moves it between them belong to the
// brush both FX tools share (brush.ts).
export const HEAL_DEFAULT_R = 0.012;
// How far the search looks, in radii, and how many directions it looks in.
const SEARCH_DISTANCES = [2.6, 4.2, 6.5];
const SEARCH_DIRS = 8;
// The taps that stand for "the light around a spot": a ring just OUTSIDE the
// brush, where the dust being repaired is not — the scale the eye reads a
// spot's surroundings at. Reading the inside instead compares a candidate with
// the very dark the repair is trying to erase, and the patch that matches a
// speck best is then the one carrying a speck of its own — which is how a
// repair ends up moving dust a few pixels instead of removing it.
//
// Twelve taps read as their MEDIAN, not their mean. The ring can only be a
// little way out — far enough out and it would be reading a light the spot does
// not sit in — so a few of its taps land on the speck's own softened edge, or
// on a neighbour's: readings in the minority, which the median drops and the
// mean would drag the whole light down by.
const RING_R = 1.15;
const RING_TAPS: [number, number][] = Array.from({ length: 12 }, (_, i) => {
  const a = (i / 12) * Math.PI * 2;
  return [Math.cos(a), Math.sin(a)] as [number, number];
});
// The taps that decide whether a candidate patch is clean: the inside of the
// patch, held against the patch's own mean. Dust is an outlier in its own
// neighbourhood; skin, however grainy, is not. The last ring is the patch's own
// edge — the circle is copied at full strength out to HEAL_FEATHER of its
// radius, so that is where a neighbour's dust leaking into the patch shows up,
// and the middle of the patch would never see it.
const INSIDE_TAPS: [number, number][] = [
  [0, 0],
  [-0.5, 0],
  [0.5, 0],
  [0, -0.5],
  [0, 0.5],
  [-0.35, -0.35],
  [0.35, -0.35],
  [-0.35, 0.35],
  [0.35, 0.35],
  ...Array.from({ length: 8 }, (_, i) => {
    const a = (i / 8) * Math.PI * 2;
    return [Math.cos(a) * HEAL_FEATHER, Math.sin(a) * HEAL_FEATHER] as [number, number];
  }),
];
// How far a candidate's light may sit from the spot's own before it is refused
// outright, in the 0-255 levels the sampler answers in. Sending a patch from a
// different light is exactly what makes a repair show up as a mark of its own —
// a light square where dark ground was, a dark one where light was — so past
// this the search gives up and the caller leaves the speck alone. The user sees
// an untouched speck, which is honest, and tries another size or another spot.
const LIGHT_GATE = 20;
// Light outweighs cleanliness in the score, so the winner is the same light
// first and the cleaner patch second; within the gate the cleanliness term only
// breaks ties between patches the eye would call the same.
const LIGHT_WEIGHT = 3;

const num = (v: unknown, fallback: number) => {
  const n = Number(v);
  return Number.isFinite(n) ? n : fallback;
};
const clamp01 = (v: number) => (v < 0 ? 0 : v > 1 ? 1 : v);

// The stored spots, made readable: numbers, inside the frame. Everything below
// reads a recipe through this, so a hand-written or older file cannot produce a
// spot the brush and the renderer disagree about.
export function readHeal(heal: HealSpot[] | undefined): HealSpot[] {
  if (!Array.isArray(heal)) return [];
  return heal
    .map((s) => ({
      x: clamp01(num(s?.x, 0)),
      y: clamp01(num(s?.y, 0)),
      r: Math.max(0, num(s?.r, 0)),
      sx: clamp01(num(s?.sx, s?.x ?? 0)),
      sy: clamp01(num(s?.sy, s?.y ?? 0)),
    }))
    .filter((s) => s.r > 0);
}

// The uniform block the shader for `n` spots reads: the circles, the patches,
// then the frame the fractions are of. Declaration order, arrays expanded —
// that is how the runtime effect wants its uniforms, and one buffer is one
// upload per render. Its length is a function of the list, not a fixed
// capacity, because the shader carries exactly the spots the recipe holds.
export function healUniforms(spots: HealSpot[], width: number, height: number): Float32Array {
  const list = readHeal(spots);
  const n = list.length;
  const u = new Float32Array((n * 2 + 1) * 4);
  for (let i = 0; i < n; i++) {
    const s = list[i];
    u.set([s.x, s.y, s.r, 0], i * 4);
    u.set([s.sx, s.sy, 0, 0], (n + i) * 4);
  }
  u.set([width, height, 0, 0], n * 2 * 4);
  return u;
}

// The ring the pasted circle is matched at, and how its taps are laid out.
// Measured just OUTSIDE the brush — RING_R of its radius, where the search
// reads a spot's light too — because that is the only ring with nothing of the
// repair in it: inside the circle is the borrowed patch at one radius and the
// speck being covered at another, and a light read off either of those is the
// dust talking rather than the photo. Outside, both sides are photographs — the
// ground the repair has to sit in, and the ground the patch was borrowed from.
// Sixteen taps are spread around it, at the SAME sixteen places on both sides,
// so a tap's pair of readings differ by the light between the two places and
// not by where in the photo they were taken. Sixteen rather than eight because
// a tap that lands on an edge bends its own direction only as far as the next
// ring tap: the finer the ring, the less of the rim one hard pixel of the photo
// can reach around.
const BLEND_TAPS = 16;
// ...and one pixel further out again than RING_R. The ring is read by a pixel
// that lies under the brush, so the reading has to clear the dust's own edge,
// and that edge is soft: a speck r pixels across spreads about a pixel past
// where it is drawn, in the pixels the shader samples. RING_R alone leaves the
// nearest tap or two inside that spread, and a tap that lands on the dust is
// the dust's light — which is exactly the blotch this reading exists to avoid.
// One pixel is the width of the spread, in the same pixels the sampler works
// in, so it is a pixel here rather than a fraction of the radius: a fraction
// would be nothing at all at the sensor-dust end of the brush.
const RING_PAD = 1;

// ...but one light for the whole circle is a light no ring has when the ring is
// not all one thing. A twig against the sky, a hairline, the edge of a table
// under the brush: a few taps of the ring land on the thing and the rest on the
// ground, and the MEAN of them all is a colour that is neither — a level the
// place does not have, pasted over the whole circle. The donor can be of exactly
// the right light and still come out as a blotch, which is the mark the repair
// was supposed to stop leaving. So the light is read once PER DIRECTION, and a
// pixel takes the correction of the two readings it lies between, interpolated
// by its own angle around the spot: the rim then meets the place all the way
// round instead of on average, and the taps that landed on the twig correct the
// part of the rim that is near the twig rather than the whole of it. It is the
// same taps and the same one draw — each pixel reads them at its own angle
// instead of sharing everyone's single number.

// One unrolled block per spot. SkSL indexes a uniform array by constant only
// (see TONE_SKSL's mixer), so the spots are written out rather than looped, and
// the shader is built for the count it is handed rather than for a capacity —
// that is what lets the list be uncapped. A count costs one RuntimeEffect to
// compile, so the renderer caches them by count (exportEngine's healEffectFor).
//
// What is pasted is the patch's own pixels SHIFTED onto the light of the place
// it lands in — the gradient of the borrowed texture, carried at the level of
// the destination: the border is made to match, and everything inside keeps the
// texture it was borrowed for. That is the whole of "seamless" the brush needs:
// at the rim the patch sits within a level or two of the photo around it,
// instead of up to LIGHT_GATE levels off it, so the repair stops reading as a
// soft blotch of its own and the feather has almost nothing left to hide. The
// correction is read per direction around the ring, so the border is matched
// where it is rather than on average, and the pixels of a twig that run through
// the ring bend the part of the rim they are near instead of the whole circle.
//
// ponytail: the correction lives on the rim — every pixel takes the two ring
// readings it lies between and interpolates — so it is the boundary of a Poisson
// edit and not its interior: a repair laid over something the patch cannot
// reproduce still carries the patch's own texture inside, bent to fit the rim.
// The other half of the solve (a correction that relaxes inside the circle, a
// Jacobi ping-pong per spot) lands when a real photo shows an interior the eye
// can find from the rim alone. A direction's level is also one tap pair on each
// side, so the grain the mean used to average away now rides the rim as a wedge
// of a level or two across the circle — a tangential average of three taps per
// direction is the rung for that, when a photo shows the wedge.
const spotBlock = (i: number) => `
  {
    float4 s = spots[${i}];
    if (s.z > 0.0) {
      float rad = s.z * size.x;
      float d = distance(pos, s.xy * size.xy);
      if (d < rad) {
        float4 t = srcs[${i}];
        float2 pd = s.xy * size.xy;
        float2 ps = t.xy * size.xy;
        float r = rad * ${HEAL_FEATHER};
        float rr = rad * ${RING_R} + ${RING_PAD.toFixed(1)};
        // Which way round the spot this pixel sits, one millionth of a pixel off
        // the exact centre so the angle is defined there too.
        float ang = atan(pos.y - pd.y + 1e-6, pos.x - pd.x + 1e-6);
        half3 corr = half3(0.0);
        float wsum = 0.0;
        for (int k = 0; k < ${BLEND_TAPS}; k++) {
          float a = float(k) * 6.283185307 / float(${BLEND_TAPS});
          float2 o = float2(cos(a), sin(a)) * rr;
          // How far this pixel's angle is from this tap's, the short way round.
          float da = ang - a;
          da = abs(da - 6.283185307 * floor(da * 0.1591549431 + 0.5));
          float w = max(0.0, 1.0 - da * float(${BLEND_TAPS}) * 0.1591549431);
          corr += half(w) * (img.eval(pd + o).rgb - img.eval(ps + o).rgb);
          wsum += w;
        }
        corr *= half(1.0 / wsum);
        half m = half(1.0 - smoothstep(r, rad, d));
        c = mix(c, img.eval(pos + ps - pd) + half4(corr, 0.0), m);
      }
    }
  }
`;

// The pass. It reads the pixels the pipeline has already built (the child is a
// snapshot of the surface) and writes the borrowed patches back over them, so a
// repair is one draw: no blur, no smoothing, and the grain and the frame land
// on top of it afterwards exactly as they land on the rest of the photo.
export function healSkSL(count: number): string {
  return `
uniform shader img;
uniform float4 spots[${count}];
uniform float4 srcs[${count}];
uniform float4 size;

half4 main(float2 pos) {
  half4 c = img.eval(pos);${Array.from({ length: count }, (_, i) => spotBlock(i)).join('')}
  return c;
}
`;
}

// The patch to borrow for a spot at (x, y) of radius r, from a sampler that
// answers fractions of the same photo. The candidates are a ring of offsets in
// eight directions at three distances — the patch has to be far enough that the
// dust is not in it, near enough that the light is the same — plus each one
// mirrored through the spot, which is the pair Lightroom's own auto-source
// leans on.
//
// Each candidate is asked two questions. "Is this the same light?" holds the
// mean of what would be pasted against the mean of the spot's own ring, the
// pixels just outside the dust — the light the repair has to sit in, read where
// the dust is not, because comparing against the spot's own inside compares a
// candidate with the very dark the repair is trying to erase, and the patch
// that matches a speck best is then the one carrying a speck of its own. "Is
// there dust on it?" holds the candidate's inside against its own inside mean:
// dust is an outlier in its own neighbourhood, grain is not. A candidate whose
// light is off by more than LIGHT_GATE is not scored at all — a patch from the
// wrong light is a mark of its own, and moving it a few pixels for a wrong
// clone is not a repair.
//
// Returns null when the frame is too small to hold any candidate, or when every
// candidate is refused: the caller then leaves the speck where it is rather
// than pasting a patch it cannot stand behind.
export function findHealSource(
  sample: (fx: number, fy: number) => { r: number; g: number; b: number },
  x: number,
  y: number,
  r: number
): { sx: number; sy: number } | null {
  if (!(r > 0)) return null;
  const inside = (cx: number, cy: number) => cx - r >= 0 && cx + r <= 1 && cy - r >= 0 && cy + r <= 1;
  const rgb = (v: { r: number; g: number; b: number }) => (v.r + v.g + v.b) / 3;
  // The light the repair has to sit in: the median of the spot's ring, where
  // the dust is not. Measured once, because it is the same number for every
  // candidate.
  const ring = RING_TAPS.map(([dx, dy]) => rgb(sample(clamp01(x + dx * RING_R * r), clamp01(y + dy * RING_R * r)))).sort(
    (a, b) => a - b
  );
  const target = (ring[RING_TAPS.length / 2 - 1] + ring[RING_TAPS.length / 2]) / 2;
  // The score of one candidate: null when it is refused for its light, else how
  // far off that light is times its weight, plus how much the patch varies
  // inside itself.
  const score = (cx: number, cy: number): number | null => {
    let content = 0;
    const inside: number[] = [];
    for (const [dx, dy] of INSIDE_TAPS) {
      const v = rgb(sample(clamp01(cx + dx * r), clamp01(cy + dy * r)));
      inside.push(v);
      content += v;
    }
    content /= INSIDE_TAPS.length;
    const light = Math.abs(target - content);
    if (light > LIGHT_GATE) return null;
    let dirty = 0;
    for (const v of inside) dirty += Math.abs(v - content);
    return light * LIGHT_WEIGHT + dirty / INSIDE_TAPS.length;
  };
  let best: { sx: number; sy: number; score: number } | null = null;
  for (let d = 0; d < SEARCH_DIRS; d++) {
    const a = (d / SEARCH_DIRS) * Math.PI * 2;
    for (const dist of SEARCH_DISTANCES) {
      const cx = x + Math.cos(a) * dist * r;
      const cy = y + Math.sin(a) * dist * r;
      for (const [px, py] of [
        [cx, cy],
        [2 * x - cx, 2 * y - cy],
      ]) {
        if (!inside(px, py)) continue;
        const s = score(px, py);
        if (s === null) continue;
        // A tie keeps the earlier candidate: the ring is walked from the right,
        // so the patch nearest the spot wins — the one most likely to share its
        // light.
        if (!best || s < best.score) best = { sx: px, sy: py, score: s };
      }
    }
  }
  return best ? { sx: best.sx, sy: best.sy } : null;
}