// Perlin & Simplex noise compute shader
// Generates animated 2D noise and writes it to a storage texture.
// Supports FBM layering and optional domain warping.

struct Uniforms {
  time: f32,
  scale: f32,
  offset_x: f32,
  offset_y: f32,
  octaves: u32,
  noise_type: u32,    // 0 = perlin, 1 = simplex
  domain_warp: u32,   // 0 = off, 1 = on
  warp_strength: f32,
  persistence: f32,
  lacunarity: f32,
  width: f32,
  height: f32,
}

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

// --- Hash functions ---
// Based on hash without sine by Dave Hoskins
// https://www.shadertoy.com/view/4djSRW

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

fn hash1(p: vec2f) -> f32 {
  var p3 = fract(vec3f(p.x, p.y, p.x) * 0.1031);
  p3 += dot(p3, p3.yzx + 33.33);
  return fract((p3.x + p3.y) * p3.z);
}

// --- Perlin noise ---
// Classic gradient noise with quintic interpolation

fn grad2(hash: vec2f) -> vec2f {
  let angle = hash.x * 6.283185;
  return vec2f(cos(angle), sin(angle));
}

fn perlin(p: vec2f) -> f32 {
  let i = floor(p);
  let f = fract(p);

  // Quintic interpolation curve (smoother than cubic)
  let u = f * f * f * (f * (f * 6.0 - 15.0) + 10.0);

  // Four corner gradients
  let g00 = dot(grad2(hash2(i + vec2f(0.0, 0.0))), f - vec2f(0.0, 0.0));
  let g10 = dot(grad2(hash2(i + vec2f(1.0, 0.0))), f - vec2f(1.0, 0.0));
  let g01 = dot(grad2(hash2(i + vec2f(0.0, 1.0))), f - vec2f(0.0, 1.0));
  let g11 = dot(grad2(hash2(i + vec2f(1.0, 1.0))), f - vec2f(1.0, 1.0));

  let x0 = mix(g00, g10, u.x);
  let x1 = mix(g01, g11, u.x);
  return mix(x0, x1, u.y);
}

// --- Simplex noise ---
// 2D simplex noise based on the approach by Ian McEwan / Ashima Arts

const SKEW: f32 = 0.36602540378;    // (sqrt(3) - 1) / 2
const UNSKEW: f32 = 0.21132486540;  // (3 - sqrt(3)) / 6

fn simplex(p: vec2f) -> f32 {
  // Skew input space to determine which simplex cell we're in
  let s = (p.x + p.y) * SKEW;
  let i = floor(p + s);
  let t = (i.x + i.y) * UNSKEW;
  let x0 = p - (i - t);

  // Determine which simplex triangle
  var i1: vec2f;
  if (x0.x > x0.y) {
    i1 = vec2f(1.0, 0.0);
  } else {
    i1 = vec2f(0.0, 1.0);
  }

  let x1 = x0 - i1 + UNSKEW;
  let x2 = x0 - 1.0 + 2.0 * UNSKEW;

  // Gradient contributions from the three corners
  var n = vec3f(0.0);

  var d0 = 0.5 - dot(x0, x0);
  if (d0 > 0.0) {
    d0 = d0 * d0;
    n.x = d0 * d0 * dot(grad2(hash2(i)), x0);
  }

  var d1 = 0.5 - dot(x1, x1);
  if (d1 > 0.0) {
    d1 = d1 * d1;
    n.y = d1 * d1 * dot(grad2(hash2(i + i1)), x1);
  }

  var d2 = 0.5 - dot(x2, x2);
  if (d2 > 0.0) {
    d2 = d2 * d2;
    n.z = d2 * d2 * dot(grad2(hash2(i + 1.0)), x2);
  }

  // Scale to [-1, 1] range (approximately)
  return 70.0 * (n.x + n.y + n.z);
}

// --- Noise selector ---

fn noise(p: vec2f) -> f32 {
  if (u.noise_type == 0u) {
    return perlin(p);
  }
  return simplex(p);
}

// --- Fractal Brownian Motion ---
// Layers multiple octaves of noise at increasing frequency and decreasing amplitude

fn fbm(p: vec2f) -> f32 {
  var value = 0.0;
  var amplitude = 0.5;
  var frequency = 1.0;
  var coord = p;

  for (var i = 0u; i < u.octaves; i++) {
    value += amplitude * noise(coord * frequency);
    frequency *= u.lacunarity;
    amplitude *= u.persistence;
  }

  return value;
}

// --- Domain warping ---
// Distorts UV coordinates using a second noise evaluation

fn domain_warp(p: vec2f) -> vec2f {
  let warp_x = fbm(p + vec2f(0.0, 0.0));
  let warp_y = fbm(p + vec2f(5.2, 1.3));
  return p + vec2f(warp_x, warp_y) * u.warp_strength;
}

@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, offset, and time-based animation
  var p = vec2f(f32(id.x), f32(id.y)) / vec2f(u.width, u.height);
  p = p * u.scale + vec2f(u.offset_x, u.offset_y) + vec2f(u.time * 0.1, u.time * 0.07);

  if (u.domain_warp == 1u) {
    p = domain_warp(p);
  }

  // Map noise from [-1, 1] to [0, 1]
  let value = fbm(p) * 0.5 + 0.5;

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

// Perlin / Simplex noise module
// Generates animated 2D noise on the GPU and writes it to a caller-provided GPUTexture.
//
// Default WGSL loading uses a ?raw import (works with Vite, esbuild, Webpack).
// Alternative: load via fetch — see README.md for details.
import shaderSource from './noise.wgsl?raw';

export type NoiseType = 'perlin' | 'simplex';

export interface NoiseUpdateOptions {
  time?: number;
  scale?: number;
  offset?: [number, number];
  octaves?: number;
  persistence?: number;
  lacunarity?: number;
  type?: NoiseType;
  domainWarp?: boolean;
  warpStrength?: number;
}

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

// Uniform buffer layout (std140-aligned):
// f32 time, f32 scale, f32 offset_x, f32 offset_y,
// u32 octaves, u32 noise_type, u32 domain_warp, f32 warp_strength,
// f32 persistence, f32 lacunarity, f32 width, f32 height
const UNIFORM_SIZE = 48; // 12 x 4 bytes

export function createNoise(device: GPUDevice): Noise {
  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: NoiseUpdateOptions = {}): 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;

    f32[0] = opts.time ?? 0;
    f32[1] = opts.scale ?? 2.0;
    f32[2] = opts.offset?.[0] ?? 0;
    f32[3] = opts.offset?.[1] ?? 0;
    u32[4] = opts.octaves ?? 6;
    u32[5] = (opts.type ?? 'simplex') === 'perlin' ? 0 : 1;
    u32[6] = opts.domainWarp ? 1 : 0;
    f32[7] = opts.warpStrength ?? 1.0;
    f32[8] = opts.persistence ?? 0.5;
    f32[9] = opts.lacunarity ?? 2.0;
    f32[10] = width;
    f32[11] = height;

    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 };
}