// Voronoi / Worley noise compute shader
// Generates cellular noise patterns by finding distances to random feature points.
// Supports F1 (nearest), F2 (second nearest), and F2-F1 (cell edges) output modes,
// with Euclidean, Manhattan, and Chebyshev distance metrics.

struct Uniforms {
  time: f32,
  scale: f32,
  jitter: f32,
  width: f32,
  height: f32,
  mode: u32,    // 0 = F1, 1 = F2, 2 = F2-F1
  metric: u32,  // 0 = euclidean, 1 = manhattan, 2 = chebyshev
  _pad: u32,
  offset_x: f32,
  offset_y: f32,
  _pad2: u32,
  _pad3: u32,
}

@group(0) @binding(0) var<uniform> u: Uniforms;
@group(0) @binding(1) var output: texture_storage_2d<rgba8unorm, write>;

// Hash function for deterministic random points per cell.
// Based on hash-without-sine by Dave Hoskins.
fn hash2(p: vec2f) -> vec2f {
  var p3 = fract(vec3f(p.x, p.y, p.x) * vec3f(0.1031, 0.1030, 0.0973));
  p3 += dot(p3, p3.yzx + 33.33);
  return fract((p3.xx + p3.yz) * p3.zy);
}

// Distance calculation using the selected metric
fn cell_distance(a: vec2f, b: vec2f) -> f32 {
  let d = abs(a - b);
  if (u.metric == 1u) {
    // Manhattan distance
    return d.x + d.y;
  }
  if (u.metric == 2u) {
    // Chebyshev distance
    return max(d.x, d.y);
  }
  // Euclidean distance
  return length(a - b);
}

// Voronoi evaluation: searches 3x3 cell neighborhood around the sample point.
// Returns vec2f(F1, F2) — distances to the nearest and second nearest feature points.
fn voronoi(p: vec2f) -> vec2f {
  let cell = floor(p);
  let local = fract(p);

  var f1 = 999.0;
  var f2 = 999.0;

  // Search the 3x3 neighborhood of cells
  for (var y = -1; y <= 1; y++) {
    for (var x = -1; x <= 1; x++) {
      let neighbor = vec2f(f32(x), f32(y));

      // Deterministic random point within this cell.
      // The hash seed is fixed per cell — animation comes from sine/cosine offset, not re-hashing.
      let cell_coord = cell + neighbor;
      let rand = hash2(cell_coord);

      // Animate by smoothly orbiting around the base position
      let phase = hash2(cell_coord + vec2f(127.1, 311.7));
      let offset = vec2f(
        sin(u.time + phase.x * 6.28) * 0.5,
        cos(u.time + phase.y * 6.28) * 0.5
      );

      // Jitter controls how far the point can move from the cell center.
      // jitter=0 places points at cell centers (regular grid); jitter=1 is fully random.
      let point = neighbor + (rand + offset * 0.3) * u.jitter;

      let dist = cell_distance(local, point);

      // Maintain sorted order: f1 <= f2
      if (dist < f1) {
        f2 = f1;
        f1 = dist;
      } else if (dist < f2) {
        f2 = dist;
      }
    }
  }

  return vec2f(f1, f2);
}

@compute @workgroup_size(8, 8)
fn main(@builtin(global_invocation_id) id: vec3u) {
  let dims = vec2u(u32(u.width), u32(u.height));
  if (id.x >= dims.x || id.y >= dims.y) {
    return;
  }

  // Normalize pixel coordinates and apply scale
  let uv = vec2f(f32(id.x), f32(id.y)) / vec2f(u.width, u.height);
  let p = uv * u.scale + vec2f(u.offset_x, u.offset_y);

  let distances = voronoi(p);

  // Select output value based on mode
  var value: f32;
  if (u.mode == 0u) {
    // F1: distance to nearest point (dark at cell centers, bright at edges)
    value = distances.x;
  } else if (u.mode == 1u) {
    // F2: distance to second nearest point
    value = distances.y;
  } else {
    // F2-F1: highlights cell boundaries (thin bright lines between cells)
    value = distances.y - distances.x;
  }

  // Normalize to [0, 1] range — distances can exceed 1.0 depending on metric and scale,
  // so clamp after a reasonable normalization
  value = clamp(value, 0.0, 1.0);

  textureStore(output, vec2i(id.xy), vec4f(value, value, value, 1.0));
}

