import * as THREE from "three";
import RAPIER from "@dimforge/rapier3d-compat";
import {computeNoiseOffset} from "../utils/common.js";

export class PhysicsWorld {
  constructor() {
    this.world = null;
    this.rigidBodies = [];
    this.cachedReverseData = [];
    this.groundBody = null;
    this.mouseBody = null;
    this.joint = null;
    this.activeRigidBody = null;
    this.dropFlags = null;
    this.cachedUVs = null;
    this.noiseOffsets = null;
  }

  async init() {
    await RAPIER.init({});
    this.world = new RAPIER.World({x: 0.0, y: 0.0, z: 0.0});
    this._createBoundaries();
  }

  _createBoundaries() {
    const floorThickness = 20.0;
    const groundDesc = RAPIER.RigidBodyDesc.fixed();
    this.groundBody = this.world.createRigidBody(groundDesc);
    const groundCollider = RAPIER.ColliderDesc.cuboid(1000.0, floorThickness / 2, 1000.0)
      .setTranslation(0, -(floorThickness / 2), 0)
      .setFriction(0.1)
      .setRestitution(0.9);
    this.world.createCollider(groundCollider, this.groundBody);

    const mouseDesc = RAPIER.RigidBodyDesc.kinematicPositionBased();
    this.mouseBody = this.world.createRigidBody(mouseDesc);
  }

  setGravity(y) {
    if (this.world) {
      this.world.gravity = {x: 0, y: y, z: 0};
    }
  }

  createGridBodies(gridInfo, params, settings) {
    this.clearBodies();

    const {cols, rows, width, height, spacing, cellSize, count} = gridInfo;
    const startX = -width * 0.5 + spacing * 0.5;
    const startY = -height * 0.5 + spacing * 0.5 + settings.gridOffsetY;
    const halfSize = cellSize / 2;

    this.dropFlags = new Uint8Array(count);
    this.cachedUVs = new Float32Array(count * 2);
    this.noiseOffsets = new Float32Array(count);
    const randomDepths = new Float32Array(count);

    for (let i = 0; i < count; i++) {
      const col = i % cols;
      const row = Math.floor(i / cols);
      const x = startX + col * spacing;
      const y = startY + row * spacing;

      const u = (col + 0.5) / cols;
      const v = (row + 0.5) / rows;
      this.cachedUVs[i * 2] = u;
      this.cachedUVs[i * 2 + 1] = v;

      const noiseOffset = computeNoiseOffset(i);
      this.noiseOffsets[i] = noiseOffset;

      const depthScale = params.shape === "random" ? (Math.random() < 0.2 ? 1.0 : 1.0 + Math.random() * 1.5) : 1.0;
      randomDepths[i] = depthScale;

      const rbDesc = RAPIER.RigidBodyDesc.fixed().setTranslation(x, y, 0).setLinearDamping(settings.bodyDamping).setAngularDamping(settings.bodyDamping);
      const rb = this.world.createRigidBody(rbDesc);

      const clDesc = RAPIER.ColliderDesc.cuboid(halfSize, halfSize, halfSize * depthScale)
        .setFriction(0.0)
        .setRestitution(settings.bodyRestitution);
      this.world.createCollider(clDesc, rb);

      rb.userData = {
        index: i,
        initPos: {x, y, z: 0},
        initRot: {x: 0, y: 0, z: 0, w: 1},
        depthScale,
        noiseOffset,
      };

      this.rigidBodies.push(rb);

      this.cachedReverseData.push({
        rb: rb,
        startPos: new THREE.Vector3(),
        startRot: new THREE.Quaternion(),
        endPos: new THREE.Vector3(x, y, 0),
        endRot: new THREE.Quaternion(0, 0, 0, 1),
        delay: i / count,
      });
    }
    return {randomDepths};
  }

  clearBodies() {
    if (this.joint) {
      this.world.removeImpulseJoint(this.joint, true);
      this.joint = null;
      this.activeRigidBody = null;
    }

    this.rigidBodies.forEach((rb) => {
      if (this.world.getRigidBody(rb.handle)) {
        rb.wakeUp();
        this.world.removeRigidBody(rb);
      }
    });

    this.rigidBodies = [];
    this.cachedReverseData = [];
    this.dropFlags = null;
    this.cachedUVs = null;
    this.noiseOffsets = null;
  }

  checkAndActivateDrop(rippleCenter, currentProgress, aspectVecX, organicMode, spread, spacing, count, onDropCallback) {
    if (!this.dropFlags) return;
    const spreadThreshold = spread * 0.98;

    for (let i = 0; i < count; i++) {
      if (this.dropFlags[i] === 1) continue;

      const u = this.cachedUVs[i * 2];
      const v = this.cachedUVs[i * 2 + 1];
      const noiseOffset = organicMode ? this.noiseOffsets[i] : 0;
      const dx = (u - rippleCenter.x) * aspectVecX;
      const dy = v - rippleCenter.y;

      const distSq = dx * dx + dy * dy;
      const threshold = currentProgress - spreadThreshold - noiseOffset;

      if (threshold > 0 && threshold * threshold > distSq) {
        this.dropFlags[i] = 1;
        const force = spacing * 2.0;
        this.activateBody(i, force);
        onDropCallback(i);
      }
    }
  }

