/**
 * Custom WebGPU crystal/gem particle animation for Gemma 4 intro.
 * Uses compute shader for particle simulation + render pipeline for visualization.
 */

const PARTICLE_COUNT = 2000;

const COMPUTE_SHADER = /* wgsl */ `
struct Particle {
  pos: vec2f,
  vel: vec2f,
  life: f32,
  size: f32,
  hue: f32,
  _pad: f32,
}

struct Uniforms {
  time: f32,
  delta: f32,
  width: f32,
  height: f32,
  mouseX: f32,
  mouseY: f32,
  _pad0: f32,
  _pad1: f32,
}

@group(0) @binding(0) var<storage, read_write> particles: array<Particle>;
@group(0) @binding(1) var<uniform> uniforms: Uniforms;

fn hash(p: f32) -> f32 {
  var s = fract(p * 0.1031);
  s = s * (s + 33.33);
  s = s * (s + s);
  return fract(s);
}

@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) id: vec3u) {
  let i = id.x;
  if (i >= ${PARTICLE_COUNT}u) { return; }

  var p = particles[i];
  let t = uniforms.time;
  let dt = uniforms.delta;

  // Update life
  p.life -= dt * 0.15;

  if (p.life <= 0.0) {
    // Respawn from center with random direction
    let seed = f32(i) + t * 7.31;
    let angle = hash(seed) * 6.283185;
    let speed = hash(seed + 1.0) * 0.3 + 0.1;
    p.pos = vec2f(0.5, 0.5);
    p.vel = vec2f(cos(angle) * speed, sin(angle) * speed);
    p.life = hash(seed + 2.0) * 0.8 + 0.4;
    p.size = hash(seed + 3.0) * 3.0 + 1.0;
    p.hue = hash(seed + 4.0);
  }

  // Gentle attraction to mouse
  let mousePos = vec2f(uniforms.mouseX / uniforms.width, uniforms.mouseY / uniforms.height);
  let toMouse = mousePos - p.pos;
  let dist = length(toMouse);
  if (dist > 0.01) {
    p.vel += normalize(toMouse) * 0.02 * dt / (dist + 0.2);
  }

  // Swirl effect
  let center = vec2f(0.5, 0.5);
  let toCenter = center - p.pos;
  let perpendicular = vec2f(-toCenter.y, toCenter.x);
  p.vel += perpendicular * 0.05 * dt;

  // Damping
  p.vel *= 0.995;

  // Update position
  p.pos += p.vel * dt;

  // Soft boundary bounce
  if (p.pos.x < 0.0 || p.pos.x > 1.0) { p.vel.x *= -0.5; p.pos.x = clamp(p.pos.x, 0.0, 1.0); }
  if (p.pos.y < 0.0 || p.pos.y > 1.0) { p.vel.y *= -0.5; p.pos.y = clamp(p.pos.y, 0.0, 1.0); }

  particles[i] = p;
}
`;

const VERTEX_SHADER = /* wgsl */ `
struct Particle {
  pos: vec2f,
  vel: vec2f,
  life: f32,
  size: f32,
  hue: f32,
  _pad: f32,
}

struct Uniforms {
  time: f32,
  delta: f32,
  width: f32,
  height: f32,
  mouseX: f32,
  mouseY: f32,
  _pad0: f32,
  _pad1: f32,
}

struct VSOut {
  @builtin(position) pos: vec4f,
  @location(0) alpha: f32,
  @location(1) hue: f32,
  @location(2) uv: vec2f,
}

@group(0) @binding(0) var<storage, read> particles: array<Particle>;
@group(0) @binding(1) var<uniform> uniforms: Uniforms;

// Quad vertices for instancing
const quadVerts = array<vec2f, 6>(
  vec2f(-1, -1), vec2f(1, -1), vec2f(-1, 1),
  vec2f(-1, 1), vec2f(1, -1), vec2f(1, 1),
);

@vertex
fn main(
  @builtin(vertex_index) vi: u32,
  @builtin(instance_index) ii: u32,
) -> VSOut {
  let p = particles[ii];
  let qv = quadVerts[vi];

  let aspect = uniforms.width / uniforms.height;
  let size = p.size * 0.004 * (0.3 + p.life * 0.7);

  var clipPos = vec2f(
    (p.pos.x * 2.0 - 1.0) + qv.x * size / aspect,
    (1.0 - p.pos.y * 2.0) + qv.y * size,
  );

  var out: VSOut;
  out.pos = vec4f(clipPos, 0.0, 1.0);
  out.alpha = p.life;
  out.hue = p.hue;
  out.uv = qv * 0.5 + 0.5;
  return out;
}
`;

