import {ArcRotateCamera, ComputeShader, Engine, GroundMesh, MeshBuilder, Scene, ShaderLanguage, ShaderMaterial, StorageBuffer, WebGPUEngine} from '@babylonjs/core';
import {ClothConfig, ClothData} from '../cloth.model';
import {
  CLOTH_FRAGMENT_SHADER_WGSL,
  CLOTH_INTEGRATE_COMPUTE_WGSL,
  CLOTH_SOLVE_COMPUTE_WGSL,
  CLOTH_VELOCITY_COMPUTE_WGSL,
  CLOTH_VERTEX_SHADER_WGSL
} from '../cloth.shader';
import {ClothSimulationParams, ClothSimulationStrategy} from './cloth-simulation.strategy';

interface GpuBuffers {
  positions: StorageBuffer;
  prevPositions: StorageBuffer;
  velocities: StorageBuffer;
  params: StorageBuffer;
  constraints: StorageBuffer[];
}

interface GpuPipelines {
  integrate: ComputeShader;
  solvers: ComputeShader[];
  velocity: ComputeShader;
}

export class ClothGpuStrategy implements ClothSimulationStrategy {
  private clothMesh: GroundMesh | null = null;
  private scene: Scene | null = null;
  private paramsData = new Float32Array(9);
  private buffers: GpuBuffers | null = null;
  private pipelines: GpuPipelines | null = null;
  private dispatchXConstraints: number[] = [];
  private dispatchXVertices = 0;
  private numVertices = 0;

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

