// Verlet position integration
// Computes new positions using: new_pos = pos + (pos - old_pos) * damping + acceleration * dt^2
// Pinned points are skipped (their positions remain fixed).

struct Uniforms {
  dt: f32,
  damping: f32,
  gravity_x: f32,
  gravity_y: f32,
  point_count: u32,
  bounds_width: f32,
  bounds_height: f32,
  _pad: u32,
}

@group(0) @binding(0) var<uniform> u: Uniforms;
@group(0) @binding(1) var<storage, read_write> positions: array<vec2f>;
@group(0) @binding(2) var<storage, read_write> old_positions: array<vec2f>;
@group(0) @binding(3) var<storage, read> pins: array<u32>;

@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3u) {
  let idx = gid.x;
  if (idx >= u.point_count) {
    return;
  }

  // Pinned points don't move — keep old_pos in sync so they have zero velocity
  if (pins[idx] == 1u) {
    old_positions[idx] = positions[idx];
    return;
  }

  let pos = positions[idx];
  let old_pos = old_positions[idx];

  // Verlet integration: velocity is implicit as (pos - old_pos)
  let velocity = (pos - old_pos) * u.damping;
  let acceleration = vec2f(u.gravity_x, u.gravity_y);
  let dt_sq = u.dt * u.dt;

  var new_pos = pos + velocity + acceleration * dt_sq;

  // Boundary constraints — keep points within the simulation area.
  // Bounds are centered at origin.
  let hw = u.bounds_width * 0.5;
  let hh = u.bounds_height * 0.5;

  if (hw > 0.0 && hh > 0.0) {
    new_pos.x = clamp(new_pos.x, -hw, hw);
    new_pos.y = clamp(new_pos.y, -hh, hh);
  }

  old_positions[idx] = pos;
  positions[idx] = new_pos;
}

// Distance constraint relaxation for Verlet integration
// Each constraint connects two points with a rest length. The shader moves both
// points equally toward or away from each other to satisfy the constraint.
// Pinned points are not moved — the full correction is applied to the other point.
//
// Run this shader multiple iterations per frame for stable constraint solving.

struct Uniforms {
  constraint_count: u32,
  offset: u32,   // start index (0 for even pass, 1 for odd pass)
  stride: u32,   // step between constraints (2 for even/odd splitting)
  _pad: u32,
}

struct Constraint {
  index_a: u32,
  index_b: u32,
  rest_length: f32,
  _pad: u32,
}

@group(0) @binding(0) var<uniform> u: Uniforms;
@group(0) @binding(1) var<storage, read_write> positions: array<vec2f>;
@group(0) @binding(2) var<storage, read> constraints: array<Constraint>;
@group(0) @binding(3) var<storage, read> pins: array<u32>;

@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3u) {
  let actual_idx = u.offset + gid.x * u.stride;
  if (actual_idx >= u.constraint_count) {
    return;
  }

  let c = constraints[actual_idx];
  let pos_a = positions[c.index_a];
  let pos_b = positions[c.index_b];

  let delta = pos_b - pos_a;
  let dist = length(delta);

  // Avoid division by zero for overlapping points
  if (dist < 0.0001) {
    return;
  }

  // How far off the rest length we are
  let error = (dist - c.rest_length) / dist;
  let correction = delta * error;

  let pin_a = pins[c.index_a];
  let pin_b = pins[c.index_b];

  // Distribute correction based on pin state:
  // - Both free: split correction 50/50
  // - One pinned: full correction to the free point
  // - Both pinned: no correction
  if (pin_a == 0u && pin_b == 0u) {
    positions[c.index_a] = pos_a + correction * 0.5;
    positions[c.index_b] = pos_b - correction * 0.5;
  } else if (pin_a == 0u) {
    positions[c.index_a] = pos_a + correction;
  } else if (pin_b == 0u) {
    positions[c.index_b] = pos_b - correction;
  }
}

// Verlet integration for ropes, cloth, and soft bodies
// Two compute shaders handle position integration and distance constraint solving.
// Points are 2D vec2<f32>. Constraints connect pairs of points with a rest length.
//
// Default WGSL loading uses a ?raw import (works with Vite, esbuild, Webpack).
// Alternative: load via fetch — see README.md for details.
import integrateSource from './integrate.wgsl?raw';
import constraintsSource from './constraints.wgsl?raw';

export interface VerletConstraint {
  a: number;
  b: number;
  restLength: number;
}

export interface VerletOptions {
  points: number;
  constraints?: VerletConstraint[];
  gravity?: [number, number];
  damping?: number;
  bounds?: { width: number; height: number };
  pins?: number[];
  constraintIterations?: number;
}

export interface Verlet {
  positions: GPUBuffer;
  count: number;
  update(deltaTime: number): void;
  destroy(): void;
}

const INTEGRATE_UNIFORM_SIZE = 32; // 8 x 4 bytes
const CONSTRAINT_UNIFORM_SIZE = 16; // 4 x 4 bytes
const WORKGROUP_SIZE = 64;

