import { AnimationClip, Bone,Interpolant, KeyframeTrack, Matrix4, Object3D, PropertyBinding, Quaternion, Vector3 } from "three";

import { isDevEnvironment, showBalloonWarning } from "../../../../engine/debug/debug.js";
import { getParam } from "../../../../engine/engine_utils.js";
import { Animator } from "../../../Animator.js";
import { GameObject } from "../../../Component.js";
import type { IUSDExporterExtension } from "../Extension.js";
import { buildMatrix, findStructuralNodesInBoneHierarchy, getPathToSkeleton,usdNumberFormatting as fn, USDObject, USDWriter, USDZExporterContext } from "../ThreeUSDZExporter.js";

const debug = getParam("debugusdzanimation");
const debugSerialization = getParam("debugusdzanimationserialization");

export interface UsdzAnimation {
    createAnimation(ext: AnimationExtension, model: USDObject, context);
}

export type AnimationClipCollection = Array<{ root: Object3D, clips: Array<AnimationClip> }>;

export class RegisteredAnimationInfo {

    private _start?: number;
    get start(): number { 
        if (this._start === undefined) {
            this._start = this.ext.getStartTimeByClip(this.clip);
        }
        return this._start;
    }
    get duration(): number { return this.clip?.duration ?? TransformData.restPoseClipDuration; }
    get nearestAnimatedRoot(): Object3D | undefined { return this._nearestAnimatedRoot; }
    get clipName(): string { return this.clip?.name ?? "rest"; }

    private ext: AnimationExtension;
    private root: Object3D;
    private _nearestAnimatedRoot?: Object3D = undefined;
    private clip: AnimationClip | null;

    // Playback speed. Does not affect how the animation is written, just how fast actions play it back.
    speed?: number;

    constructor(ext: AnimationExtension, root: Object3D, clip: AnimationClip | null) {
        this.ext = ext;
        this.root = root;
        this.clip = clip;

        this._nearestAnimatedRoot = this.getNearestAnimatedRoot();
    }

    private static isDescendantOf(parent: Object3D, child: Object3D) {
        let current: Object3D | null = child;
        if (!current || !parent) return false;
        while (current) {
            if (!current) return false;
            if (current === parent) return true;
            current = current.parent;
        }
        return false;
    }

    /** Finds the nearest actually animated object under root based on the tracks in the AnimationClip. */
    getNearestAnimatedRoot() {
        let highestRoot: Object3D | undefined = undefined;
        try {
            for (const track of this.clip?.tracks ?? []) {
                const parsedPath = PropertyBinding.parseTrackName(track.name);
                let animationTarget = PropertyBinding.findNode(this.root, parsedPath.nodeName);
                if (animationTarget) {
                    if (!highestRoot) highestRoot = animationTarget;
                    else {
                        if (animationTarget === highestRoot) continue;
                        if (RegisteredAnimationInfo.isDescendantOf(highestRoot, animationTarget)) continue;
                        if (!RegisteredAnimationInfo.isDescendantOf(animationTarget, highestRoot)) {
                            // TODO test this, should find the nearest common ancestor
                            while (!RegisteredAnimationInfo.isDescendantOf(animationTarget, highestRoot) && animationTarget.parent) {
                                animationTarget = animationTarget.parent;
                            }
                            if (!RegisteredAnimationInfo.isDescendantOf(animationTarget, highestRoot)) {
                                console.error("USDZExporter: Animation clip targets multiple roots that are not parent/child. Please report a bug", this.root, this.clip, highestRoot, animationTarget);
                            }
                        }
                        highestRoot = animationTarget;
                    }
                }
            }
        } catch (e) {
            console.error("USDZExporter: Exception when trying to find nearest animated root. Please report a bug", e);
            highestRoot = undefined;
        }
        return highestRoot;
    }
}

export class TransformData {
    clip: AnimationClip | null;
    pos?: KeyframeTrack;
    rot?: KeyframeTrack;
    scale?: KeyframeTrack;

    private root: Object3D | null;
    private target: Object3D;
    private duration = 0;
    private useRootMotion = false;

    /** This value can theoretically be anything – a value of 1 is good to clearly see animation gaps.
     * For production, a value of 1/60 is enough, since the files can then still properly play back at 60fps.
     */
    static frameRate = 60;
    static animationDurationPadding = 6 / 60;
    static restPoseClipDuration = 6 / 60;

    constructor(root: Object3D | null, target: Object3D, clip: AnimationClip | null) {
        this.root = root;
        this.target = target;
        this.clip = clip;

        // this is a rest pose clip.
        // we assume duration 1/60 and no tracks, and when queried for times we just return [0, duration]
        if (!clip) {
            this.duration = TransformData.restPoseClipDuration;
        }
        else {
            this.duration = clip.duration;
        }

        // warn if the duration does not equal the maximum time value in the tracks
        if (clip && clip.tracks) {
            const maxTime = Math.max(...clip.tracks.map((t) => t.times[t.times.length - 1]));
            if (maxTime !== this.duration) {
                console.warn("USDZExporter: Animation clip duration does not match the maximum time value in the tracks.", clip, maxTime, this.duration);
                // We need to override the duration, otherwise we end up with gaps in the exported animation
                // where there are actually no written animation values
                this.duration = maxTime;
            }
        }

        const animator = GameObject.getComponent(root, Animator);
        if (animator) this.useRootMotion = animator.applyRootMotion;
    }

    addTrack(track: KeyframeTrack) {
        if (!this.clip) {
            console.error("This is a rest clip but you're trying to add tracks to it – this is likely a bug");
            return;
        }

        if (track.name.endsWith("position")) this.pos = track;
        else if (track.name.endsWith("quaternion")) this.rot = track;
        else if (track.name.endsWith("scale")) this.scale = track;
        else {
            if (track.name.endsWith("activeSelf")) {
                /*
                // Construct a scale track, then apply to the existing scale track.
                // Not supported right now, because it would also require properly tracking that these objects need to be enabled
                // at animation start because they're animating the enabled state... otherwise they just stay disabled from scene start
                const newValues = [...track.values].map((v) => v ? [1,1,1] : [0,0,0]).flat();
                const scaleTrack = new KeyframeTrack(track.name.replace(".activeSelf", ".scale"), track.times, newValues, InterpolateDiscrete); 
                if (!this.scale) 
                {
                    this.scale = scaleTrack;
                    console.log("Mock scale track", this.scale);
                }
                */
                console.warn("[USDZ] Animation of enabled/disabled state is not supported for USDZ export and will NOT be exported: " + track.name + " on " + (this.root?.name ?? this.target.name) + ". Animate scale 0/1 instead.");
            }
            else {
                console.warn("[USDZ] Animation track type not supported for USDZ export and will NOT be exported: " + track.name + " on " + (this.root?.name ?? this.target.name) + ". Only .position, .rotation, .scale are supported.");
            }

            if (isDevEnvironment()) showBalloonWarning("[USDZ] Some animations can't be exported. See console for details.");
        }
    }

