import {Constants, Engine, MeshBuilder, RenderTargetTexture, Scene, ShaderMaterial, Texture, Vector2, WebGPUEngine} from '@babylonjs/core';
import {PENDULUM_DISPLAY_FRAGMENT_GLSL, PENDULUM_FEEDBACK_FRAGMENT_GLSL, PENDULUM_VERTEX_SHADER_GLSL} from '../pendulum-glsl.shader';
import {PendulumState, stepPendulumPhysics} from './pendulum-cpu-physics';
import {PendulumSimParams, PendulumSimulationStrategy} from './pendulum-simulation.strategy';

export class PendulumCpuStrategy implements PendulumSimulationStrategy {
  private scene: Scene | null = null;
  private state: PendulumState = {theta1: Math.PI / 4, theta2: Math.PI / 2, v1: 0, v2: 0};
  private params: PendulumSimParams | null = null;
  private feedbackMaterial: ShaderMaterial | null = null;
  private displayMaterial: ShaderMaterial | null = null;
  private rttA: RenderTargetTexture | null = null;
  private rttB: RenderTargetTexture | null = null;
  private pingPong = false;

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

    const existingPlane = scene.getMeshByName("plane");
    if (existingPlane) {
      scene.removeMesh(existingPlane);
    }

    const width = engine.getRenderWidth();
    const height = engine.getRenderHeight();

    this.rttA = new RenderTargetTexture("rttA", {width, height}, scene, false, true, Constants.TEXTURETYPE_FLOAT);
    this.rttB = new RenderTargetTexture("rttB", {width, height}, scene, false, true, Constants.TEXTURETYPE_FLOAT);

    this.rttA.wrapU = Texture.CLAMP_ADDRESSMODE;
    this.rttA.wrapV = Texture.CLAMP_ADDRESSMODE;
    this.rttB.wrapU = Texture.CLAMP_ADDRESSMODE;
    this.rttB.wrapV = Texture.CLAMP_ADDRESSMODE;

    // Size 2 maps positions to -1..1 which fills clip space (vertex shader bypasses camera)
    const feedbackPlane = MeshBuilder.CreatePlane("feedbackPlane", {size: 2}, scene);
    feedbackPlane.alwaysSelectAsActiveMesh = true;
    // Use layer mask to exclude from main scene rendering (only rendered via RTT)
    feedbackPlane.layerMask = 0x10000000;

    this.feedbackMaterial = new ShaderMaterial("feedbackMat", scene, {
      vertexSource: PENDULUM_VERTEX_SHADER_GLSL,
      fragmentSource: PENDULUM_FEEDBACK_FRAGMENT_GLSL
    }, {
      attributes: ["position", "uv"],
      uniforms: ["resolution", "theta1", "theta2", "l1", "l2", "trailDecay"],
      samplers: ["previousFrame"]
    });
    this.feedbackMaterial.backFaceCulling = false;
    feedbackPlane.material = this.feedbackMaterial;

    this.rttA.renderList!.push(feedbackPlane);
    this.rttB.renderList!.push(feedbackPlane);

    // Size 2 maps positions to -1..1 which fills clip space (vertex shader bypasses camera)
    const displayPlane = MeshBuilder.CreatePlane("displayPlane", {size: 2}, scene);
    displayPlane.alwaysSelectAsActiveMesh = true;

    this.displayMaterial = new ShaderMaterial("displayMat", scene, {
      vertexSource: PENDULUM_VERTEX_SHADER_GLSL,
      fragmentSource: PENDULUM_DISPLAY_FRAGMENT_GLSL
    }, {
      attributes: ["position", "uv"],
      uniforms: ["resolution", "theta1", "theta2", "l1", "l2"],
      samplers: ["trailTexture"]
    });
    this.displayMaterial.backFaceCulling = false;
    this.displayMaterial.disableDepthWrite = true;
    displayPlane.material = this.displayMaterial;

    // RTTs are rendered manually via writeTarget.render() -- do NOT add to customRenderTargets
    this.startRenderLoop(scene, width, height);
  }

  updateParams(params: PendulumSimParams): void {
    this.params = params;
  }

  dispose(): void {
    if (this.scene) {
      this.scene.onBeforeRenderObservable.clear();
    }
    this.rttA?.dispose();
    this.rttB?.dispose();
    this.feedbackMaterial?.dispose();
    this.displayMaterial?.dispose();
    this.scene = null;
  }

  private startRenderLoop(scene: Scene, width: number, height: number): void {
    const resolution = new Vector2(width, height);

    scene.onBeforeRenderObservable.clear();
    scene.onBeforeRenderObservable.add(() => {
      if (!this.params || !this.feedbackMaterial || !this.displayMaterial || !this.rttA || !this.rttB) {
        return;
      }

      stepPendulumPhysics(this.state, {
        dt: this.params.dt,
        g: this.params.g,
        m1: this.params.m1,
        m2: this.params.m2,
        l1: this.params.l1,
        l2: this.params.l2,
        damping: this.params.damping,
        impulseM1: this.params.impulseM1,
        impulseM2: this.params.impulseM2
      });

      const readTarget = this.pingPong ? this.rttA : this.rttB;
      const writeTarget = this.pingPong ? this.rttB : this.rttA;

      this.feedbackMaterial.setTexture("previousFrame", readTarget);
      this.feedbackMaterial.setVector2("resolution", resolution);
      this.feedbackMaterial.setFloat("theta1", this.state.theta1);
      this.feedbackMaterial.setFloat("theta2", this.state.theta2);
      this.feedbackMaterial.setFloat("l1", this.params.l1);
      this.feedbackMaterial.setFloat("l2", this.params.l2);
      this.feedbackMaterial.setFloat("trailDecay", this.params.trailDecay);

      writeTarget.render();

      this.displayMaterial.setTexture("trailTexture", writeTarget);
      this.displayMaterial.setVector2("resolution", resolution);
      this.displayMaterial.setFloat("theta1", this.state.theta1);
      this.displayMaterial.setFloat("theta2", this.state.theta2);
      this.displayMaterial.setFloat("l1", this.params.l1);
      this.displayMaterial.setFloat("l2", this.params.l2);

      this.pingPong = !this.pingPong;
    });
  }
}