const FRAGMENT_SHADER = /* wgsl */ `
struct FSIn {
  @location(0) alpha: f32,
  @location(1) hue: f32,
  @location(2) uv: vec2f,
}

fn hsl2rgb(h: f32, s: f32, l: f32) -> vec3f {
  let c = (1.0 - abs(2.0 * l - 1.0)) * s;
  let x = c * (1.0 - abs(((h * 6.0) % 2.0) - 1.0));
  let m = l - c * 0.5;
  var rgb: vec3f;
  let h6 = h * 6.0;
  if (h6 < 1.0) { rgb = vec3f(c, x, 0.0); }
  else if (h6 < 2.0) { rgb = vec3f(x, c, 0.0); }
  else if (h6 < 3.0) { rgb = vec3f(0.0, c, x); }
  else if (h6 < 4.0) { rgb = vec3f(0.0, x, c); }
  else if (h6 < 5.0) { rgb = vec3f(x, 0.0, c); }
  else { rgb = vec3f(c, 0.0, x); }
  return rgb + m;
}

@fragment
fn main(in: FSIn) -> @location(0) vec4f {
  // Soft circle
  let d = length(in.uv - 0.5) * 2.0;
  if (d > 1.0) { discard; }
  let softEdge = 1.0 - smoothstep(0.4, 1.0, d);

  // Gemma brand hue range: blue (0.6) → purple (0.75) → teal (0.47)
  let h = mix(0.55, 0.78, in.hue);
  let color = hsl2rgb(h, 0.85, 0.65);

  let alpha = softEdge * in.alpha * 0.7;
  return vec4f(color * alpha, alpha);
}
`;

export interface GemmaParticleSystem {
  render: () => void;
  updateMouse: (x: number, y: number) => void;
  resize: (w: number, h: number) => void;
  destroy: () => void;
}