    getFrames(): number {
        if (!this.clip) return 2;
        return Math.max(this.pos?.times?.length ?? 0, this.rot?.times?.length ?? 0, this.scale?.times?.length ?? 0);
    }

    getDuration(): number {
        return this.duration;
    }

    getSortedTimesArray(generatePos: boolean = true, generateRot: boolean = true, generateScale: boolean = true) {
        if (!this.clip) return [0, this.duration];

        const posTimesArray = this.pos?.times;
        const rotTimesArray = this.rot?.times;
        const scaleTimesArray = this.scale?.times;
        
        // timesArray is the sorted union of all time values
        const timesArray: number[] = [];
        if (generatePos && posTimesArray) for (const t of posTimesArray) timesArray.push(t);
        if (generateRot && rotTimesArray) for (const t of rotTimesArray) timesArray.push(t);
        if (generateScale && scaleTimesArray) for (const t of scaleTimesArray) timesArray.push(t);

        // We also need to make sure we have start and end times for these tracks
        // We already ensure this is the case for duration – it's always the maximum time value in the tracks
        // (see constructor)
        if (!timesArray.includes(0)) timesArray.push(0);

        // sort times so it's increasing
        timesArray.sort((a, b) => a - b);
        // make sure time values are unique
        return [...new Set(timesArray)];
    }

    /**
     * Returns an iterator that yields the values for each time sample. 
     * Values are reused objects - if you want to append them to some array 
     * instead of processing them right away, clone() them.
     * @param timesArray 
     * @param generatePos 
     * @param generateRot 
     * @param generateScale 
     */
    *getValues(timesArray: number[], generatePos: boolean = true, generateRot: boolean = true, generateScale: boolean = true) {

        const translation = new Vector3();
        const rotation = new Quaternion();
        const scale = new Vector3(1, 1, 1);
        const object = this.target;

        const positionInterpolant: Interpolant | undefined = generatePos ? this.pos?.createInterpolant() : undefined;
        const rotationInterpolant: Interpolant | undefined = generateRot ? this.rot?.createInterpolant() : undefined;
        const scaleInterpolant: Interpolant | undefined = generateScale ? this.scale?.createInterpolant() : undefined;

        if (!positionInterpolant) translation.set(object.position.x, object.position.y, object.position.z);
        if (!rotationInterpolant) rotation.set(object.quaternion.x, object.quaternion.y, object.quaternion.z, object.quaternion.w);
        if (!scaleInterpolant) scale.set(object.scale.x, object.scale.y, object.scale.z);

        // WORKAROUND because it seems that sometimes we get a quaternion interpolant with a valueSize
        // of its own length... which then goes into an endless loop
        if (positionInterpolant && positionInterpolant.valueSize !== 3)
            positionInterpolant.valueSize = 3;
        if (rotationInterpolant && rotationInterpolant.valueSize !== 4)
            rotationInterpolant.valueSize = 4;
        if (scaleInterpolant && scaleInterpolant.valueSize !== 3)
            scaleInterpolant.valueSize = 3;

        // We're optionally padding with one extra time step at the beginning and the end
        // So that we don't get unwanted interpolations between clips (during the padding time)
        const extraFrame = 0; // TransformData.animationDurationPadding > 1 / 60 ? 1 : 0;
        for (let index = 0 - extraFrame; index < timesArray.length + extraFrame; index++) {
            let time = 0;
            let returnTime = 0;
            if (index < 0) {
                time = timesArray[0];
                returnTime = time - (TransformData.animationDurationPadding / 2) + 1 / 60;
            }
            else if (index >= timesArray.length) {
                time = timesArray[timesArray.length - 1];
                returnTime = time + (TransformData.animationDurationPadding / 2) - 1 / 60;
            }
            else {
                time = timesArray[index];
                returnTime = time;
            }

            if (positionInterpolant) {
                const pos = positionInterpolant.evaluate(time);
                translation.set(pos[0], pos[1], pos[2]);
            }
            if (rotationInterpolant) {
                const quat = rotationInterpolant.evaluate(time);
                rotation.set(quat[0], quat[1], quat[2], quat[3]);
            }
            if (scaleInterpolant) {
                const scaleVal = scaleInterpolant.evaluate(time);
                scale.set(scaleVal[0], scaleVal[1], scaleVal[2]);
            }

            // Apply basic root motion offset – non-animated transformation data is applied to the node again.
            // We're doing this because clips that animate their own root are (typically) not in world space,
            // but in local space and moved to a specific spot in the world.
            if (this.useRootMotion && object === this.root) {
                const rootMatrix = new Matrix4();
                rootMatrix.compose(translation, rotation, scale);
                rootMatrix.multiply(object.matrix);
                rootMatrix.decompose(translation, rotation, scale);
            }

            yield { time: returnTime, translation, rotation, scale, index };
        }
    }
}

export class AnimationExtension implements IUSDExporterExtension {

    get extensionName(): string { return "animation" }
    get animationData() { return this.dict; }
    get registeredClips() { return this.clipToStartTime.keys(); }
    get animatedRoots() { return this.rootTargetMap.keys(); }
    get holdClipMap() { return this.clipToHoldClip; }

