import { RigidBody, type RigidBodyProps, type RapierRigidBody } from "@react-three/rapier";
import { useFrame } from "@react-three/fiber";
import { useRef } from "react";
import { Euler, Quaternion } from "three";
import type { R3, V3 } from "../types";

export type SyncRigidBodyData = {
    position: V3;
    rotation: R3;
    linearVelocity: V3;
    radialVelocity: R3;
}

export type SyncRigidBodyOnSyncFunction = (data: SyncRigidBodyData) => void;

type SyncRigidBodyProps = RigidBodyProps & {
    onSync: SyncRigidBodyOnSyncFunction;
};

const _q = new Quaternion();
const _e = new Euler();
const EPS = 1e-6;

// indices: 0-2 position, 3-5 rotation, 6-8 linearVelocity, 9-11 radialVelocity
const PREV_INIT = new Float64Array(12).fill(Infinity);

function syncRigidBodyDataToArray(data: SyncRigidBodyData, arr: Float64Array): void {
    arr[0] = data.position[0];      arr[1] = data.position[1];      arr[2] = data.position[2];
    arr[3] = data.rotation[0];      arr[4] = data.rotation[1];      arr[5] = data.rotation[2];
    arr[6] = data.linearVelocity[0]; arr[7] = data.linearVelocity[1]; arr[8] = data.linearVelocity[2];
    arr[9] = data.radialVelocity[0]; arr[10] = data.radialVelocity[1]; arr[11] = data.radialVelocity[2];
}

function compareTwoFloatArrays(a: Float64Array, b: Float64Array, epsilon: number): boolean {
    for (let i = 0; i < a.length; i++)
        if (Math.abs(a[i] - b[i]) > epsilon)
            return true;
    return false;
}

export function SyncRigidBody({ onSync, children, ...props }: SyncRigidBodyProps) {
    const rbRef = useRef<RapierRigidBody>(null);
    const prevData = useRef<Float64Array>(PREV_INIT.slice());
    const currentData = useRef<Float64Array>(PREV_INIT.slice());

    useFrame(() => {
        const body = rbRef.current;
        if (!body)
            return;

        const { x, y, z } = body.translation();
        const rot = body.rotation();
        _q.set(rot.x, rot.y, rot.z, rot.w);
        _e.setFromQuaternion(_q);
        const lv = body.linvel();
        const av = body.angvel();

        const data: SyncRigidBodyData = {
            position: [x, y, z],
            rotation: [_e.x, _e.y, _e.z],
            linearVelocity: [lv.x, lv.y, lv.z],
            radialVelocity: [av.x, av.y, av.z],
        };
        
        syncRigidBodyDataToArray(data, currentData.current);

        if (compareTwoFloatArrays(currentData.current, prevData.current, EPS)) {
            prevData.current.set(currentData.current);
            onSync(data);
        }
    });

    return <RigidBody ref={rbRef} {...props}>{children}</RigidBody>;
}
