

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 { MaterialXEnvironment, state } from "../materialx.js";
import { MaterialXMaterial } from "../materialx.material.js";

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

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

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

    onBeforeRenderThree = () => {
        // 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;

        const env = state.materialXEnvironment;

        if (Array.isArray(material)) {
            for (const entry of material) {
                if (entry && entry instanceof MaterialXMaterial) {
                    entry.updateUniforms(this.context, env, this.gameObject, camera);
                }
            }
        } else if (material instanceof MaterialXMaterial) {
            material.updateUniforms(this.context, env, this.gameObject, camera);
        }
    }
}

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

    private loader: 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.loader = new MaterialXLoader(p, url, context);
            return this.loader;
        });
    };

    onLoaded = (url: string, gltf: GLTF, _context: Context) => {
        if (debug) console.log("[MaterialX] MaterialXLoaderPlugin: glTF loaded", { url, scene: gltf.scene, materialX_root_data: this.loader?.materialX_root_data });

        // If we don't have MaterialX data in the loaded glTF we don't need to do anything else here
        if (!this.loader?.materialX_root_data) {
            return;
        }
        // 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 instanceof MaterialXMaterial) {
                    if (debug) console.log("[MaterialX] Adding MaterialX uniform update component to:", child.name);
                    child.addComponent(MaterialXUniformUpdate);
                }
            }
        });

        if (debug) console.log("[MaterialX] Loaded: ", this.loader);
    };

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

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