    /** For each animated object, contains time/pos/rot/scale samples in the format that USD needs,
     *  ready to be written to the .usda file.
     */
    private dict: Map<Object3D, Array<TransformData>> = new Map();
    /** Map of all roots (Animation/Animator or scene) and all targets that they animate.
     *  We need that info so that we can ensure that each target has the same number of TransformData entries
     *  so that switching between animations doesn't result in data "leaking" to another clip.
     */
    private rootTargetMap = new Map<Object3D, Array<Object3D>>();
    private rootAndClipToRegisteredAnimationMap = new Map<string, RegisteredAnimationInfo>();
    /** Clips registered for each root */
    private rootToRegisteredClip = new Map<Object3D, Array<AnimationClip>>();
    private lastClipEndTime = 0;
    private clipToStartTime: Map<AnimationClip, number> = new Map();
    private clipToHoldClip: Map<AnimationClip, AnimationClip> = new Map();

    private serializers: SerializeAnimation[] = [];
    
    /** Determines if we inject a rest pose clip for each root - only makes sense for QuickLook */
    injectRestPoses = false;
    /** Determines if we inject a PlayAnimationOnClick component with "scenestart" trigger - only makes sense for QuickLook */
    injectImplicitBehaviours = false;

    constructor(quickLookCompatible: boolean) {
        this.injectRestPoses = quickLookCompatible;
        this.injectImplicitBehaviours = quickLookCompatible;
    }

    getStartTimeCode() {
        // return the end time + padding of the rest clip, if any exists
        if (!this.injectRestPoses) return 0;
        if (this.rootAndClipToRegisteredAnimationMap.size === 0) return 0;
        // return 0;
        return (TransformData.restPoseClipDuration + TransformData.animationDurationPadding) * 60;
    }

    /** Returns the end time code, based on 60 frames per second, for all registered animations.
     * This matches the highest time value in the USDZ file. */
    getEndTimeCode() {
        let max = 0;
        for (const [_, info] of this.rootAndClipToRegisteredAnimationMap) {
            const end = info.start + info.duration;
            if (end > max) max = end;
        }
        return max * 60;
    }

    getClipCount(root: Object3D): number {
        const currentCount = this.rootToRegisteredClip.get(root)?.length ?? 0;
        // The rest pose is not part of rootToRegisteredClip
        // if (this.injectRestPoses) currentCount = currentCount ? currentCount - 1 : 0;
        return currentCount ?? 0;
    }

    /*
    // TODO why do we have this here and on TransformData? Can RegisteredAnimationInfo not cache this value?
    // TODO we probably want to assert here that this is the same value on all nodes
    getStartTime01(root: Object3D, clip: AnimationClip | null) {
        // This is a rest pose clip, it always starts at 0
        if (!clip) return 0;

        const targets = this.rootTargetMap.get(root);
        if (!targets) return 0;
        const transformDatas = this.dict.get(targets[0]);
        if (!transformDatas) {
            console.error("Trying to get start time for root that has no animation data", root, clip, ...this.dict);
            return 0;
        }

        let currentStartTime = 0;
        for (let i = 0; i < transformDatas.length; i++) {
            if (transformDatas[i].clip === clip) break;
            currentStartTime += transformDatas[i].getDuration() + TransformData.animationDurationPadding;
        }

        return currentStartTime;
    }
    */

    getStartTimeByClip(clip: AnimationClip | null) {
        if (!clip) return 0;
        if (!this.clipToStartTime.has(clip)) {
            console.error("USDZExporter: Missing start time for clip – please report a bug.", clip);
            return 0;
        } 
        const time = this.clipToStartTime.get(clip)!;
        return time;
    }

    // The same clip could be registered for different roots. All of them need written animation data then.
    // The same root could have multiple clips registered to it. If it does, the clips need to write
    // independent time data, so that playing back an animation on that root doesn't result in data "leaking"/"overlapping".
    // The structure we need is:
    // - MyRoot
    //   Animator
    //     - Clip1: CubeScale (only animates MyCube), duration: 3s
    //     - Clip2: SphereRotation (only animates MySphere), duration: 2s
    //   - MyCube
    //   - MySphere
    // Results in:
    // - MyRoot
    //   - MyCube
    //     - # rest clip (0..0.1)
    //     - # CubeScale (0.2..3.2)
    //     - # rest clip for SphereRotation (3.3..5.3)
    //   - MySphere
    //     - # rest clip (0..0.1)
    //     - # rest clip for CubeScale (0.2..3.2)
    //     - # SphereRotation (3.3..5.3)

