import { AxesHelper, Group, Material, Mesh, Object3D, XRHandSpace } from "three";
import { GLTFLoader } from "three/examples/jsm/loaders/GLTFLoader.js";
import { XRControllerModelFactory } from "three/examples/jsm/webxr/XRControllerModelFactory.js";
import { XRHandMeshModel } from "three/examples/jsm/webxr/XRHandMeshModel.js";

import { AssetReference } from "../../../engine/engine_addressables.js";
import { setDontDestroy } from "../../../engine/engine_gameobject.js";
import { Gizmos } from "../../../engine/engine_gizmos.js";
import { getLoader } from "../../../engine/engine_gltf.js";
import { addDracoAndKTX2Loaders } from "../../../engine/engine_loaders.gltf.js";
import { serializable } from "../../../engine/engine_serialization_decorator.js";
import type { IGameObject, SourceIdentifier } from "../../../engine/engine_types.js";
import { getParam } from "../../../engine/engine_utils.js";
import { NeedleXRController, type NeedleXRControllerEventArgs, type NeedleXREventArgs, NeedleXRSession } from "../../../engine/engine_xr.js";
import { registerComponentExtension, registerExtensions } from "../../../engine/extensions/extensions.js";
import { NEEDLE_progressive } from "../../../engine/extensions/NEEDLE_progressive.js";
import { flipForwardMatrix } from "../../../engine/xr/internal.js";
import { Behaviour, Component, GameObject } from "../../Component.js"

const debug = getParam("debugwebxr");

const handsJointBuffer = new Float32Array(16 * 25);
const renderingUpdateTimings = new Array<number>();

/**
 * XRControllerModel is a component that allows to display controller models or hand models in an XR session.  
 * It automatically loads the appropriate model for the connected controller or hand.  
 * 
 * You can configure if controller models or hand models should be created.
 * 
 * @summary Displays controller or hand models in XR
 * @category XR
 * @group Components
 */
export class XRControllerModel extends Behaviour {

    /**
     * If true, the controller model will be created when a controller is added/connected
     * @default true
     */
    @serializable()
    createControllerModel: boolean = true;

    /**
     * If true, the hand model will be created when a hand is "added"/tracked
     * @default true
     */
    @serializable()
    createHandModel: boolean = true;

    /** assign a model or model url to create custom hand models */
    @serializable(AssetReference)
    customLeftHand?: AssetReference;
    /** assign a model or model url to create custom hand models */
    @serializable(AssetReference)
    customRightHand?: AssetReference;


    static readonly factory: XRControllerModelFactory = new XRControllerModelFactory();

    supportsXR(mode: XRSessionMode): boolean {
        return mode === "immersive-vr" || mode === "immersive-ar";
    }

    private readonly _models = new Array<{ model?: IGameObject, controller: NeedleXRController, handmesh?: XRHandMeshModel }>();


    async onXRControllerAdded(args: NeedleXRControllerEventArgs) {
        // TODO we may want to treat controllers differently in AR/Passthrough mode
        const isSupportedSession = args.xr.isVR || args.xr.isPassThrough;
        if (!isSupportedSession) return;

        console.debug("XR Controller Added", args.controller.side, args.controller.index);
        // return;

        const { controller } = args;

        if (this.createControllerModel || this.createHandModel) {
            if (controller.hand) {
                if (this.createHandModel) {
                    const res = await this.loadHandModel(this, controller);
                    // check if the model doesnt exist, the hand disconnected or it's suddenly a controller
                    if (!res || !controller.connected || !controller.isHand) {
                        if (res?.handObject) setDontDestroy(res.handObject, false);
                        res?.handObject?.destroy();
                        return;
                    }
                    this._models.push({ controller: controller, model: res.handObject, handmesh: res.handmesh });
                    this._models.sort((a, b) => a.controller.index - b.controller.index);
                    this.scene.add(res.handObject);
                    controller.model = res.handObject;
                }
            }
            else {
                if (this.createControllerModel) {
                    const assetUrl = await controller.getModelUrl();
                    if (assetUrl) {
                        const model = await this.loadModel(controller, assetUrl);
                        // check if the model doesnt exist, the controller disconnected or it's suddenly a hand
                        if (!model || !controller.connected || controller.isHand) {
                            return;
                        }
                        this._models.push({ controller: controller, model });
                        this._models.sort((a, b) => a.controller.index - b.controller.index);
                        this.scene.add(model);
                        // The controller mesh should by default inherit layers.
                        model.traverse(child => {
                            child.layers.set(2);
                            // disable auto update on controller objects. No need to do this every frame
                            child.matrixAutoUpdate = false;
                            child.updateMatrix();
                        });
                        controller.model = model;
                    }
                    else if (controller.targetRayMode !== "transient-pointer") {
                        console.warn("XRControllerModel: no model found for " + controller.side);
                    }
                }
            }
        }
    }

