import MaterialX from "../bin/JsMaterialXGenShader.js";
import { debug, waitForNetworkIdle } from "./utils.js";
import { renderPMREMToEquirect } from "./utils.texture.js";
import { Light, Mesh, MeshBasicMaterial, Object3D, PlaneGeometry, PMREMGenerator, Scene, Texture, WebGLRenderer } from "three";
import { registerLights, getLightData } from "./materialx.helper.js";
import { whiteTexture } from "./utils.texture.js";


/**
 * Preloads the MaterialX WebAssembly module.
 * @type {import("./materialx.js").preloadWasm}
 */
export async function preloadWasm(trigger) {
    if (trigger === "immediately") {
        // Load the WASM module immediately
        return ready();
    } else if (trigger === "network_idle") {
        // Wait for network to be idle before loading
        return waitForNetworkIdle().then(ready);
    }
}



export const state = new class {
    /** @type {import("./materialx.types.js").MaterialX.MODULE | null} */
    materialXModule = null;
    /** @type {any} */
    materialXGenerator = null;
    /** @type {any} */
    materialXGenContext = null;
    /** @type {any} */
    materialXStdLib = null;
    /** @type {Promise<void> | null} */
    materialXInitPromise = null;
}

/** 
 * Wait for the MaterialX WASM module to be ready.
 * @returns {Promise<void>}
 */