  forceActivateDrop(onDropCallback) {
    if (!this.dropFlags) return;
    const force = 2.0;
    for (let i = 0; i < this.dropFlags.length; i++) {
      if (this.dropFlags[i] === 1) continue;
      this.dropFlags[i] = 1;
      this.activateBody(i, force);
      if (onDropCallback) onDropCallback(i);
    }
  }

  activateBody(index, initialForceScale = 1.0) {
    const rb = this.rigidBodies[index];
    if (!rb) return;

    rb.setBodyType(RAPIER.RigidBodyType.Dynamic);
    rb.setGravityScale(0.01);
    rb.setLinearDamping(1.2);
    rb.wakeUp();

    const force = initialForceScale;
    rb.setLinvel(
      {
        x: (Math.random() - 0.5) * force,
        y: (Math.random() - 0.5) * force,
        z: (Math.random() - 0.5) * force,
      },
      true
    );
    rb.setAngvel(
      {
        x: Math.random() - 0.5,
        y: Math.random() - 0.5,
        z: Math.random() - 0.5,
      },
      true
    );
  }

  resetDropFlags() {
    if (this.dropFlags) {
      this.dropFlags.fill(0);
    }
  }

  updateAllPhysicsParams(bodyDamping, bodyRestitution) {
    this.rigidBodies.forEach((rb) => {
      rb.setLinearDamping(bodyDamping);
      rb.setAngularDamping(bodyDamping);
      const collider = rb.collider(0);
      if (collider) collider.setRestitution(bodyRestitution);
      rb.wakeUp();
    });
  }

  resetBodyParameters(index, damping) {
    const rb = this.rigidBodies[index];
    if (rb && this.world.getRigidBody(rb.handle)) {
      rb.setGravityScale(1.0);
      rb.setLinearDamping(damping);
      rb.setAngularDamping(damping);
      rb.wakeUp();
    }
  }

  step() {
    if (this.world) {
      this.world.step();
    }
  }

  syncMesh(mesh, forceUpdate = false) {
    let needsUpdate = false;
    for (let i = 0; i < this.rigidBodies.length; i++) {
      const rb = this.rigidBodies[i];
      if (!rb.isSleeping() || forceUpdate) {
        mesh.setMatrixAt(i, rb.translation(), rb.rotation());
        needsUpdate = true;
      }
    }
    if (needsUpdate) mesh.requestMatrixUpdate();
  }

  startDrag(hitPoint, instanceId) {
    this.activeRigidBody = this.rigidBodies[instanceId];
    if (!this.activeRigidBody || this.activeRigidBody.bodyType() !== RAPIER.RigidBodyType.Dynamic) {
      return false;
    }

    this.mouseBody.setNextKinematicTranslation(hitPoint);

    const bodyPos = this.activeRigidBody.translation();
    const bodyRot = this.activeRigidBody.rotation();
    const vecPoint = new THREE.Vector3(hitPoint.x, hitPoint.y, hitPoint.z);
    const vecBody = new THREE.Vector3(bodyPos.x, bodyPos.y, bodyPos.z);
    const quatBody = new THREE.Quaternion(bodyRot.x, bodyRot.y, bodyRot.z, bodyRot.w);
    const localAnchor = vecPoint.sub(vecBody).applyQuaternion(quatBody.clone().invert());

    this.joint = this.world.createImpulseJoint(
      RAPIER.JointData.spherical({x: localAnchor.x, y: localAnchor.y, z: localAnchor.z}, {x: 0, y: 0, z: 0}),
      this.activeRigidBody,
      this.mouseBody,
      true
    );
    this.activeRigidBody.wakeUp();
    return true;
  }

  moveDrag(targetPos) {
    if (this.joint) {
      this.mouseBody.setNextKinematicTranslation(targetPos);
      if (this.activeRigidBody) this.activeRigidBody.wakeUp();
    }
  }

  endDrag() {
    if (this.joint) {
      this.world.removeImpulseJoint(this.joint, true);
      this.joint = null;
      this.activeRigidBody = null;
    }
  }

  prepareForReverse() {
    this.world.gravity = {x: 0, y: 0, z: 0};
    this.cachedReverseData.forEach((data) => {
      const pos = data.rb.translation();
      const rot = data.rb.rotation();
      data.startPos.set(pos.x, pos.y, pos.z);
      data.startRot.set(rot.x, rot.y, rot.z, rot.w);
      data.rb.setBodyType(RAPIER.RigidBodyType.Fixed);
      data.rb.sleep();
    });

    return this.cachedReverseData;
  }
}