    /** Register an AnimationClip for a specific root object.
     * @param root The root object that the animation clip is targeting.
     * @param clip The animation clip to register. If null, a rest pose is registered.
     * @returns The registered animation info, which contains the start time and duration of the clip.
     */
    registerAnimation(root: Object3D, clip: AnimationClip | null): RegisteredAnimationInfo | null {
        if(!root) return null;
        if (!this.rootTargetMap.has(root)) this.rootTargetMap.set(root, []);

        // if we registered that exact pair already, just return the info
        // no proper tuples in JavaScript, but we can use the uuids for a unique key here
        const hash = root.uuid + (clip?.uuid ?? "-rest");
        if (this.rootAndClipToRegisteredAnimationMap.has(hash)) {
            return this.rootAndClipToRegisteredAnimationMap.get(hash)!;
        }

        if (debug) console.log("registerAnimation", root, clip);

        // When injecting a rest clip, the rest clip has ALL animated nodes as targets.
        // So all other nodes will already have at least one animated clip registered, and their own
        // animations need to start at index 1. Otherwise we're getting overlap where everything is
        // in animation 0 and some data overrides each other incorrectly.
        // When we don't inject a rest clip, we start at 0.
        const startIndex = this.injectRestPoses ? 1 : 0;
        const currentCount = (this.rootToRegisteredClip.get(root)?.length ?? 0) + startIndex;

        const targets = this.rootTargetMap.get(root);
        const unregisteredNodesForThisClip = new Set(targets);
        if (clip && clip.tracks) {
        // We could sort so that supported tracks come first, this allows us to support some additional tracks by 
        // modifying what has already been written for the supported ones (e.g. enabled -> modify scale).
        // Only needed if we're actually emulating some unsupported track types in addTrack(...).
        // const sortedTracks = clip.tracks.filter(x => !!x).sort((a, _b) => a.name.endsWith("position") || a.name.endsWith("quaternion") || a.name.endsWith("scale") ? -1 : 1);
        for (const track of clip.tracks) {
            const parsedPath = PropertyBinding.parseTrackName(track.name);
            const animationTarget = PropertyBinding.findNode(root, parsedPath.nodeName);
            
            if (!animationTarget) {
                console.warn("no object found for track", track.name, "using " + root.name + " instead");
                continue;
                // // if no object was found it might be that we have a component that references an animation clip but wants to target another object
                // // in that case UnityGLTF writes the name of the component as track targets because it doesnt know of the intented target
                // animationTarget = root;
            }
            if (!this.dict.has(animationTarget)) {
                this.dict.set(animationTarget, []);
            }
            const transformDataForTarget = this.dict.get(animationTarget);
            if (!transformDataForTarget) {
                console.warn("no transform data found for target ", animationTarget, "at slot " + currentCount + ", this is likely a bug");
                continue;
            }

            // this node has animation data for this clip – no need for additional padding
            unregisteredNodesForThisClip.delete(animationTarget);

            // Since we're interleaving animations, we need to ensure that for the same root,
            // all clips that "touch" it are written in the same order for all animated nodes.
            // this means we need to pad the TransformData array with empty entries when a particular
            // node inside that root is not animated by a particular clip.
            // It also means that when we encounter a clip that contains animation data for a new node,
            // We need to pad that one's array as well so it starts at the same point.
            // TODO most likely doesn't work for overlapping clips (clips where a root is a child of another root)
            // TODO in that case we may need to pad globally, not per root

            // Inject a rest pose if we don't have it already
            if (this.injectRestPoses && !transformDataForTarget[0]) {
                console.log("Injecting rest pose", animationTarget, clip, "at slot", currentCount);
                transformDataForTarget[0] = new TransformData(null, animationTarget, null);
            }
            // These all need to be at the same index, otherwise our padding went wrong
            let model = transformDataForTarget[currentCount];
            
            if (!model) {
                model = new TransformData(root, animationTarget, clip);
                transformDataForTarget[currentCount] = model;
            }
            model.addTrack(track);

            // We're keeping track of all animated nodes per root, needed for proper padding
            if (!targets?.includes(animationTarget)) targets?.push(animationTarget);
        }
        }

        if (debug) console.log("Unregistered nodes for this clip", unregisteredNodesForThisClip, "clip", clip, "at slot", currentCount, "for root", root, "targets", targets)

        // add padding for nodes that are not animated by this clip
        for (const target of unregisteredNodesForThisClip) {
            const transformDataForTarget = this.dict.get(target);
            if (!transformDataForTarget) continue;

            // Inject rest pose if these nodes don't have it yet for some reason – this is likely a bug
            if (this.injectRestPoses && !transformDataForTarget[0]) {
                console.warn("Adding rest pose for ", target, clip, "at slot", currentCount, "This is likely a bug, should have been added earlier.");
                const model = new TransformData(null, target, null);
                transformDataForTarget[0] = model;
            }

            let model = transformDataForTarget[currentCount];
            if (!model) {
                if (debug) console.log("Adding padding clip for ", target, clip, "at slot", currentCount);
                model = new TransformData(root, target, clip);
                transformDataForTarget[currentCount] = model;
            }
        }

        // get the entry for this object. 
        // This doesnt work if we have clips animating multiple objects
        const info = new RegisteredAnimationInfo(this, root, clip);
        this.rootAndClipToRegisteredAnimationMap.set(hash, info);
        
        if (debug) console.log({root, clip, info});
        if (clip) {
            const registered = this.rootToRegisteredClip.get(root);
            if (!registered) this.rootToRegisteredClip.set(root, [clip]);
            else registered.push(clip);

            const lastClip = this.clipToStartTime.get(clip);
            if (!lastClip) {
                if (this.lastClipEndTime == null) this.lastClipEndTime = TransformData.restPoseClipDuration;
                
                let newStartTime = this.lastClipEndTime + TransformData.animationDurationPadding;
                let newEndTime = newStartTime + clip.duration;

                // Round these times, makes it easier to understand what happens in the file
                const roundedStartTime = Math.round(newStartTime * 60) / 60;
                const roundedEndTime = Math.round(newEndTime * 60) / 60;
                if (Math.abs(roundedStartTime - newStartTime) < 0.01) newStartTime = roundedStartTime;
                if (Math.abs(roundedEndTime - newEndTime) < 0.01) newEndTime = roundedEndTime;
                // Round newStartTime up to the next frame
                newStartTime = Math.ceil(newStartTime);
                newEndTime = newStartTime + clip.duration;
                
                this.clipToStartTime.set(clip, newStartTime);
                this.lastClipEndTime = newEndTime;
            }
        }
        return info;
    }

    onAfterHierarchy(_context) {
        if (debug) console.log("Animation clips per animation target node", this.dict);
    }

    onAfterBuildDocument(_context: any) {

        // Validation: go through all roots, check if their data and lengths are consistent.
        // There are some cases that we can patch up here, especially if non-overlapping animations have resulted in "holes"
        // in TransformData where only now we know what the matching animation clip actually is.
        if (debug) console.log("Animation data", { dict: this.dict, rootTargetMap: this.rootTargetMap, rootToRegisteredClip: this.rootToRegisteredClip});

        for (const root of this.rootTargetMap.keys()) {
            const targets = this.rootTargetMap.get(root);

            if (!targets) continue;

            // The TransformData[] arrays here should have the same length, and the same durations
            let arrayLength: number | undefined = undefined;
            const durations: Array<number | undefined> = [];
            for (const target of targets) {
                const datas = this.dict.get(target);
                if (!datas) {
                    console.error("No data found for target on USDZ export – please report a bug!", target);
                    continue;
                }

                if (arrayLength === undefined) arrayLength = datas?.length;
                if (arrayLength !== datas?.length) console.error("Different array lengths for targets – please report a bug!", datas);

                for (let i = 0; i < datas.length; i++) {
                    let data = datas[i];
                    if (!data) {
                        // If we don't have TransformData for this object yet, we're emitting a rest pose with the duration
                        // of the matching clip that was registered for this root. 
                        const index = i - (this.injectRestPoses ? 1 : 0);
                        datas[i] = new TransformData(null, target, this.rootToRegisteredClip.get(root)![index]);
                        data = datas[i];
                    }

                    const duration = data.getDuration();
                    if (durations[i] === undefined) durations[i] = duration;
                    else if (durations[i] !== duration) {
                        console.error("Error during UDSZ export: Encountered different animation durations for animated targets. Please report a bug!", {datas, target});
                        durations[i] = duration;
                        continue;
                    }
                }
            }
        }

        for (const ser of this.serializers) {
            const parent = ser.model?.parent;
            const isEmptyParent = parent?.isDynamic === true;
            if (debugSerialization) console.log(isEmptyParent, ser.model?.parent);
            if (isEmptyParent) {
                ser.registerCallback(parent);
            }
        }
    }

