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) and feathers the edge, 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.
//
// 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.
export const HEAL_MAX = 16;
// 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.
export const HEAL_FEATHER = 0.55;
// Brush sizes, S/M/L, as fractions of the photo's width — the labels the chips
// wear sit beside them so the two cannot drift apart. The default is the
// sensor-dust end: a speck is a few thousandths of the frame.
export const HEAL_SIZES = [0.012, 0.03, 0.07];
export const HEAL_SIZE_LABELS = ['S', 'M', 'L'];
export const HEAL_DEFAULT_R = HEAL_SIZES[0];
// 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 patch" when two patches are compared: the centre
// and eight points on a ring at half the radius — the scale the eye reads a
// spot's surroundings at.
const PATCH_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],
];

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, at most HEAL_MAX
// of them. 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
    .slice(0, HEAL_MAX)
    .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 HEAL_SKSL 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.
export function healUniforms(spots: HealSpot[], width: number, height: number): Float32Array {
  const u = new Float32Array((HEAL_MAX * 2 + 1) * 4);
  const list = readHeal(spots);
  for (let i = 0; i < list.length; i++) {
    const s = list[i];
    u.set([s.x, s.y, s.r, 0], i * 4);
    u.set([s.sx, s.sy, 0, 0], (HEAL_MAX + i) * 4);
  }
  u.set([width, height, HEAL_FEATHER, 0], HEAL_MAX * 2 * 4);
  return u;
}

// 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 buffer is the same HEAL_MAX long whatever the recipe holds.
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}];
        half m = half(1.0 - smoothstep(rad * size.z, rad, d));
        c = mix(c, img.eval(pos + (t.xy - s.xy) * size.xy), 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 const HEAL_SKSL = `
uniform shader img;
uniform float4 spots[${HEAL_MAX}];
uniform float4 srcs[${HEAL_MAX}];
uniform float4 size;

half4 main(float2 pos) {
  half4 c = img.eval(pos);${Array.from({ length: HEAL_MAX }, (_, 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. The winner is the one whose surroundings look most like the spot's.
// Returns null when the frame is too small to hold any candidate: the caller
// then leaves the spot where it is rather than inventing a patch off the photo.
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 score = (cx: number, cy: number) => {
    let sum = 0;
    for (const [dx, dy] of PATCH_TAPS) {
      const a = sample(clamp01(x + dx * r), clamp01(y + dy * r));
      const b = sample(clamp01(cx + dx * r), clamp01(cy + dy * r));
      sum += (Math.abs(a.r - b.r) + Math.abs(a.g - b.g) + Math.abs(a.b - b.b)) / 3;
    }
    return sum / PATCH_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);
        // 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;
}