    onXRControllerRemoved(args: NeedleXRControllerEventArgs): void {
        console.debug("XR Controller Removed", args.controller.side, args.controller.index);
        // we need to find the index by the controller because if controller 0 is removed first then args.controller.index 1 will be at index 0
        const indexInArray = this._models.findIndex(m => m.controller === args.controller);
        const entry = this._models[indexInArray];
        if (!entry) return;

        this._models.splice(indexInArray, 1);

        if (entry.model) {
            setDontDestroy(entry.model, false);
            entry.model.destroy();
            entry.model = undefined;
        }
    }

    onBeforeXR(_mode: XRSessionMode, args: XRSessionInit & { trackedImages: Array<any> }): void {
        // When a custom hand model is used, we want to ensure we're requesting hand tracking,
        // even when the platform default doesn't include it (for example, on VisionOS we don't
        // request hand tracking by default because there's an additional permissions dialogue.
        if (this.createHandModel && (this.customLeftHand || this.customRightHand)) {
            args.optionalFeatures = args.optionalFeatures || [];
            if (!args.optionalFeatures.includes("hand-tracking")) {
                args.optionalFeatures.push("hand-tracking");
            }
        }
    }

    onLeaveXR(_args: NeedleXREventArgs): void {
        for (const entry of this._models) {
            if (!entry) continue;

            if (entry.model) {
                setDontDestroy(entry.model, false);
                entry.model.destroy();
                entry.model = undefined;
            }

            // Unassign the model from the controller when this script becomes inactive
            if (entry.controller.model === entry.model) {
                entry.controller.model = null;
            }
        }
        this._models.length = 0;
    }

    onBeforeRender() {
        if (!NeedleXRSession.active) return;

        if (debug) renderingUpdateTimings[0] = Date.now();
        // update model
        this.updateRendering(NeedleXRSession.active);

        if (debug) {
            const dt = Date.now() - renderingUpdateTimings[0];
            renderingUpdateTimings.push(dt);
            if (renderingUpdateTimings.length >= 30) {
                renderingUpdateTimings[0] = 0;
                const avrg = renderingUpdateTimings.reduce((a, b) => a + b, 0) / renderingUpdateTimings.length;
                renderingUpdateTimings.length = 0;
                // console.log("[XRControllerModel] " + avrg.toFixed(2) + " ms");
            }
        }
    }

    private updateRendering(xr: NeedleXRSession) {

        for (let i = 0; i < this._models.length; i++) {
            const entry = this._models[i];
            if (!entry) continue;
            const ctrl = entry.controller;
            if (!ctrl.connected) {
                // the actual removal of the model happens in onXRControllerRemoved
                if (debug) console.warn("XRControllerModel.onUpdateXR: controller is not connected anymore", ctrl.side, ctrl.hand);
                continue;
            }


            // do we have a controller model?
            if (entry.model && !entry.handmesh) {
                entry.model.matrixAutoUpdate = false;
                entry.model.matrix.copy(ctrl.gripMatrix);
                entry.model.visible = ctrl.isTracking;
                // ensure that controller models are in rig space
                xr.rig?.gameObject.add(entry.model);
            }
            // do we have a hand mesh?
            else if (ctrl.inputSource.hand && entry.handmesh) {
                const referenceSpace = xr.referenceSpace;
                const hand = this.context.renderer.xr.getHand(ctrl.index);
                // if (referenceSpace && xr.frame.fillPoses) {
                //     xr.frame.fillPoses(ctrl.inputSource.hand.values(), referenceSpace, handsJointBuffer);
                //     let j = 0;
                //     for (const space of ctrl.inputSource.hand.values()) {
                //         const joint = hand.joints[space.jointName];
                //         if (joint) {
                //             joint.matrix.fromArray(handsJointBuffer, j * 16);
                //             joint.matrix.decompose(joint.position, joint.quaternion, joint.scale);
                //             joint.visible = true;
                //         }
                //         j++;
                //     }
                // }
                // else 
                if (referenceSpace && xr.frame.getJointPose) {
                    for (const inputjoint of ctrl.inputSource.hand.values()) {
                        // The transform of this joint will be updated with the joint pose on each frame
                        const joint = hand.joints[inputjoint.jointName];
                        if (joint) {
                            // Update the joints groups with the XRJoint poses
                            const jointPose = ctrl.getHandJointPose(inputjoint);
                            if (jointPose) {
                                const position = jointPose.transform.position;
                                const quaternion = jointPose.transform.orientation;
                                joint.position.copy(position);
                                joint.quaternion.copy(quaternion);
                                joint.matrixAutoUpdate = false;
                            }
                            joint.visible = jointPose != null;
                        }
                    }
                    // ensure that the hand renders in rig space
                    if (entry.model) {
                        entry.model.visible = ctrl.isTracking;
                        if (entry.model.visible && entry.model.parent !== xr.rig?.gameObject) {
                            xr.rig?.gameObject.add(entry.model);
                        }
                    }

                    if (entry.model?.visible) {
                        entry.handmesh?.updateMesh();
                        entry.model.matrixAutoUpdate = false;
                        entry.model.matrix.identity();
                        entry.model.applyMatrix4(flipForwardMatrix);
                    }
                }
            }
        }
    }