    onExportObject(object, model: USDObject, _context) {

        GameObject.foreachComponent(object, (comp) => {
            const c = comp as unknown as UsdzAnimation;
            if (typeof c.createAnimation === "function") {
                c.createAnimation(this, model, _context);
            }
        }, false);

        // we need to be able to retarget serialization to empty parents before actually serializing (we do that in another callback)
        const ser = new SerializeAnimation(object, this);
        this.serializers.push(ser);
        ser.registerCallback(model);
    }

}

declare type TransformDataByObject = Map<Object3D, TransformData[]>;
declare type AnimationClipFrameTimes = { pos: number[], rot: number[], scale: number[], timeOffset: number };

class SerializeAnimation {

    model: USDObject | undefined = undefined;
    
    private object: Object3D;
    private animationData: Map<Object3D, Array<TransformData>>;
    private ext: AnimationExtension;

    private callback?: (writer: USDWriter, context: USDZExporterContext) => void;

    constructor(object: Object3D, ext: AnimationExtension) {
        this.object = object;
        this.animationData = ext.animationData;
        this.ext = ext;
    }

    registerCallback(model: USDObject) {
        if (this.model && this.callback) {
            this.model.removeEventListener("serialize", this.callback);
        }
        if (!this.callback)
            this.callback = this.onSerialize.bind(this);
        if (debugSerialization) console.log("REPARENT", model);
        this.model = model;
        
        if (this.callback)
            this.model.addEventListener("serialize", this.callback);
    }

