import { Context, delay, isDevEnvironment, ObjectUtils } from "@needle-tools/engine";
import MaterialX from "../bin/JsMaterialXGenShader.js";
import { debug } from "./utils.js";
import { renderPMREMToEquirect } from "./textureHelper.js";
import { Light, MeshBasicMaterial, Object3D, PMREMGenerator } from "three";
import { registerLights } from "./helper.js";


// Global MaterialX module instance - initialized lazily
export const state = new class {
    materialXModule: any = null;
    materialXGenerator: any = null;
    materialXGenContext: any = null;
    materialXStdLib: any = null;
    materialXInitPromise: Promise<void> | null = null;

    // Global MaterialX environment instance
    private _materialXEnvironment: MaterialXEnvironment | null = null;
    get materialXEnvironment() {
        this._materialXEnvironment ??= new MaterialXEnvironment();
        return this._materialXEnvironment;
    }
}

// Initialize MaterialX WASM module lazily
export async function initializeMaterialX(): Promise<void> {
    if (state.materialXInitPromise) {
        return state.materialXInitPromise;
    }
    return state.materialXInitPromise = (async () => {
        if (state.materialXModule) return; // Already initialized
        if (debug) console.log("Initializing MaterialX WASM module...");
        try {

            const urls = await Promise.all([
                /** @ts-ignore */
                import(`../bin/JsMaterialXCore.wasm?url`),
                /** @ts-ignore */
                import(`../bin/JsMaterialXGenShader.wasm?url`),
                /** @ts-ignore */
                import(`../bin/JsMaterialXGenShader.data?url`),
            ]);
            const [JsMaterialXCore, JsMaterialXGenShader, JsMaterialXGenShader_data] = urls;

            const module = await MaterialX({
                locateFile: (path: string, scriptDirectory: string) => {
                    if (debug) console.debug("MaterialX locateFile called:", { path, scriptDirectory });

                    if (path.includes("JsMaterialXCore.wasm")) {
                        return JsMaterialXCore; // Use the URL for the core WASM file
                    }
                    else if (path.includes("JsMaterialXGenShader.wasm")) {
                        return JsMaterialXGenShader; // Use the URL for the shader WASM file
                    }
                    else if (path.includes("JsMaterialXGenShader.data")) {
                        return JsMaterialXGenShader_data; // Use the URL for the shader data file
                    }

                    return scriptDirectory + path;
                },
            });
            if (debug) console.log("MaterialXLoader module loaded", module);
            state.materialXModule = module;

            // Initialize shader generator and context
            state.materialXGenerator = module.EsslShaderGenerator.create();
            state.materialXGenContext = new module.GenContext(state.materialXGenerator);

            // Load standard libraries
            const tempDoc = module.createDocument();
            state.materialXStdLib = module.loadStandardLibraries(state.materialXGenContext);
            tempDoc.setDataLibrary(state.materialXStdLib);

            // Initialize basic lighting with default light rig
            const defaultLightRigXml = `<?xml version="1.0"?>
<materialx version="1.39">
  <!-- Default directional light -->
  <directional_light name="default_light" type="lightshader">
    <input name="direction" type="vector3" value="0.0, -1.0, -0.5" />
    <input name="color" type="color3" value="1.0, 1.0, 1.0" />
    <input name="intensity" type="float" value="1.0" />
  </directional_light>
</materialx>`;

            // This prewarms the shader generation context to have all light types
            await registerLights(state.materialXModule, [], state.materialXGenContext);

            if (debug) console.log("MaterialX generator initialized successfully");
        } catch (error) {
            console.error("Failed to load MaterialX module:", error);
            throw error;
        }
    })();
}

// MaterialX Environment Manager - handles lighting and environment setup
export class MaterialXEnvironment {
    private lights: any[] = [];
    private lightData: any = null;
    private radianceTexture: any = null;
    private irradianceTexture: any = null;
    private context: Context | null = null;
    private initialized: boolean = false;

    constructor() {
        if (debug) console.log("MaterialX Environment created");
    }

    setContext(context: Context) {
        this.context = context;
    }

