

import { Group, Camera, Material, Mesh, Object3D } from "three";
import { Context, GLTF, addCustomExtensionPlugin, Component, INeedleGLTFExtensionPlugin } from "@needle-tools/engine";
import type { GLTFLoader } from "three/examples/jsm/loaders/GLTFLoader.js";
import type { GLTFExporter } from "three/examples/jsm/exporters/GLTFExporter.js";
import { MaterialXLoader } from "./loader.three.js";
import { debug } from "../utils.js";
import { state } from "../materialx.js";

//@dont-generate-component
export class MaterialXUniformUpdate extends Component {

    static updateMaterial(mat: Material | Material[], object: Object3D, camera: Camera) {
        if (Array.isArray(mat)) {
            mat.forEach(m => {
                if (m.userData?.updateUniforms) {
                    m.userData.updateUniforms(object, camera);
                }
            });
        } else if (mat.userData?.updateUniforms) {
            mat.userData.updateUniforms(object, camera);
        }
    }

    onEnable(): void {
        this.context.addBeforeRenderListener(this.gameObject, this._onBeforeRender);
    }

    onDisable(): void {
        this.context.removeBeforeRenderListener(this.gameObject, this._onBeforeRender);
    }

    _onBeforeRender = () => {
        // Update uniforms or perform any pre-render logic here
        const gameObject = this.gameObject as any as Mesh;
        const material = gameObject?.material;

        const camera = this.context.mainCamera;
        if (!camera) return;

        MaterialXUniformUpdate.updateMaterial(material, gameObject, camera);

        // If this is a Group, we need to update all direct children
        if ((gameObject as any as Group).isGroup) {
            gameObject.children.forEach((child: Object3D) => {
                if (child instanceof Mesh && child.material)
                    MaterialXUniformUpdate.updateMaterial(child.material, child, camera);
            });
        }
    }
}

export class MaterialXLoaderPlugin implements INeedleGLTFExtensionPlugin {
    name = "MaterialXLoaderPlugin";

    mtlxLoader: MaterialXLoader | null = null;

    onImport = (loader: GLTFLoader, url: string, context: Context) => {
        if (debug) console.log("MaterialXLoaderPlugin: Registering MaterialX extension for", url);

        // Register the MaterialX loader extension
        // Environment initialization is now handled in the MaterialXLoader constructor
        loader.register(p => {
            this.mtlxLoader = new MaterialXLoader(p, context);
            return this.mtlxLoader;
        });
    };

    onLoaded = (url: string, gltf: GLTF, _context: Context) => {
        if (debug) console.log("MaterialXLoaderPlugin: glTF loaded", url, gltf.scene);

        // Set up onBeforeRender callbacks for objects with MaterialX materials
        // This ensures uniforms are updated properly during rendering
        gltf.scene.traverse((child) => {
            if ((child as any).isMesh) {
                const mesh = child as Mesh;
                const material = mesh.material as Material;

                if (material?.userData?.updateUniforms) {
                    if (debug) console.log("Adding MaterialX uniform update component to:", child.name);
                    child.addComponent(MaterialXUniformUpdate);
                }
            }
        });

        if (debug) console.log("Loaded: ", this.mtlxLoader);

        // Initialize MaterialX lighting system with scene data
        const environment = state.materialXEnvironment;
        environment.initializeFromContext().then(() => {
            if (this.mtlxLoader) {
                this.mtlxLoader.updateLightingFromEnvironment(environment);
            }
        });
    };

    onExport = (_exporter: GLTFExporter, _context: Context) => {
        console.log("TODO: MaterialXLoaderPlugin: Setting up export extensions");
        // TODO: Add MaterialX export functionality if needed
    };
}

export function registerNeedleLoader() {
    addCustomExtensionPlugin(new MaterialXLoaderPlugin());
}