    skinnedMeshExport(writer: USDWriter, _context: USDZExporterContext, ext: AnimationExtension) {
        const model = this.model;
        const dict = this.animationData;
        if (!model) return;

        if ( model.skinnedMesh ) {

            const skeleton = model.skinnedMesh.skeleton;

            const boneAndInverse = new Array<{bone: Object3D, inverse: Matrix4}>();

            const sortedBones:Bone[] = [];
            const uuidsFound:string[] = [];
            for(const bone of skeleton.bones ) {
                // if (bone.parent!.type !== 'Bone')
                {
                    sortedBones.push(bone);
                    uuidsFound.push(bone.uuid);
                    const inverse = skeleton.boneInverses[skeleton.bones.indexOf(bone)];
                    boneAndInverse.push({bone, inverse});
                }
            }

            let maxSteps = 10_000;
            while (uuidsFound.length < skeleton.bones.length && maxSteps-- > 0) {
                for (const sortedBone of sortedBones){
                    const children = sortedBone.children as Bone[];
                    for (const childBone of children){
                        if (uuidsFound.indexOf(childBone.uuid) === -1 && skeleton.bones.indexOf(childBone) !== -1){
                            sortedBones.push(childBone);
                            uuidsFound.push(childBone.uuid);
                            const childInverse = skeleton.boneInverses[skeleton.bones.indexOf(childBone)];
                            boneAndInverse.push({bone: childBone, inverse: childInverse});
                        }
                    }
                }
            }

            if (maxSteps <= 0) console.error("Failed to sort bones in skinned mesh", model.skinnedMesh, skeleton.bones, uuidsFound);

            for (const structuralNode of findStructuralNodesInBoneHierarchy(skeleton.bones)) {
                boneAndInverse.push({bone: structuralNode, inverse: structuralNode.matrixWorld.clone().invert()});
            }
            
            // sort bones by path – need to be sorted in the same order as during mesh export
            const assumedRoot = boneAndInverse[0].bone.parent!;
            if (!assumedRoot) console.error("No bone parent found for skinned mesh during USDZ export", model.skinnedMesh);
            boneAndInverse.sort((a, b) => getPathToSkeleton(a.bone, assumedRoot) > getPathToSkeleton(b.bone, assumedRoot) ? 1 : -1);

            function createVector3TimeSampleLines_( values: Map<number, Vector3[]> ) {

                const lines:string[] = []
                for (const [frame, frameValues] of values) {
                    let line = `${frame} : [`;
                    const boneRotations: Array<string> = [];
                    for ( const v of frameValues ) {
                        boneRotations.push( `(${fn( v.x )}, ${fn( v.y )}, ${fn( v.z )})` );
                    }
                    line = line.concat( boneRotations.join( ', ' ) );
                    line = line.concat( '],' );
                    lines.push( line );
                }

                return lines;
                
            }

            function createVector4TimeSampleLines_( rotations: Map<number, Quaternion[]> ) {

                const lines:string[] = []

                for (const [frame, frameRotations] of rotations) {
                    let line = `${frame} : [`;
                    const boneRotations: Array<string> = [];
                    for ( const v of frameRotations ) {
                        boneRotations.push( `(${fn( v.w )}, ${fn( v.x )}, ${fn( v.y )}, ${fn( v.z )})` );
                    }
                    line = line.concat( boneRotations.join( ', ' ) );
                    line = line.concat( '],' );
                    lines.push( line );
                }

                return lines;

            }

            function getSortedFrameTimes( boneToTransformData: TransformDataByObject ): AnimationClipFrameTimes[] {
        
                // We should have a proper rectangular array,
                // Where for each bone we have the same number of TransformData entries.
                let numberOfEntries: number | undefined = undefined;
                let allBonesHaveSameNumberOfTransformDataEntries = true;
                const undefinedBoneEntries = new Map<Object3D, Array<number>>();
                for (const [bone, transformDatas] of boneToTransformData) {
                    if (numberOfEntries === undefined) numberOfEntries = transformDatas.length;
                    if (numberOfEntries !== transformDatas.length) {
                        allBonesHaveSameNumberOfTransformDataEntries = false;
                    }
                    let index = 0;
                    for (const transformData of transformDatas) {
                        index++;
                        if (!transformData) {
                            if (!undefinedBoneEntries.has(bone)) undefinedBoneEntries.set(bone, []);
                            undefinedBoneEntries.get(bone)!.push(index);
                        }
                    }
                }

                // TODO not working yet for multiple skinned characters at the same time
                if (debug) {
                    console.log("Bone count: ", boneToTransformData.size, "TransformData entries per bone: ", numberOfEntries, "Undefined bone entries: ", undefinedBoneEntries);
                }
                console.assert(allBonesHaveSameNumberOfTransformDataEntries, "All bones should have the same number of TransformData entries", boneToTransformData);
                console.assert(undefinedBoneEntries.size === 0, "All TransformData entries should be set", undefinedBoneEntries);
                const times: AnimationClipFrameTimes[] = [];

                for (const [bone, transformDatas] of boneToTransformData) {
                    
                    /*
                    // calculate start times from the transformDatas
                    const startTimes = new Array<number>();
                    let currentStartTime = 0;
                    for (let i = 0; i < transformDatas.length; i++) {
                        startTimes.push(currentStartTime);
                        currentStartTime += transformDatas[i].getDuration() + TransformData.animationDurationPadding;
                    }
                    */

                    for (let i = 0; i < transformDatas.length; i++) {
                        const transformData = transformDatas[i];
                        // const timeOffset = transformData.getStartTime(dict);
                        const timeOffset = ext.getStartTimeByClip(transformData.clip);
                        if (times.length <= i) {
                            times.push({pos: [], rot: [], scale: [], timeOffset});
                        }

                        const perTransfromDataTimes = times[i];

                        perTransfromDataTimes.pos.push( ...transformData.getSortedTimesArray(true, false, false));
                        perTransfromDataTimes.rot.push( ...transformData.getSortedTimesArray(false, true, false));
                        perTransfromDataTimes.scale.push( ...transformData.getSortedTimesArray(false, false, true));
                    }
                }

                for (const perTransfromDataTimes of times) {

                    /*
                    // TODO we're doing that in animation export as well
                    if (!times.pos.includes(0)) times.pos.push(0);
                    if (!times.rot.includes(0)) times.rot.push(0);
                    if (!times.scale.includes(0)) times.scale.push(0);
                    */

                    // sort times so it's increasing
                    perTransfromDataTimes.pos.sort((a, b) => a - b);
                    perTransfromDataTimes.rot.sort((a, b) => a - b);
                    perTransfromDataTimes.scale.sort((a, b) => a - b);

                    // make sure time values are unique
                    perTransfromDataTimes.pos = [...new Set(perTransfromDataTimes.pos)];
                    perTransfromDataTimes.rot = [...new Set(perTransfromDataTimes.rot)];
                    perTransfromDataTimes.scale = [...new Set(perTransfromDataTimes.scale)];
                }

                return times;
            }

            function createTimeSamplesObject_( data: TransformDataByObject, sortedComponentFrameNumbers: AnimationClipFrameTimes[], bones: Array<Object3D> ) {
                const positionTimeSamples = new Map<number, Array<Vector3>>();
                const quaternionTimeSamples = new Map<number, Array<Quaternion>>();
                const scaleTimeSamples = new Map<number, Array<Vector3>>();

                const count = sortedComponentFrameNumbers.length;

                // return sampled data for each bone
                for ( const bone of bones ) {
                    const boneEntryInData = data.get( bone );

                    let emptyTransformData: TransformData | undefined = undefined;
                    // if we have animation data for this bone, check that it's the right amount.
                    if (boneEntryInData)
                        console.assert(boneEntryInData.length === count, "We should have the same number of TransformData entries for each bone", boneEntryInData, sortedComponentFrameNumbers);
                    // if we don't have animation data, create an empty one – 
                    // it will automatically map to the rest pose, albeit inefficiently
                    else
                        emptyTransformData = new TransformData(null, bone, null);

                    for (let i = 0; i < count; i++) {
                        const transformData = boneEntryInData ? boneEntryInData[i] : emptyTransformData!;
                        const timeData = sortedComponentFrameNumbers[i];

                        for (const { time, translation } of transformData.getValues(timeData.pos, true, false, false)) {
                            const shiftedTime = time + timeData.timeOffset;
                            const t = shiftedTime * 60;

                            if (!positionTimeSamples.has(t)) positionTimeSamples.set(t, new Array<Vector3>());
                            positionTimeSamples.get(t)!.push( translation.clone() );
                        }
                        for (const { time, rotation } of transformData.getValues(timeData.rot, false, true, false)) {
                            const shiftedTime = time + timeData.timeOffset;
                            const t = shiftedTime * 60;

                            if (!quaternionTimeSamples.has(t)) quaternionTimeSamples.set(t, new Array<Quaternion>());
                            quaternionTimeSamples.get(t)!.push( rotation.clone() );
                        }
                        for (const { time, scale } of transformData.getValues(timeData.scale, false, false, true)) {
                            const shiftedTime = time + timeData.timeOffset;
                            const t = shiftedTime * 60;

                            if (!scaleTimeSamples.has(t)) scaleTimeSamples.set(t, new Array<Vector3>());
                            scaleTimeSamples.get(t)!.push( scale.clone() );
                        }
                    }
                }

                return {
                    position: positionTimeSamples.size == 0 ? undefined : positionTimeSamples,
                    quaternion: quaternionTimeSamples.size == 0 ? undefined: quaternionTimeSamples,
                    scale: scaleTimeSamples.size == 0 ? undefined : scaleTimeSamples,
                };
            }

            function buildVector3Array_( array: Array<Vector3> ) {
                const lines: Array<string> = [];
                for ( const v of array ) {
                    lines.push( `(${fn( v.x )}, ${fn( v.y )}, ${fn( v.z )})` );
                }
                return lines.join( ', ' );
            }

            function buildVector4Array_( array: Array<Quaternion> ) {
                const lines: Array<string> = [];
                for ( const v of array ) {
                    lines.push( `(${fn( v.w )}, ${fn( v.x )}, ${fn( v.y )}, ${fn( v.z )})` );
                }
                return lines.join( ', ' );
            }

            function getPerBoneTransformData( bones: Array<Object3D> ): TransformDataByObject {

                const boneToTransformData = new Map<Object3D, TransformData[]>();                
                if (debug) {
                    const logData = new Array<string>();
                    for (const [key, val] of dict) {
                        logData.push(key.uuid + ": " + val.length + " " + val.map(x => x.clip?.uuid.substring(0, 6)).join(" "));
                    }
                    console.log("getPerBoneTransformData\n" + logData.join("\n"));
                }

                for (const bone of bones) {
                    const data = dict.get(bone);
                    if (!data) continue;
                    boneToTransformData.set(bone, data);
                }

                return boneToTransformData;
            }

            function createAllTimeSampleObjects( bones: Array<Object3D> ) {
                const perBoneTransformData = getPerBoneTransformData( bones );
                const sortedFrameNumbers = getSortedFrameTimes( perBoneTransformData );
                return createTimeSamplesObject_( perBoneTransformData, sortedFrameNumbers, bones );
            }

            const sanitizeRestPose = _context.quickLookCompatible;

            const rest: Array<Matrix4> = [];
            const translations: Array<Vector3> = [];
            const rotations: Array<Quaternion> = [];
            const scales: Array<Vector3> = [];
            for ( const { bone } of boneAndInverse ){

                // Workaround for FB13808839: Rest poses must be decomposable in QuickLook
                if (sanitizeRestPose) {
                    const scale = bone.scale;
                    if (scale.x == 0) scale.x = 0.00001;
                    if (scale.y == 0) scale.y = 0.00001;
                    if (scale.z == 0) scale.z = 0.00001;
                    rest.push( new Matrix4().compose( bone.position, bone.quaternion, bone.scale )) ;
                } else {
                    rest.push( bone.matrix.clone() );
                }
                translations.push( bone.position ) ;
                rotations.push( bone.quaternion );
                scales.push( bone.scale );
            }

            const bonesArray = boneAndInverse.map( x => "\"" + getPathToSkeleton(x.bone, assumedRoot) + "\"" ).join( ', ' );
            const bindTransforms = boneAndInverse.map( x => buildMatrix( x.inverse.clone().invert() ) ).join( ', ' );
            
            writer.beginBlock( `def Skeleton "Rig"`);

            writer.appendLine( `uniform matrix4d[] bindTransforms = [${bindTransforms}]` );
            writer.appendLine( `uniform token[] joints = [${bonesArray}]` );
            writer.appendLine( `uniform token purpose = "guide"` );
            writer.appendLine( `uniform matrix4d[] restTransforms = [${rest.map( m => buildMatrix( m ) ).join( ', ' )}]` );
            
            // In glTF, transformations on the Skeleton are ignored (NODE_SKINNED_MESH_LOCAL_TRANSFORMS validator warning)
            // So here we don't write anything to get an identity transform.
            // writer.appendLine( `matrix4d xformOp:transform = ${buildMatrix( new Matrix4() )}` );
            // writer.appendLine( `uniform token[] xformOpOrder = ["xformOp:transform"]` );

            const timeSampleObjects = createAllTimeSampleObjects ( boneAndInverse.map( x => x.bone ) );

            if (debug) {
                // find the first..last value in the time samples
                let min = 10000000;
                let max = 0;
                for (const key of timeSampleObjects.position?.keys() ?? []) {
                    min = Math.min(min, key);
                    max = Math.max(max, key);
                }
                console.log("Time samples", min, max, timeSampleObjects);
            }

            writer.beginBlock( `def SkelAnimation "_anim"` );

            // TODO if we include blendshapes we likely need subdivision?
            // writer.appendLine( `uniform token[] blendShapes` )
            // writer.appendLine( `float[] blendShapeWeights` )

            writer.appendLine( `uniform token[] joints = [${bonesArray}]` );

            writer.appendLine( `quatf[] rotations = [${buildVector4Array_( rotations )}]` );
            
            if ( timeSampleObjects && timeSampleObjects.quaternion ) {
                writer.beginBlock( `quatf[] rotations.timeSamples = {`, '');
                const rotationTimeSampleLines = createVector4TimeSampleLines_( timeSampleObjects['quaternion'] );
                for ( const line of rotationTimeSampleLines ) {
                    writer.appendLine( line );
                }
                writer.closeBlock();

            }

            writer.appendLine( `half3[] scales = [${buildVector3Array_( scales )}]` );

            if ( timeSampleObjects && timeSampleObjects.scale ) {
                writer.beginBlock( `half3[] scales.timeSamples = {`, '' );
                const scaleTimeSampleLines = createVector3TimeSampleLines_( timeSampleObjects['scale'] );

                for ( const line of scaleTimeSampleLines ) {
                    writer.appendLine( line );
                }
                writer.closeBlock();

            }

            writer.appendLine( `float3[] translations = [${buildVector3Array_( translations )}]` );

            if ( timeSampleObjects && timeSampleObjects.position) {
                writer.beginBlock( `float3[] translations.timeSamples = {`, '' );
                const positionTimeSampleLines = createVector3TimeSampleLines_( timeSampleObjects['position'] );

                for ( const line of positionTimeSampleLines ) {
                    writer.appendLine( line );
                }
                writer.closeBlock();

            }

            writer.closeBlock();
            writer.closeBlock();

        }
    }

