import {Camera, ComputeShader, Engine, MeshBuilder, Scene, ShaderLanguage, ShaderMaterial, StorageBuffer, WebGPUEngine} from '@babylonjs/core';
import {PENDULUM_FRAGMENT_SHADER_WGSL, PENDULUM_PHYSIC_COMPUTE_SHADER_WGSL, PENDULUM_RENDER_COMPUTE_SHADER_WGSL, PENDULUM_VERTEX_SHADER_WGSL} from '../pendulum.shader';
import {PendulumSimParams, PendulumSimulationStrategy} from './pendulum-simulation.strategy';

export class PendulumGpuStrategy implements PendulumSimulationStrategy {
  private scene: Scene | null = null;
  private paramsData = new Float32Array(14);
  private paramsBuffer: StorageBuffer | null = null;

  init(scene: Scene, engine: WebGPUEngine | Engine): void {
    this.scene = scene;
    const gpuEngine = engine as WebGPUEngine;
    gpuEngine.resize();

    const width = gpuEngine.getRenderWidth();
    const height = gpuEngine.getRenderHeight();
    const totalPixels = width * height;

    const pixelBuffer = new StorageBuffer(gpuEngine, totalPixels * 4);

    const stateBuffer = new StorageBuffer(gpuEngine, 4 * 4);
    stateBuffer.update(new Float32Array([Math.PI / 4, Math.PI / 2, 0, 0]));

    this.paramsBuffer = new StorageBuffer(gpuEngine, 14 * 4);

    const csPhysics = new ComputeShader("physics", gpuEngine,
      {computeSource: PENDULUM_PHYSIC_COMPUTE_SHADER_WGSL},
      {bindingsMapping: {"state": {group: 0, binding: 0}, "p": {group: 0, binding: 1}}}
    );
    csPhysics.setStorageBuffer("state", stateBuffer);
    csPhysics.setStorageBuffer("p", this.paramsBuffer);

    const csRender = new ComputeShader("render", gpuEngine,
      {computeSource: PENDULUM_RENDER_COMPUTE_SHADER_WGSL},
      {bindingsMapping: {"pixelBuffer": {group: 0, binding: 0}, "p": {group: 0, binding: 1}, "state": {group: 0, binding: 2}}}
    );
    csRender.setStorageBuffer("pixelBuffer", pixelBuffer);
    csRender.setStorageBuffer("p", this.paramsBuffer);
    csRender.setStorageBuffer("state", stateBuffer);

    // Create WGSL display plane and material
    const displayMaterial = new ShaderMaterial("pendulumWgslMat", scene, {
      vertexSource: PENDULUM_VERTEX_SHADER_WGSL,
      fragmentSource: PENDULUM_FRAGMENT_SHADER_WGSL
    }, {
      attributes: ["position"],
      uniforms: [],
      storageBuffers: ["pixelBuffer", "p"],
      shaderLanguage: ShaderLanguage.WGSL
    });
    displayMaterial.setStorageBuffer("pixelBuffer", pixelBuffer);
    displayMaterial.setStorageBuffer("p", this.paramsBuffer);
    displayMaterial.backFaceCulling = false;
    displayMaterial.disableDepthWrite = true;

    const plane = MeshBuilder.CreatePlane("pendulumPlane", {size: 100}, scene);
    const camera = scene.activeCamera as Camera;
    if (camera) {
      plane.lookAt(camera.position);
    }
    plane.alwaysSelectAsActiveMesh = true;
    plane.material = displayMaterial;

    scene.onBeforeRenderObservable.clear();
    scene.onBeforeRenderObservable.add(() => {
      if (!this.paramsBuffer) {
        return;
      }

      const currentWidth = gpuEngine.getRenderWidth();
      const currentHeight = gpuEngine.getRenderHeight();
      this.paramsData[0] = currentWidth;
      this.paramsData[1] = currentHeight;

      this.paramsBuffer.update(this.paramsData);

      csPhysics.dispatch(1, 1, 1);

      const dispatchCount = Math.ceil((currentWidth * currentHeight) / 64);
      csRender.dispatch(dispatchCount, 1, 1);
    });
  }

  updateParams(params: PendulumSimParams): void {
    // paramsData[0] (width) and paramsData[1] (height) are set in the render loop
    this.paramsData[2] = params.time;
    this.paramsData[3] = params.dt;
    this.paramsData[4] = params.g;
    this.paramsData[5] = params.m1;
    this.paramsData[6] = params.m2;
    this.paramsData[7] = params.l1;
    this.paramsData[8] = params.l2;
    this.paramsData[9] = params.damping;
    this.paramsData[10] = params.trailDecay;
    this.paramsData[11] = params.impulseM1;
    this.paramsData[12] = params.impulseM2;
    this.paramsData[13] = 0;
  }

  dispose(): void {
    if (this.scene) {
      this.scene.onBeforeRenderObservable.clear();
      const plane = this.scene.getMeshByName("pendulumPlane");
      if (plane) {
        this.scene.removeMesh(plane);
        plane.dispose();
      }
    }
    this.paramsBuffer = null;
    this.scene = null;
  }
}