export async function createGemmaParticles(
  canvas: HTMLCanvasElement,
): Promise<GemmaParticleSystem | null> {
  if (!navigator.gpu) return null;

  const adapter = await navigator.gpu.requestAdapter();
  if (!adapter) return null;

  const device = await adapter.requestDevice();
  const ctx = canvas.getContext("webgpu") as GPUCanvasContext | null;
  if (!ctx) return null;

  const format = navigator.gpu.getPreferredCanvasFormat();
  ctx.configure({ device, format, alphaMode: "premultiplied" });

  // Particle buffer
  const particleData = new Float32Array(PARTICLE_COUNT * 8); // 8 floats per particle
  for (let i = 0; i < PARTICLE_COUNT; i++) {
    const off = i * 8;
    particleData[off + 0] = Math.random(); // pos.x
    particleData[off + 1] = Math.random(); // pos.y
    const angle = Math.random() * Math.PI * 2;
    const speed = Math.random() * 0.2 + 0.05;
    particleData[off + 2] = Math.cos(angle) * speed; // vel.x
    particleData[off + 3] = Math.sin(angle) * speed; // vel.y
    particleData[off + 4] = Math.random() * 0.8 + 0.2; // life
    particleData[off + 5] = Math.random() * 3 + 1; // size
    particleData[off + 6] = Math.random(); // hue
    particleData[off + 7] = 0; // pad
  }

  const particleBuffer = device.createBuffer({
    size: particleData.byteLength,
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
  });
  device.queue.writeBuffer(particleBuffer, 0, particleData);

  // Uniforms buffer: time, delta, width, height, mouseX, mouseY, pad, pad
  const uniformData = new Float32Array([0, 0.016, canvas.width, canvas.height, canvas.width / 2, canvas.height / 2, 0, 0]);
  const uniformBuffer = device.createBuffer({
    size: uniformData.byteLength,
    usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
  });
  device.queue.writeBuffer(uniformBuffer, 0, uniformData);

  // Compute pipeline
  const computePipeline = device.createComputePipeline({
    layout: "auto",
    compute: {
      module: device.createShaderModule({ code: COMPUTE_SHADER }),
      entryPoint: "main",
    },
  });

  const computeBindGroup = device.createBindGroup({
    layout: computePipeline.getBindGroupLayout(0),
    entries: [
      { binding: 0, resource: { buffer: particleBuffer } },
      { binding: 1, resource: { buffer: uniformBuffer } },
    ],
  });

  // Render pipeline
  const renderPipeline = device.createRenderPipeline({
    layout: "auto",
    vertex: {
      module: device.createShaderModule({ code: VERTEX_SHADER }),
      entryPoint: "main",
    },
    fragment: {
      module: device.createShaderModule({ code: FRAGMENT_SHADER }),
      entryPoint: "main",
      targets: [
        {
          format,
          blend: {
            color: { srcFactor: "one", dstFactor: "one-minus-src-alpha" },
            alpha: { srcFactor: "one", dstFactor: "one-minus-src-alpha" },
          },
        },
      ],
    },
    primitive: { topology: "triangle-list" },
  });

  const renderBindGroup = device.createBindGroup({
    layout: renderPipeline.getBindGroupLayout(0),
    entries: [
      { binding: 0, resource: { buffer: particleBuffer } },
      { binding: 1, resource: { buffer: uniformBuffer } },
    ],
  });

  let mouseX = canvas.width / 2;
  let mouseY = canvas.height / 2;
  let time = 0;
  let lastTime = performance.now();
  let destroyed = false;

  function render() {
    if (destroyed) return;

    const now = performance.now();
    const dt = Math.min((now - lastTime) / 1000, 0.05);
    lastTime = now;
    time += dt;

    uniformData[0] = time;
    uniformData[1] = dt;
    uniformData[2] = canvas.width;
    uniformData[3] = canvas.height;
    uniformData[4] = mouseX;
    uniformData[5] = mouseY;
    device.queue.writeBuffer(uniformBuffer, 0, uniformData);

    const encoder = device.createCommandEncoder();

    // Compute pass
    const computePass = encoder.beginComputePass();
    computePass.setPipeline(computePipeline);
    computePass.setBindGroup(0, computeBindGroup);
    computePass.dispatchWorkgroups(Math.ceil(PARTICLE_COUNT / 64));
    computePass.end();

    // Render pass
    const textureView = ctx!.getCurrentTexture().createView();
    const renderPass = encoder.beginRenderPass({
      colorAttachments: [
        {
          view: textureView,
          clearValue: { r: 0.039, g: 0.039, b: 0.102, a: 1 },
          loadOp: "clear",
          storeOp: "store",
        },
      ],
    });
    renderPass.setPipeline(renderPipeline);
    renderPass.setBindGroup(0, renderBindGroup);
    renderPass.draw(6, PARTICLE_COUNT);
    renderPass.end();

    device.queue.submit([encoder.finish()]);
    requestAnimationFrame(render);
  }

  return {
    render() {
      requestAnimationFrame(render);
    },
    updateMouse(x: number, y: number) {
      mouseX = x;
      mouseY = y;
    },
    resize(w: number, h: number) {
      canvas.width = w;
      canvas.height = h;
      ctx!.configure({ device, format, alphaMode: "premultiplied" });
    },
    destroy() {
      destroyed = true;
      device.destroy();
    },
  };
}