    onSerialize(writer: USDWriter, _context: USDZExporterContext) {
        if (!this.model) return;

        // Workaround: Sanitize TransformData for this object.
        // This works around an issue with wrongly detected animation roots, where some of the indices
        // in the TransformData array are not property set. Reproduces with golem_yokai.glb
        const arr0 = this.animationData.get(this.object);
        if (arr0) {
            for (let i = 0; i < arr0.length; i++) {
                if (arr0[i] !== undefined) continue;
                arr0[i] = new TransformData(null, this.object, null);
            }
        }

        const ext = this.ext;

        this.skinnedMeshExport(writer, _context, ext);

        const object = this.object;
        const model = this.model;

        // do we have animation data for this node? if not, return
        const animationData = this.animationData.get(object);
        if (!animationData) return;

        // Skinned meshes are handled separately by the method above. 
        // They need to be handled first (before checking for animation data) because animation needs to be exported
        // as part of the skinned mesh and that may not be animated at all – if any bone is animated we need to export.
        //@ts-ignore
        if (object.isSkinnedMesh) return;

        // Excluding the concrete bone xform hierarchy animation his is mostly useful for debugging, 
        // as otherwise we're getting a ton of extra animation data per-bone xform
        // TODO if this turns out to be slow/large (lots of duplicated animation data)
        // we can look into optimizing it, and only including both the xform hierarchy and the animation
        // only when we need it – when anywhere below it in the hierarchy are animated visible meshes.
        //@ts-ignore
        // if (object.isBone) return;

        if (debugSerialization) console.log("SERIALIZE", this.model.name, this.object.type, animationData);

        // const composedTransform = new Matrix4();
        // writer.appendLine("matrix4d xformOp:transform.timeSamples = {");
        // writer.indent++;

        // TransformData is a collection of clips (position, rotation and scale) for a particular node
        // We need to make sure that the same underlying animation clip ends up
        // at the same start time in the USD file, and that we're not getting overlaps to other clips.
        // That means that the same clip (transformData.clip) should end up at the same start time for all nodes.

        // calculate start times
        // calculate start times from the transformDatas
        /*
        const startTimes = new Array<number>();
        let currentStartTime = 0;
        for (let i = 0; i < animationData.length; i++) {
            startTimes.push(currentStartTime);
            if (animationData[i] === undefined) {
                console.error("Got an undefined transform data, this is likely a bug.", object, animationData);
                continue;
            }
            currentStartTime += animationData[i].getDuration() + TransformData.animationDurationPadding;
        }
        */

        const formatter = Intl.NumberFormat("en-US", {
            maximumFractionDigits: 3,
            minimumFractionDigits: 0,
            useGrouping: false,
        });

        function writeAnimationTimesamples(arr: TransformData[], type: "position" | "rotation" | "scale") {

            const hasAnimationData = arr.some(x => x && {
                position: x.pos,
                rotation: x.rot,
                scale: x.scale
            }[type]);
            if (!hasAnimationData) {
                return;
            }

            switch (type) {
                case "position":
                    model.needsTranslate = true;
                    writer.beginBlock(`double3 xformOp:translate.timeSamples = {`, '');
                    break;
                case "rotation":
                    model.needsOrient = true;
                    writer.beginBlock(`quatf xformOp:orient.timeSamples = {`, '');
                    break;
                case "scale":
                    model.needsScale = true;
                    writer.beginBlock(`double3 xformOp:scale.timeSamples = {`, '');
                    break;
            }

            for (let i = 0; i < arr.length; i++) {
                const transformData = arr[i];
                if (!transformData) continue;
                const startTime = ext.getStartTimeByClip(transformData.clip);
                const timesArray = transformData.getSortedTimesArray(type === "position", type === "rotation", type === "scale");

                if (!timesArray || timesArray.length === 0) {
                    console.error("got an animated object but no time values?", object, transformData);
                    continue;
                }

                const isRestPose = !transformData.clip;
                const hasPos = type === "position" && (transformData.pos || isRestPose);
                const hasRot = type === "rotation" && (transformData.rot || isRestPose);
                const hasScale = type === "scale" && (transformData.scale || isRestPose);

                // Writing out the clip name and duration.
                // If this is a rest pose clip, we still want to write it out, as it's a valid clip.
                if (hasPos || hasRot || hasScale)
                {
                    const clipName = transformData.clip?.name ?? "rest";
                    const duration = transformData.getDuration();
                    if (debug) console.log("Write .timeSamples:", clipName, startTime, duration, arr);
                    writer.appendLine("# " + clipName + ": start=" + formatter.format(startTime * TransformData.frameRate) + ", length=" + formatter.format(duration * TransformData.frameRate) + ", frames=" + transformData.getFrames());
                }

                if (hasPos) {
                    for (const { time, translation } of transformData.getValues(timesArray, true, false, false)) {
                        const timeStr = formatter.format((startTime + time) * TransformData.frameRate);
                        const line = `${timeStr}: (${fn(translation.x)}, ${fn(translation.y)}, ${fn(translation.z)}),`;
                        writer.appendLine(line);
                    }
                }
                if (hasRot) {
                    for (const { time, rotation } of transformData.getValues(timesArray, false, true, false)) {
                        const timeStr = formatter.format((startTime + time) * TransformData.frameRate);
                        const line = `${timeStr}: (${fn(rotation.w)}, ${fn(rotation.x)}, ${fn(rotation.y)}, ${fn(rotation.z)}),`;
                        writer.appendLine(line);
                    }
                }
                if (hasScale) {
                    for (const { time, scale } of transformData.getValues(timesArray, false, false, true)) {
                        const timeStr = formatter.format((startTime + time) * TransformData.frameRate);
                        const line = `${timeStr}: (${fn(scale.x)}, ${fn(scale.y)}, ${fn(scale.z)}),`;
                        writer.appendLine(line);
                    }
                }
            }
            writer.closeBlock();
        }
        writeAnimationTimesamples(animationData, "position");
        writeAnimationTimesamples(animationData, "rotation");
        writeAnimationTimesamples(animationData, "scale");

        // We must be careful here that we don't overwrite the xformOpOrder that the object already had.
        // it _might_ not even have any (if the position/quaternion/scale are all indentity values)
        // so in some cases we end up writing the same xformOpOrder multiple times here...
        // So right now, we just always write the xformOpOrder when traversing the hierarchy, in case something is animated.
        
        // if (xFormOrder.length > 0) {
            // const xformUnique = [...new Set(xFormOrder)];
            // writer.appendLine(`uniform token[] xformOpOrder = [${xformUnique.map(x => `"${x}"`).join(', ')}]`);
            // writer.appendLine('uniform token[] xformOpOrder = ["xformOp:translate", "xformOp:orient", "xformOp:scale"]');
        // }

        // for (let i = 0; i < arr.length; i++) {
        //     const transformData = arr[i];
        //     if (!transformData) continue;
        //     const startTime = startTimes[i];
        //     const timesArray = transformData.getSortedTimesArray();

        //     if (!timesArray || timesArray.length === 0) {
        //         console.error("got an animated object but no time values?", object, transformData);
        //         continue;
        //     }
        //     for (const { time, translation, rotation, scale } of transformData.getValues(timesArray)) {
        //         composedTransform.compose(translation, rotation, scale);


        //         const timeStr = formatter.format((startTime + time) * transformData.frameRate);
        //         const line = `${timeStr}: ${buildMatrix(composedTransform)},`;
        //         writer.appendLine(line);
        //     }
        // }

        // writer.indent--;
        // writer.appendLine("}");
    }
}