import type { MaterialX as MX } from "./materialx.types.js";
import MaterialX from "../bin/JsMaterialXGenShader.js";
import { debug } from "./utils.js";
import { renderPMREMToEquirect } from "./textureHelper.js";
import { Light, Mesh, MeshBasicMaterial, Object3D, PlaneGeometry, PMREMGenerator, Scene, Texture, WebGLRenderer } from "three";
import { registerLights, getLightData, LightData } from "./materialx.helper.js";
import type { MaterialXMaterial } from "./materialx.material.js";

export type MaterialXContext = {
    getTime?(): number,
    getFrame?(): number,
}


export const state = new class {
    materialXModule: MX.MODULE | null = null;
    materialXGenerator: any = null;
    materialXGenContext: any = null;
    materialXStdLib: any = null;
    materialXInitPromise: Promise<void> | null = null;
}


/** 
 * Wait for the MaterialX WASM module to be ready.
 */
export async function ready(): Promise<void> {
    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 {

            const urls: Array<string> = await Promise.all([
                /** @ts-ignore */
                import(`../bin/JsMaterialXCore.wasm?url`).then(m => m.default || m),
                /** @ts-ignore */
                import(`../bin/JsMaterialXGenShader.wasm?url`).then(m => m.default || m),
                /** @ts-ignore */
                import(`../bin/JsMaterialXGenShader.data.txt?url`).then(m => m.default || m),
            ]);
            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("[MaterialX] module loaded", module);
            state.materialXModule = module as MX.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;

            // 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;
        }
    })();
}

type EnvironmentTextureSet = {
    radianceTexture: Texture | null;
    irradianceTexture: Texture | null;
}


/**
 * MaterialXEnvironment manages the environment settings for MaterialX materials.
 */
export class MaterialXEnvironment {

    static get(scene: Scene): MaterialXEnvironment | null {
        return this.getEnvironment(scene);
    }
    private static _environments: WeakMap<Scene, MaterialXEnvironment> = new Map();
    private static getEnvironment(scene: Scene): MaterialXEnvironment {
        if (this._environments.has(scene)) {
            return this._environments.get(scene)!;
        }
        const env = new MaterialXEnvironment(scene);
        this._environments.set(scene, env);
        return env;
    }


    private _lights: Array<Light> = [];
    private _lightData: null | LightData[] = null;
    private _lightCount: number = 0;

    private _initializePromise: Promise<boolean> | null = null;
    private _isInitialized: boolean = false;
    private _lastUpdateFrame: number = -1;

    constructor(private _scene: Scene) {
        if (debug) console.log("[MaterialX] Environment created");
    }

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

    update(frame: number, scene: Scene, renderer: WebGLRenderer): void {
        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, radianceMat);
            scene.add(radianceCube);
            scene.add(irradianceCube);
            radianceCube.name = "MaterialXRadianceCube";
            radianceCube.position.set(.6, 1, 0);
            irradianceCube.name = "MaterialXIrradianceCube";
            irradianceCube.position.set(-.6, 1, 0);
            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 lightData() {
        return this._lightData;
    }
    get lightCount() { return this._lightCount || 0; }
    getTextures(material: MaterialXMaterial) {
        if (material.envMap) {
            // If the material has its own envMap, we don't use the irradiance texture
            return this._getTextures(material.envMap);
        }
        return this._getTextures(this._scene.environment);
    }

    private _pmremGenerator: PMREMGenerator | null = null;
    private _renderer: WebGLRenderer | null = null;
    private readonly _texturesCache: Map<Texture | null, EnvironmentTextureSet> = new Map();

    private async _initialize(renderer: WebGLRenderer): Promise<boolean> {
        this._isInitialized = false;
        this._pmremGenerator = new PMREMGenerator(renderer);
        this._renderer = renderer;
        this.updateLighting(true);
        this._isInitialized = true;
        return true;
    }

    private _getTextures(texture: Texture | null | undefined): {
        radianceTexture: Texture | null,
        irradianceTexture: Texture | null
    } {
        let res: EnvironmentTextureSet | undefined = 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;
    }

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

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