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.
//
// 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. The wheel walks
// between the bounds multiplicatively, so a trackpad's small deltas and a
// mouse's 100px notch are the same gesture at two speeds; the default is the
// sensor-dust end, where a speck is a few thousandths of the frame.
export const HEAL_DEFAULT_R = 0.012;
export const HEAL_MIN_R = 0.003;
export const HEAL_MAX_R = 0.25;
// How far apart a drag lays its spots, in radii: any closer and the run is the
// same repair written twice, which is work for nothing.
export const HEAL_SPACING = 0.6;
// 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 brush radius after a wheel notch. Multiplicative, so the brush grows by
// the same proportion wherever it starts, and bounded — past HEAL_MAX_R a spot
// would borrow a patch from off the photo, and below HEAL_MIN_R it would be
// finer than the pixels it is drawn on at preview scale.
export function wheelHealR(r: number, deltaY: number): number {
  const next = r * Math.exp(-deltaY * 0.0015);
  return next < HEAL_MIN_R ? HEAL_MIN_R : next > HEAL_MAX_R ? HEAL_MAX_R : next;
}

// 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, HEAL_FEATHER, 0], n * 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 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).
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 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. 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;
}