// Voronoi / Worley noise texture generator
// Generates cellular noise patterns on the GPU via a compute shader.
// Supports multiple distance metrics and output modes.
//
// Default WGSL loading uses a ?raw import (works with Vite, esbuild, Webpack).
// Alternative: load via fetch — see README.md for details.
import shaderSource from './voronoi.wgsl?raw';

export type VoronoiMode = 'f1' | 'f2' | 'f2-f1';
export type VoronoiMetric = 'euclidean' | 'manhattan' | 'chebyshev';

export interface VoronoiUpdateOptions {
  time?: number;
  scale?: number;
  jitter?: number;
  offset?: [number, number];
  mode?: VoronoiMode;
  metric?: VoronoiMetric;
}

export interface Voronoi {
  update(target: GPUTexture, options?: VoronoiUpdateOptions): void;
  destroy(): void;
}

const MODE_MAP: Record<VoronoiMode, number> = { f1: 0, f2: 1, 'f2-f1': 2 };
const METRIC_MAP: Record<VoronoiMetric, number> = { euclidean: 0, manhattan: 1, chebyshev: 2 };

// Uniform buffer layout (std140-aligned):
// f32 time, f32 scale, f32 jitter, f32 width,
// f32 height, u32 mode, u32 metric, u32 _pad,
// f32 offset_x, f32 offset_y, u32 _pad2, u32 _pad3
const UNIFORM_SIZE = 48; // 12 x 4 bytes

export function createVoronoi(device: GPUDevice): Voronoi {
  const uniformBuffer = device.createBuffer({
    size: UNIFORM_SIZE,
    usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST
  });

  const shaderModule = device.createShaderModule({ code: shaderSource });

  const bindGroupLayout = device.createBindGroupLayout({
    entries: [
      {
        binding: 0,
        visibility: GPUShaderStage.COMPUTE,
        buffer: { type: 'uniform' }
      },
      {
        binding: 1,
        visibility: GPUShaderStage.COMPUTE,
        storageTexture: { format: 'rgba8unorm', access: 'write-only' }
      }
    ]
  });

  const pipeline = device.createComputePipeline({
    layout: device.createPipelineLayout({ bindGroupLayouts: [bindGroupLayout] }),
    compute: { module: shaderModule, entryPoint: 'main' }
  });

  let lastTarget: GPUTexture | null = null;
  let bindGroup: GPUBindGroup | null = null;

  const uniformData = new ArrayBuffer(UNIFORM_SIZE);
  const f32 = new Float32Array(uniformData);
  const u32 = new Uint32Array(uniformData);

  function update(target: GPUTexture, opts: VoronoiUpdateOptions = {}): void {
    if (lastTarget !== target) {
      bindGroup = device.createBindGroup({
        layout: bindGroupLayout,
        entries: [
          { binding: 0, resource: { buffer: uniformBuffer } },
          { binding: 1, resource: target.createView() }
        ]
      });
      lastTarget = target;
    }

    const width = target.width;
    const height = target.height;

    const offset = opts.offset ?? [0, 0];
    f32[0] = opts.time ?? 0;
    f32[1] = opts.scale ?? 8.0;
    f32[2] = opts.jitter ?? 0.9;
    f32[3] = width;
    f32[4] = height;
    u32[5] = MODE_MAP[opts.mode ?? 'f1'];
    u32[6] = METRIC_MAP[opts.metric ?? 'euclidean'];
    u32[7] = 0; // padding
    f32[8] = offset[0];
    f32[9] = offset[1];
    u32[10] = 0; // padding
    u32[11] = 0; // padding

    device.queue.writeBuffer(uniformBuffer, 0, uniformData);

    const encoder = device.createCommandEncoder();
    const pass = encoder.beginComputePass();
    pass.setPipeline(pipeline);
    pass.setBindGroup(0, bindGroup!);
    pass.dispatchWorkgroups(Math.ceil(width / 8), Math.ceil(height / 8));
    pass.end();
    device.queue.submit([encoder.finish()]);
  }

  function destroy(): void {
    uniformBuffer.destroy();
  }

  return { update, destroy };
}