// Default rope configuration: chain of N points connected sequentially, first point pinned
function defaultRopeConfig(pointCount: number): {
  constraints: VerletConstraint[];
  pins: number[];
  initialPositions: Float32Array;
} {
  const constraints: VerletConstraint[] = [];
  const restLength = 1.0;

  for (let i = 0; i < pointCount - 1; i++) {
    constraints.push({ a: i, b: i + 1, restLength });
  }

  const pins = [0]; // Pin the first point

  // Lay out points horizontally, starting from origin
  const initialPositions = new Float32Array(pointCount * 2);
  for (let i = 0; i < pointCount; i++) {
    initialPositions[i * 2] = i * restLength;
    initialPositions[i * 2 + 1] = 0;
  }

  return { constraints, pins, initialPositions };
}

export function createVerlet(device: GPUDevice, options: VerletOptions): Verlet {
  const pointCount = options.points;
  const gravity = options.gravity ?? [0, 9.8];
  const damping = options.damping ?? 0.99;
  const boundsWidth = options.bounds?.width ?? 0;
  const boundsHeight = options.bounds?.height ?? 0;
  const constraintIterations = options.constraintIterations ?? 8;

  // Use explicit constraints or default to a rope
  const ropeDefaults = defaultRopeConfig(pointCount);
  const constraints = options.constraints ?? ropeDefaults.constraints;
  const pinIndices = options.pins ?? ropeDefaults.pins;

  const posBufferSize = pointCount * 2 * 4; // count * vec2f * sizeof(f32)
  const constraintCount = constraints.length;

  // Initialize positions
  const initialPositions = options.constraints
    ? new Float32Array(pointCount * 2)
    : ropeDefaults.initialPositions;

  // Build pin mask (1 = pinned, 0 = free)
  const pinData = new Uint32Array(pointCount);
  for (const idx of pinIndices) {
    if (idx >= 0 && idx < pointCount) {
      pinData[idx] = 1;
    }
  }

  // Build constraint buffer data: { indexA: u32, indexB: u32, restLength: f32, _pad: u32 }
  const constraintData = new ArrayBuffer(constraintCount * 16);
  const constraintU32 = new Uint32Array(constraintData);
  const constraintF32 = new Float32Array(constraintData);
  for (let i = 0; i < constraintCount; i++) {
    const c = constraints[i];
    constraintU32[i * 4] = c.a;
    constraintU32[i * 4 + 1] = c.b;
    constraintF32[i * 4 + 2] = c.restLength;
    constraintU32[i * 4 + 3] = 0; // padding
  }

  // GPU Buffers
  const positionsBuffer = device.createBuffer({
    size: posBufferSize,
    usage:
      GPUBufferUsage.STORAGE |
      GPUBufferUsage.COPY_DST |
      GPUBufferUsage.VERTEX |
      GPUBufferUsage.COPY_SRC,
    mappedAtCreation: true
  });
  new Float32Array(positionsBuffer.getMappedRange()).set(initialPositions);
  positionsBuffer.unmap();

  const oldPositionsBuffer = device.createBuffer({
    size: posBufferSize,
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
    mappedAtCreation: true
  });
  new Float32Array(oldPositionsBuffer.getMappedRange()).set(initialPositions);
  oldPositionsBuffer.unmap();

  const pinsBuffer = device.createBuffer({
    size: pointCount * 4,
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
    mappedAtCreation: true
  });
  new Uint32Array(pinsBuffer.getMappedRange()).set(pinData);
  pinsBuffer.unmap();

  const constraintBuffer = device.createBuffer({
    size: Math.max(constraintCount * 16, 16), // minimum 16 bytes for empty constraint list
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
    mappedAtCreation: true
  });
  if (constraintCount > 0) {
    new Uint8Array(constraintBuffer.getMappedRange()).set(new Uint8Array(constraintData));
  }
  constraintBuffer.unmap();

  const integrateUniformBuffer = device.createBuffer({
    size: INTEGRATE_UNIFORM_SIZE,
    usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST
  });

  // Two uniform buffers for even/odd constraint passes
  const evenConstraintUniformBuffer = device.createBuffer({
    size: CONSTRAINT_UNIFORM_SIZE,
    usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST
  });
  const oddConstraintUniformBuffer = device.createBuffer({
    size: CONSTRAINT_UNIFORM_SIZE,
    usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST
  });

  // Even pass: offset=0, stride=2 — processes constraints 0, 2, 4, ...
  // Odd pass:  offset=1, stride=2 — processes constraints 1, 3, 5, ...
  const evenCount = Math.ceil(constraintCount / 2);
  const oddCount = Math.floor(constraintCount / 2);
  device.queue.writeBuffer(
    evenConstraintUniformBuffer,
    0,
    new Uint32Array([constraintCount, 0, 2, 0])
  );
  device.queue.writeBuffer(
    oddConstraintUniformBuffer,
    0,
    new Uint32Array([constraintCount, 1, 2, 0])
  );

  // Shader modules
  const integrateModule = device.createShaderModule({ code: integrateSource });
  const constraintsModule = device.createShaderModule({ code: constraintsSource });

  // Integration pipeline
  const integrateLayout = device.createBindGroupLayout({
    entries: [
      { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'uniform' } },
      { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
      { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
      { binding: 3, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'read-only-storage' } }
    ]
  });

  const integratePipeline = device.createComputePipeline({
    layout: device.createPipelineLayout({ bindGroupLayouts: [integrateLayout] }),
    compute: { module: integrateModule, entryPoint: 'main' }
  });

  const integrateBindGroup = device.createBindGroup({
    layout: integrateLayout,
    entries: [
      { binding: 0, resource: { buffer: integrateUniformBuffer } },
      { binding: 1, resource: { buffer: positionsBuffer } },
      { binding: 2, resource: { buffer: oldPositionsBuffer } },
      { binding: 3, resource: { buffer: pinsBuffer } }
    ]
  });

  // Constraint pipeline
  const constraintLayout = device.createBindGroupLayout({
    entries: [
      { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'uniform' } },
      { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
      { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'read-only-storage' } },
      { binding: 3, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'read-only-storage' } }
    ]
  });

  const constraintPipeline = device.createComputePipeline({
    layout: device.createPipelineLayout({ bindGroupLayouts: [constraintLayout] }),
    compute: { module: constraintsModule, entryPoint: 'main' }
  });

  const evenConstraintBG = device.createBindGroup({
    layout: constraintLayout,
    entries: [
      { binding: 0, resource: { buffer: evenConstraintUniformBuffer } },
      { binding: 1, resource: { buffer: positionsBuffer } },
      { binding: 2, resource: { buffer: constraintBuffer } },
      { binding: 3, resource: { buffer: pinsBuffer } }
    ]
  });

  const oddConstraintBG = device.createBindGroup({
    layout: constraintLayout,
    entries: [
      { binding: 0, resource: { buffer: oddConstraintUniformBuffer } },
      { binding: 1, resource: { buffer: positionsBuffer } },
      { binding: 2, resource: { buffer: constraintBuffer } },
      { binding: 3, resource: { buffer: pinsBuffer } }
    ]
  });

  const integrateUniformData = new ArrayBuffer(INTEGRATE_UNIFORM_SIZE);
  const intF32 = new Float32Array(integrateUniformData);
  const intU32 = new Uint32Array(integrateUniformData);

  function update(deltaTime: number): void {
    // Clamp dt to prevent instability from large time steps
    const dt = Math.min(deltaTime, 1 / 30);

    intF32[0] = dt;
    intF32[1] = damping;
    intF32[2] = gravity[0];
    intF32[3] = gravity[1];
    intU32[4] = pointCount;
    intF32[5] = boundsWidth;
    intF32[6] = boundsHeight;
    intU32[7] = 0; // padding

    device.queue.writeBuffer(integrateUniformBuffer, 0, integrateUniformData);

    const encoder = device.createCommandEncoder();

    // Step 1: Verlet integration (apply velocity + gravity)
    const integratePass = encoder.beginComputePass();
    integratePass.setPipeline(integratePipeline);
    integratePass.setBindGroup(0, integrateBindGroup);
    integratePass.dispatchWorkgroups(Math.ceil(pointCount / WORKGROUP_SIZE));
    integratePass.end();

    // Step 2: Constraint relaxation — alternate even/odd passes to avoid race conditions
    if (constraintCount > 0) {
      for (let i = 0; i < constraintIterations; i++) {
        // Even constraints (0, 2, 4, ...)
        const evenPass = encoder.beginComputePass();
        evenPass.setPipeline(constraintPipeline);
        evenPass.setBindGroup(0, evenConstraintBG);
        evenPass.dispatchWorkgroups(Math.ceil(evenCount / WORKGROUP_SIZE));
        evenPass.end();

        // Odd constraints (1, 3, 5, ...)
        if (oddCount > 0) {
          const oddPass = encoder.beginComputePass();
          oddPass.setPipeline(constraintPipeline);
          oddPass.setBindGroup(0, oddConstraintBG);
          oddPass.dispatchWorkgroups(Math.ceil(oddCount / WORKGROUP_SIZE));
          oddPass.end();
        }
      }
    }

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

  function destroy(): void {
    positionsBuffer.destroy();
    oldPositionsBuffer.destroy();
    pinsBuffer.destroy();
    constraintBuffer.destroy();
    integrateUniformBuffer.destroy();
    evenConstraintUniformBuffer.destroy();
    oddConstraintUniformBuffer.destroy();
  }

  return {
    positions: positionsBuffer,
    count: pointCount,
    update,
    destroy
  };
}