export async function ready() {
    if (state.materialXInitPromise) {
        return state.materialXInitPromise;
    }
    return state.materialXInitPromise = (async () => {
        if (state.materialXModule) return; // Already initialized
        if (debug) console.log("[MaterialX] Initializing WASM module...");
        try {

            /** @type {Array<string>} */
            const urls = await Promise.all([
                /** @ts-ignore */
                import( /* @vite-ignore */ `../bin/JsMaterialXCore.wasm?url`).then(m => m.default || m),
                /** @ts-ignore */
                import( /* @vite-ignore */ `../bin/JsMaterialXGenShader.wasm?url`).then(m => m.default || m),
                /** @ts-ignore */
                import( /* @vite-ignore */ `../bin/JsMaterialXGenShader.data.txt?url`).then(m => m.default || m),
            ]);
            const [JsMaterialXCore, JsMaterialXGenShader, JsMaterialXGenShader_data] = urls;

            const module = await MaterialX({
                locateFile: (/** @type {string} */ path, /** @type {string} */ scriptDirectory) => {
                    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("[MaterialX] module loaded", module);
            state.materialXModule = /** @type {import("./materialx.types.js").MaterialX.MODULE} */ (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);

            // TODO ShaderInterfaceType.SHADER_INTERFACE_REDUCED would be better, but doesn't actually seem to be supported in the MaterialX javascript bindings
            state.materialXGenContext.getOptions().shaderInterfaceType = state.materialXModule.ShaderInterfaceType.SHADER_INTERFACE_COMPLETE;

            // SPECULAR_ENVIRONMENT_NONE: Do not use specular environment maps.
            // SPECULAR_ENVIRONMENT_FIS: Use Filtered Importance Sampling for specular environment/indirect lighting.
            // SPECULAR_ENVIRONMENT_PREFILTER: Use pre-filtered environment maps for specular environment/indirect lighting.
            state.materialXGenContext.getOptions().hwSpecularEnvironmentMethod = state.materialXModule.HwSpecularEnvironmentMethod.SPECULAR_ENVIRONMENT_FIS;

            // TRANSMISSION_REFRACTION: Use a refraction approximation for transmission rendering.
            // TRANSMISSION_OPACITY: Use opacity for transmission rendering.
            // state.materialXGenContext.getOptions().hwTransmissionRenderMethod = state.materialXModule.HwTransmissionRenderMethod.TRANSMISSION_REFRACTION;

            // Turned off because we're doing color space conversion the three.js way
            state.materialXGenContext.getOptions().hwSrgbEncodeOutput = false;

            // Enables the generation of a prefiltered environment map.
            // TODO Would be great to use but requires setting more uniforms (like u_envPrefilterMip).
            // When set to true, the u_envRadiance map is expected to be a prefiltered environment map.
            // state.materialXGenContext.getOptions().hwWriteEnvPrefilter = true;

            // Set a reasonable default for max active lights
            state.materialXGenContext.getOptions().hwMaxActiveLightSources = 4;

            // state.materialXGenContext.getOptions().hwShadowMap = true;

            // 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("[MaterialX] Failed to load MaterialX module:", error);
            throw error;
        }
    })();
}

/**
 * @typedef {Object} EnvironmentTextureSet
 * @property {Texture | null} radianceTexture
 * @property {Texture | null} irradianceTexture
 */

/**
 * MaterialXEnvironment manages the environment settings for MaterialX materials. 
 */
// @dont-generate-component
export class MaterialXEnvironment {

    /**
     * @param {Scene} scene
     * @returns {MaterialXEnvironment | null}
     */
    static get(scene) {
        return this.getEnvironment(scene);
    }

    /** @type {WeakMap<Scene, MaterialXEnvironment>} */
    static _environments = new Map();

    /**
     * @param {Scene} scene
     * @returns {MaterialXEnvironment}
     */
    static getEnvironment(scene) {
        if (this._environments.has(scene)) {
            return /** @type {MaterialXEnvironment} */ (this._environments.get(scene));
        }
        const env = new MaterialXEnvironment(scene);
        this._environments.set(scene, env);
        return env;
    }

    /** @type {Array<Light>} */
    _lights = [];
    /** @type {import("./materialx.helper.js").LightData[] | null} */
    _lightData = null;
    /** @type {number} */
    _lightCount = 0;

    /** @type {Promise<boolean> | null} */
    _initializePromise = null;
    /** @type {boolean} */
    _isInitialized = false;
    /** @type {number} */
    _lastUpdateFrame = -1;

    /**
     * @param {Scene} _scene
     */
    constructor(_scene) {
        this._scene = _scene;
        if (debug) console.log("[MaterialX] Environment created");
    }

    /**
     * Initialize with Needle Engine context
     * @param {WebGLRenderer} renderer
     * @returns {Promise<boolean>}
     */
    async initialize(renderer) {
        if (this._initializePromise) {
            return this._initializePromise;
        }
        this._initializePromise = this._initialize(renderer);
        return this._initializePromise;
    }

    /**
     * @param {number} frame
     * @param {Scene} scene
     * @param {WebGLRenderer} renderer
     */
    update(frame, scene, renderer) {
        if (!this._initializePromise) {
            this.initialize(renderer);
            return;
        }
        if (!this._isInitialized) {
            return;
        }
        if (this._lastUpdateFrame === frame) {
            // Already updated this frame
            return;
        }
        this._lastUpdateFrame = frame;
        this.updateLighting(false);
        const textures = this._getTextures(scene.environment);

        if (debug && !this["_debug"]) {
            this["_debug"] = true;
            // Show both of them on cubes in the scene
            const unlitMat = new MeshBasicMaterial();
            unlitMat.side = 2;
            const radianceMat = unlitMat.clone();
            radianceMat.map = textures.radianceTexture;
            const planeGeometry = new PlaneGeometry(1, 1, 1, 1)
            const radianceCube = new Mesh(planeGeometry, radianceMat);
            const irradianceMat = unlitMat.clone();
            irradianceMat.map = textures.irradianceTexture;
            const irradianceCube = new Mesh(planeGeometry, irradianceMat);
            scene.add(radianceCube);
            scene.add(irradianceCube);
            radianceCube.name = "MaterialXRadianceCube";
            radianceCube.position.set(.8, 1, .01);
            radianceCube.scale.set(1.5, 1, 1);
            irradianceCube.name = "MaterialXIrradianceCube";
            irradianceCube.position.set(-.8, 1, -.01);
            irradianceCube.scale.set(1.5, 0.98, 1);
            console.log("[MaterialX] environment initialized from Needle context", { textures, radianceCube, irradianceCube });
        }
    }

    // Reset the environment to allow re-initialization
    reset() {
        if (debug) console.log("[MaterialX] Resetting environment");
        this._initializePromise = null;
        this._isInitialized = false;
        this._lastUpdateFrame = -1;
        this._lights = [];
        this._lightData = null;
        this._lightCount = 0;
        this._pmremGenerator?.dispose();
        this._pmremGenerator = null;
        this._renderer = null;
        for (const textureSet of this._texturesCache.values()) {
            textureSet.radianceTexture?.dispose();
            textureSet.irradianceTexture?.dispose();
        }
        this._texturesCache.clear();
    }

    get lights() {
        return this._lights;
    }
    get lightData() {
        return this._lightData;
    }
    get lightCount() { return this._lightCount || 0; }

    /**
     * @param {import("./materialx.material.js").MaterialXMaterial} material
     */
    getTextures(material) {
        if (material.envMap) {
            // If the material has its own envMap, we don't use the irradiance texture
            return this._getTextures(material.envMap);
        }

        // Use the scene background for lighting if no environment is available
        // If we don't do this we don't see the correct lighting for scenes exported with 'Environment Lighting: Color' and 'Environment Reflections: Skybox'  
        const skybox = this._scene.environment || this._scene.background;
        if (skybox instanceof Texture) {
            return this._getTextures(skybox);
        }
        return this._getTextures(null);
    }

    /** @type {PMREMGenerator | null} */
    _pmremGenerator = null;
    /** @type {WebGLRenderer | null} */
    _renderer = null;
    /** @type {Map<Texture | null, EnvironmentTextureSet>} */
    _texturesCache = new Map();

    /**
     * @param {WebGLRenderer} renderer
     * @returns {Promise<boolean>}
     */
    async _initialize(renderer) {
        this._isInitialized = false;
        this._pmremGenerator = new PMREMGenerator(renderer);
        this._renderer = renderer;
        this.updateLighting(true);
        this._isInitialized = true;
        return true;
    }

    /**
     * @param {Texture | null | undefined} texture
     * @returns {{radianceTexture: Texture | null, irradianceTexture: Texture | null}}
     */
    _getTextures(texture) {

        // Fallback to white texture if no texture is provided
        if (!texture) {
            texture = whiteTexture;
        }

        /** @type {EnvironmentTextureSet | undefined} */
        let res = this._texturesCache.get(texture || null);
        if (res) {
            return res;
        }

        if (this._scene && this._pmremGenerator && this._renderer && texture) {
            if (debug) console.log("[MaterialX] Generating environment textures", texture.name);
            const target = this._pmremGenerator.fromEquirectangular(texture);
            const radianceRenderTarget = renderPMREMToEquirect(this._renderer, target.texture, 0.0, 1024, 512, target.height);
            const irradianceRenderTarget = renderPMREMToEquirect(this._renderer, target.texture, 1.0, 32, 16, target.height);
            target.dispose();
            res = {
                radianceTexture: radianceRenderTarget.texture,
                irradianceTexture: irradianceRenderTarget.texture
            }
        }
        else {
            res = {
                radianceTexture: null,
                irradianceTexture: null
            }
        }
        this._texturesCache.set(texture || null, res);
        return res;
    }

    /**
     * @param {boolean} collectLights
     */
    updateLighting = (collectLights = false) => {
        if (!this._scene) return;
        // Find lights in scene
        if (collectLights) {
            /** @type {Array<Light>} */
            const lights = new Array();
            this._scene.traverse((/** @type {Object3D} */ object) => {
                if ((/** @type {Light} */ (object)).isLight && object.visible)
                    lights.push(/** @type {Light} */(object));
            });
            this._lights = lights;
        }

        if (state.materialXGenContext) {
            const { lightData, lightCount } = getLightData(this._lights, state.materialXGenContext);
            this._lightData = lightData;
            this._lightCount = lightCount;
        }
    }
}