    /*
    // Initialize MaterialX lighting system based on the reference implementation
    async initializeLighting(lightRigXml: string, renderer?: any, radianceTexture?: any, irradianceTexture?: any): Promise<void> {
        if (!materialXModule || !materialXGenContext) {
            console.warn("MaterialX module not initialized, skipping lighting setup");
            return;
        }

        registerLights(materialXModule, this.lights, materialXGenContext);
    }
    */

    // Initialize with Needle Engine context
    async initializeFromContext(): Promise<void> {
        if (!this.context) {
            console.warn("No Needle context available for MaterialX environment initialization");
            return;
        }

        // Prevent multiple initializations
        if (this.initialized) {
            if (debug) console.log("MaterialX environment already initialized, skipping");
            return;
        }

        // Clean up previous textures if they exist
        if (this.radianceTexture) {
            if (debug) console.log("Disposing previous radiance texture");
            this.radianceTexture.dispose();
            this.radianceTexture = null;
        }
        if (this.irradianceTexture) {
            if (debug) console.log("Disposing previous irradiance texture");
            this.irradianceTexture.dispose();
            this.irradianceTexture = null;
        }

        // Get renderer from context
        const renderer = this.context.renderer;

        // TODO remove this delay; we should wait for the scene lighting to be ready
        // and then update the uniforms
        let envMap = this.context.scene.environment;
        while (!envMap) {
            await delay(200);
            envMap = this.context.scene.environment;
        }
        var pmrem = new PMREMGenerator(renderer);
        const target = pmrem.fromEquirectangular(envMap);

        const radianceRenderTarget = renderPMREMToEquirect(renderer, target.texture, 0.0, 1024, 512, target.height);
        const irradianceRenderTarget = renderPMREMToEquirect(renderer, target.texture, 1.0, 32, 16, target.height);

        this.radianceTexture = radianceRenderTarget.texture;
        this.irradianceTexture = irradianceRenderTarget.texture;

        // Clean up PMREM generator and its render target
        target.dispose();
        pmrem.dispose();

        if (debug) {
            console.log({ radiance: this.radianceTexture, irradiance: this.irradianceTexture });
            // Show both of them on cubes in the scene
            const unlitMat = new MeshBasicMaterial();
            const radianceMat = unlitMat.clone();
            radianceMat.map = this.radianceTexture;
            const radianceCube = ObjectUtils.createPrimitive("Cube", { material: radianceMat });
            const irradianceMat = unlitMat.clone();
            irradianceMat.map = this.irradianceTexture;
            const irradianceCube = ObjectUtils.createPrimitive("Cube", { material: irradianceMat });
            this.context.scene.add(radianceCube);
            this.context.scene.add(irradianceCube);
            radianceCube.position.set(2, 0, 0);
            radianceCube.scale.y = 0.00001;
            irradianceCube.position.set(-2, 0, 0);
            irradianceCube.scale.y = 0.00001;
            // await this.initializeLighting(defaultLightRigXml, renderer);
            console.log("MaterialX environment initialized from Needle context", this, this.context.scene);
        }

        // Find lights in scene
        let lights = new Array<Light>();
        this.context.scene.traverse((object: Object3D) => {
            if ((object as Light).isLight) lights.push(object as Light);
        });

        this.lightData = await registerLights(state.materialXModule, lights, state.materialXGenContext);

        // Mark as initialized
        this.initialized = true;
    }

    // getLights() { return this.lights; }
    getLightData() { return this.lightData; }
    getRadianceTexture() { return this.radianceTexture; }
    getIrradianceTexture() { return this.irradianceTexture; }

    setRadianceTexture(texture: any) { this.radianceTexture = texture; }
    setIrradianceTexture(texture: any) { this.irradianceTexture = texture; }

    // Reset the environment to allow re-initialization
    reset() {
        if (debug) console.log("Resetting MaterialX environment");
        if (this.radianceTexture) {
            this.radianceTexture.dispose();
            this.radianceTexture = null;
        }
        if (this.irradianceTexture) {
            this.irradianceTexture.dispose();
            this.irradianceTexture = null;
        }
        this.initialized = false;
        // this.lights = [];
        this.lightData = null;
    }
}