    protected async loadModel(controller: NeedleXRController, url: string): Promise<IGameObject | null> {
        if (!controller.connected) {
            console.warn("XRControllerModel.onXRControllerAdded: controller is not connected anymore", controller.side);
            return null;
        }
        const assetReference = AssetReference.getOrCreate("", url);
        const model = await assetReference.instantiate() as GameObject;
        setDontDestroy(model);

        if (NeedleXRSession.active?.isPassThrough) {
            model.traverseVisible((obj: Object3D) => {
                this.makeOccluder(obj);
            })
        }
        return model as IGameObject;
    }

    protected async loadHandModel(comp: Component, controller: NeedleXRController): Promise<{ handObject: IGameObject, handmesh: XRHandMeshModel } | null> {

        const context = this.context;
        const hand = context.renderer.xr.getHand(controller.index);
        if (!hand) {
            if (debug) Gizmos.DrawLabel(controller.rayWorldPosition, "No hand found for index " + controller.index, .05, 5);
            else console.warn("No hand found for index " + controller.index);
        }

        const loader = new GLTFLoader();
        addDracoAndKTX2Loaders(loader, context);
        await registerExtensions(loader, context, this.sourceId ?? "", this.sourceId ?? "");
        const componentsExtension = registerComponentExtension(loader);

        let filename = "";

        const customHand = controller.side === "left" ? this.customLeftHand : this.customRightHand;
        if (customHand) {
            const urlWithoutExtension = customHand.url.split('.').slice(0, -1).join('.');
            filename = urlWithoutExtension;
            loader.setPath("");
        }
        else {
            // DEFAULT hands
            // XRHandmeshModel is using "<handedness>.glb" for loading the file
            filename = controller.inputSource.handedness === "left" ? "left" : "right";
            loader.setPath('https://cdn.jsdelivr.net/npm/@webxr-input-profiles/assets@1.0/dist/profiles/generic-hand/');
        }


        const handObject = new Object3D();
        setDontDestroy(handObject);
        // @ts-ignore
        const handmesh = new XRHandMeshModel(handObject, hand, loader.path, filename, loader, (object: Object3D) => {

            const gltf = componentsExtension?.gltf;
            // The XRHandMeshController removes the hand from the gltf before calling this callback
            // we need this in the GLTF scene however for creating the builtin components 
            if (gltf?.scene.children?.length === 0) {
                gltf.scene.children[0] = object;
            }

            // console.log(controller.side, componentsExtension.gltf, object, componentsExtension.gltf.scene?.children)
            if (componentsExtension?.gltf)
                getLoader().createBuiltinComponents(comp.context, comp.sourceId || filename, componentsExtension.gltf, null, componentsExtension);

            // The hand mesh should not receive raycasts
            object.traverse(child => {
                child.layers.set(2);
                if (NeedleXRSession.active?.isPassThrough && !customHand)
                    this.makeOccluder(child);
                if (child instanceof Mesh) {
                    NEEDLE_progressive.assignMeshLOD(child, 0);
                }
            });
            if (!controller.connected) {
                if (debug) Gizmos.DrawLabel(controller.rayWorldPosition, "Hand is loaded but not connected anymore", .05, 5);
                object.removeFromParent();
            }
        });

        if (debug) handObject.add(new AxesHelper(.5));

        if (controller.inputSource.hand) {
            if (debug) console.log(controller.inputSource.hand);
            for (const inputjoint of controller.inputSource.hand.values()) {
                if (hand.joints[inputjoint.jointName] === undefined) {
                    const joint = new Group();
                    joint.matrixAutoUpdate = false;
                    joint.visible = true;
                    // joint.jointRadius = 0.01;
                    // @ts-ignore
                    hand.joints[inputjoint.jointName] = joint;
                    hand.add(joint);

                }
            }
        }
        else {
            if (debug) {
                Gizmos.DrawLabel(controller.rayWorldPosition, "No inputSource.hand found for index " + controller.index, .05, 5);
            }
        }

        return { handObject: handObject as IGameObject, handmesh: handmesh };
    }

    private makeOccluder(obj: Object3D) {
        if (obj instanceof Mesh) {
            let mat = obj.material;
            if (mat instanceof Material) {
                mat = obj.material = mat.clone();
                // depth only
                mat.depthWrite = true;
                mat.depthTest = true;
                mat.colorWrite = false;
                obj.receiveShadow = false;
                obj.renderOrder = -100;
            }
        }
    }
}