    const clothData = this.generateClothData(config);
    this.buffers = this.createStorageBuffers(gpuEngine, clothData);
    this.pipelines = this.setupComputePipelines(gpuEngine, this.buffers);
    this.setupRenderMesh(scene, config, this.buffers.positions);
    this.setupDispatchSizes(config, this.buffers);
    this.startRenderLoop(scene);
  }

  updateParams(params: ClothSimulationParams): void {
    if (!this.buffers) {
      return;
    }

    const windX = params.isWindActive ? 5.0 : 0.0;
    const windZ = params.isWindActive ? 15.0 : 0.0;

    const softCompliance = 10.0;
    const rigidCompliance = 0.00001;
    const t = (params.stiffness - 1) / 99.0;
    const compliance = softCompliance * Math.pow(rigidCompliance / softCompliance, t);

    this.paramsData[0] = 0.016;
    this.paramsData[1] = -9.81;
    this.paramsData[2] = compliance;
    this.paramsData[3] = this.numVertices;
    this.paramsData[4] = windX;
    this.paramsData[5] = 0.0;
    this.paramsData[6] = windZ;
    this.paramsData[7] = params.simulationTime;
    this.paramsData[8] = params.elongation;
  }

  getMesh(): GroundMesh | null {
    return this.clothMesh;
  }

  dispose(): void {
    if (this.scene && this.clothMesh) {
      this.scene.removeMesh(this.clothMesh);
    }
    this.clothMesh = null;
    this.buffers = null;
    this.pipelines = null;
  }

  private generateClothData(config: ClothConfig): ClothData {
    const positionsData = new Float32Array(config.numVertices * 4);
    const prevPositionsData = new Float32Array(config.numVertices * 4);
    const velocitiesData = new Float32Array(config.numVertices * 4);

    const constraintsP0: number[] = [];
    const constraintsP1: number[] = [];
    const constraintsP2: number[] = [];
    const constraintsP3: number[] = [];

    const addHorizontalConstraint = (arr: number[], a: number, b: number): void => {
      arr.push(a, b, config.spacing, 1.0);
    };
    const addVerticalConstraint = (arr: number[], a: number, b: number): void => {
      arr.push(a, b, config.spacing, 2.0);
    };

    for (let y = 0; y < config.gridHeight; y++) {
      for (let x = 0; x < config.gridWidth; x++) {
        const idx = (y * config.gridWidth + x) * 4;
        positionsData[idx + 0] = (x - config.gridWidth / 2) * config.spacing;
        positionsData[idx + 1] = 5.0 - (y * config.spacing);
        positionsData[idx + 2] = 0.0;
        positionsData[idx + 3] = (y === 0) ? 0.0 : config.particleInvMass;

        prevPositionsData[idx + 0] = positionsData[idx + 0];
        prevPositionsData[idx + 1] = positionsData[idx + 1];
        prevPositionsData[idx + 2] = positionsData[idx + 2];
        prevPositionsData[idx + 3] = positionsData[idx + 3];
      }
    }

    for (let y = 0; y < config.gridHeight; y++) {
      for (let x = 0; x < config.gridWidth - 1; x += 2) { addHorizontalConstraint(constraintsP0, y * config.gridWidth + x, y * config.gridWidth + x + 1); }
      for (let x = 1; x < config.gridWidth - 1; x += 2) { addHorizontalConstraint(constraintsP1, y * config.gridWidth + x, y * config.gridWidth + x + 1); }
    }
    for (let y = 0; y < config.gridHeight - 1; y += 2) {
      for (let x = 0; x < config.gridWidth; x++) { addVerticalConstraint(constraintsP2, y * config.gridWidth + x, (y + 1) * config.gridWidth + x); }
    }
    for (let y = 1; y < config.gridHeight - 1; y += 2) {
      for (let x = 0; x < config.gridWidth; x++) { addVerticalConstraint(constraintsP3, y * config.gridWidth + x, (y + 1) * config.gridWidth + x); }
    }

    const constraintsP4: number[] = [];
    const constraintsP5: number[] = [];
    const constraintsP6: number[] = [];
    const constraintsP7: number[] = [];
    const diagSpacing = config.spacing * Math.SQRT2;
    const addDiagConstraint = (arr: number[], a: number, b: number): void => {
      arr.push(a, b, diagSpacing, 1.0);
    };

    for (let y = 0; y < config.gridHeight - 1; y++) {
      const arr = (y % 2 === 0) ? constraintsP4 : constraintsP5;
      for (let x = 0; x < config.gridWidth - 1; x++) {
        addDiagConstraint(arr, y * config.gridWidth + x, (y + 1) * config.gridWidth + (x + 1));
      }
    }

    for (let y = 0; y < config.gridHeight - 1; y++) {
      const arr = (y % 2 === 0) ? constraintsP6 : constraintsP7;
      for (let x = 0; x < config.gridWidth - 1; x++) {
        addDiagConstraint(arr, y * config.gridWidth + (x + 1), (y + 1) * config.gridWidth + x);
      }
    }

    return {
      positions: positionsData,
      prevPositions: prevPositionsData,
      velocities: velocitiesData,
      constraints: [
        constraintsP0, constraintsP1, constraintsP2, constraintsP3,
        constraintsP4, constraintsP5, constraintsP6, constraintsP7
      ],
      params: new Float32Array(9)
    };
  }

  private createStorageBuffers(engine: WebGPUEngine, data: ClothData): GpuBuffers {
    const createBuffer = (arrayData: Float32Array | number[]): StorageBuffer => {
      const buffer = new StorageBuffer(engine, arrayData.length * 4);
      buffer.update(arrayData instanceof Float32Array ? arrayData : new Float32Array(arrayData));
      return buffer;
    };

    return {
      positions: createBuffer(data.positions),
      prevPositions: createBuffer(data.prevPositions),
      velocities: createBuffer(data.velocities),
      params: createBuffer(data.params),
      constraints: data.constraints.map(cData => createBuffer(cData))
    };
  }

  private setupComputePipelines(engine: WebGPUEngine, buffers: GpuBuffers): GpuPipelines {
    const createBasicShader = (name: string, source: string) => {
      const cs = new ComputeShader(name, engine, {computeSource: source}, {
        bindingsMapping: {
          "p": {group: 0, binding: 0},
          "positions": {group: 0, binding: 1},
          "prev_positions": {group: 0, binding: 2},
          "velocities": {group: 0, binding: 3}
        }
      });
      cs.setStorageBuffer("p", buffers.params);
      cs.setStorageBuffer("positions", buffers.positions);
      cs.setStorageBuffer("prev_positions", buffers.prevPositions);
      cs.setStorageBuffer("velocities", buffers.velocities);
      return cs;
    };

    const createSolverShader = (name: string, constraintBuffer: StorageBuffer) => {
      const cs = new ComputeShader(name, engine, {computeSource: CLOTH_SOLVE_COMPUTE_WGSL}, {
        bindingsMapping: {
          "p": {group: 0, binding: 0},
          "positions": {group: 0, binding: 1},
          "constraints": {group: 0, binding: 2}
        }
      });
      cs.setStorageBuffer("p", buffers.params);
      cs.setStorageBuffer("positions", buffers.positions);
      cs.setStorageBuffer("constraints", constraintBuffer);
      return cs;
    };

    return {
      integrate: createBasicShader("integrate", CLOTH_INTEGRATE_COMPUTE_WGSL),
      solvers: buffers.constraints.map((cBuffer, i) => createSolverShader(`solve${i}`, cBuffer)),
      velocity: createBasicShader("velocity", CLOTH_VELOCITY_COMPUTE_WGSL)
    };
  }

  private setupRenderMesh(scene: Scene, config: ClothConfig, positionsBuffer: StorageBuffer): void {
    if (this.clothMesh) {
      scene.removeMesh(this.clothMesh);
    }

    this.clothMesh = MeshBuilder.CreateGround("cloth", {width: 10, height: 10, subdivisions: config.gridWidth - 1}, scene);

    const clothMaterial = new ShaderMaterial("clothMat", scene, {
      vertexSource: CLOTH_VERTEX_SHADER_WGSL,
      fragmentSource: CLOTH_FRAGMENT_SHADER_WGSL
    }, {
      attributes: ["position", "uv"],
      uniforms: ["viewProjection"],
      storageBuffers: ["positions"],
      shaderLanguage: ShaderLanguage.WGSL
    });

    clothMaterial.backFaceCulling = false;
    clothMaterial.setStorageBuffer("positions", positionsBuffer);
    this.clothMesh.material = clothMaterial;

    const camera = scene.activeCamera as ArcRotateCamera;
    if (camera) {
      camera.alpha = Math.PI / 4;
      camera.beta = Math.PI / 2.5;
      camera.radius = 15;
    }
  }

  private setupDispatchSizes(config: ClothConfig, buffers: GpuBuffers): void {
    const constraintsLength = buffers.constraints.map(b => (b as any)._buffer.capacity / 4 / 4);
    this.dispatchXConstraints = constraintsLength.map(len => Math.ceil(len / 64));
    this.dispatchXVertices = Math.ceil(config.numVertices / 64);
  }

  private startRenderLoop(scene: Scene): void {
    const substeps = 15;

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

      this.buffers.params.update(this.paramsData);

      this.pipelines.integrate.dispatch(this.dispatchXVertices, 1, 1);

      for (let i = 0; i < substeps; i++) {
        for (let phase = 0; phase < this.pipelines.solvers.length; phase++) {
          this.pipelines.solvers[phase].dispatch(this.dispatchXConstraints[phase], 1, 1);
        }
      }

      this.pipelines.velocity.dispatch(this.dispatchXVertices, 1, 1);
    });
  }
}
