import {ArcRotateCamera, Engine, GroundMesh, MeshBuilder, Scene, ShaderMaterial, VertexBuffer, WebGPUEngine} from '@babylonjs/core';
import {ClothConfig, ClothData} from '../cloth.model';
import {CLOTH_FRAGMENT_SHADER_GLSL, CLOTH_VERTEX_SHADER_GLSL} from '../cloth-glsl.shader';
import {ClothPhysicsParams, integratePositions, solveConstraints, updateVelocities} from './cloth-cpu-physics';
import {ClothSimulationParams, ClothSimulationStrategy} from './cloth-simulation.strategy';

export class ClothCpuStrategy implements ClothSimulationStrategy {
  private clothMesh: GroundMesh | null = null;
  private scene: Scene | null = null;
  private config: ClothConfig | null = null;

  private positions!: Float32Array;
  private prevPositions!: Float32Array;
  private velocities!: Float32Array;
  private constraintPhases!: Float32Array[];

  private physicsParams: ClothPhysicsParams = {
    dt: 0.016,
    gravityY: -9.81,
    compliance: 0.00001,
    numVertices: 0,
    windX: 0,
    windY: 0,
    windZ: 0,
    time: 0,
    elongation: 1.0
  };

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

    const clothData = this.generateClothData(config);
    this.positions = clothData.positions;
    this.prevPositions = clothData.prevPositions;
    this.velocities = clothData.velocities;
    this.constraintPhases = clothData.constraints.map(c => new Float32Array(c));

    this.setupRenderMesh(scene, config);
    this.startRenderLoop(scene, config);
  }

  updateParams(params: ClothSimulationParams): void {
    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;

    this.physicsParams.compliance = softCompliance * Math.pow(rigidCompliance / softCompliance, t);
    this.physicsParams.windX = windX;
    this.physicsParams.windY = 0.0;
    this.physicsParams.windZ = windZ;
    this.physicsParams.time = params.simulationTime;
    this.physicsParams.elongation = params.elongation;
  }

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

  dispose(): void {
    if (this.scene && this.clothMesh) {
      this.scene.removeMesh(this.clothMesh);
    }
    this.clothMesh = 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 setupRenderMesh(scene: Scene, config: ClothConfig): void {
    if (this.clothMesh) {
      scene.removeMesh(this.clothMesh);
    }

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

    const clothMaterial = new ShaderMaterial("clothMat", scene, {
      vertexSource: CLOTH_VERTEX_SHADER_GLSL,
      fragmentSource: CLOTH_FRAGMENT_SHADER_GLSL
    }, {
      attributes: ["position", "uv"],
      uniforms: ["viewProjection"]
    });

    clothMaterial.backFaceCulling = false;
    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;
    }
  }

  /**
   * Extracts xyz from the vec4 positions array into a vec3 array for mesh vertex update.
   */
  private extractPositionsVec3(positions: Float32Array, numVertices: number): Float32Array {
    const result = new Float32Array(numVertices * 3);
    for (let i = 0; i < numVertices; i++) {
      result[i * 3 + 0] = positions[i * 4 + 0];
      result[i * 3 + 1] = positions[i * 4 + 1];
      result[i * 3 + 2] = positions[i * 4 + 2];
    }
    return result;
  }

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

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

      integratePositions(this.positions, this.prevPositions, this.velocities, this.physicsParams);

      for (let i = 0; i < substeps; i++) {
        for (const phase of this.constraintPhases) {
          solveConstraints(this.positions, phase, this.physicsParams);
        }
      }

      updateVelocities(this.positions, this.prevPositions, this.velocities, this.physicsParams);

      const posVec3 = this.extractPositionsVec3(this.positions, config.numVertices);
      this.clothMesh.updateVerticesData(VertexBuffer.PositionKind, posVec3);
    });
  }
}
