import { Context } from "@needle-tools/engine";
import { RawShaderMaterial, Material, MeshStandardMaterial, LoadingManager, TextureLoader, Texture, NearestFilter, Matrix4, GLSL3, AddEquation, OneMinusSrcAlphaFactor, SrcAlphaFactor, DoubleSide, Matrix3, Vector3, Object3D, Camera } from "three";
import { GLTFLoaderPlugin, GLTFParser } from "three/examples/jsm/loaders/GLTFLoader.js";
import { getUniformValues } from "../helper.js";
import { initializeMaterialX, MaterialXEnvironment, state } from "../materialx.js";
import { debug } from "../utils.js";

// TypeScript interfaces matching the C# data structures
interface MaterialXData {
    version: string; // e.g. "1.39"
    name: string;    // e.g. "Material"
    mtlx: string;    // MaterialX XML content
}

interface MaterialXDataIndex {
    name: string;    // Material name reference
}

// MaterialX loader extension for js GLTFLoader
export class MaterialXLoader implements GLTFLoaderPlugin {
    name = "NEEDLE_materials_mtlx";

    private rootMaterialXData: MaterialXData | null = null;
    private parsedDocument: any = null;
    private documentParsePromise: Promise<any> | null = null;
    private rootDataInitialized = false;
    private environmentInitialized = false;
    private generatedMaterials: RawShaderMaterial[] = [];

    constructor(private parser: GLTFParser, private context: Context) {
        if (debug) console.log("MaterialXLoader created for parser");
        // Initialize MaterialX environment after MaterialX is ready
        this.initializeEnvironment();
    }

    // Initialize MaterialX environment - called once after MaterialX is ready
    private async initializeEnvironment(): Promise<void> {
        if (this.environmentInitialized) return;

        if (debug) console.log("MaterialXLoader: Initializing MaterialX environment...");

        // Ensure MaterialX is initialized first
        await initializeMaterialX();

        // Set up environment with context
        const environment = state.materialXEnvironment;
        environment.setContext(this.context);

        // Initialize the environment from context (properly awaited)
        try {
            await environment.initializeFromContext();
            this.environmentInitialized = true;
            if (debug) console.log("MaterialXLoader: Environment initialized successfully");
        } catch (error) {
            console.warn("MaterialXLoader: Failed to initialize MaterialX environment:", error);
        }
    }

    // Initialize root data from parser.json.extensions once
    private initializeRootData(): void {
        if (this.rootDataInitialized) return;

        const gltfExtensions = this.parser.json.extensions;
        if (gltfExtensions?.[this.name]) {
            if (debug) console.log("MaterialX extension found in root:", gltfExtensions[this.name]);

            const materialXExtension = gltfExtensions[this.name];
            this.rootMaterialXData = materialXExtension as MaterialXData;
        }
        this.rootDataInitialized = true;
    }

    // Parse the MaterialX document once and cache it
    private async parseRootDocument(): Promise<any> {
        if (this.documentParsePromise) {
            return this.documentParsePromise;
        }

        this.documentParsePromise = (async () => {
            if (this.parsedDocument) return this.parsedDocument;

            if (debug) console.log("Parsing MaterialX root document...");

            // Ensure MaterialX is initialized
            await initializeMaterialX();

            if (!state.materialXModule) {
                throw new Error("MaterialX module failed to initialize");
            }

            // Create MaterialX document and parse ALL the XML data from root
            const doc = state.materialXModule.createDocument();
            doc.setDataLibrary(state.materialXStdLib);

            // Parse all MaterialX XML strings from the root data
            if (this.rootMaterialXData) {
                if (debug) console.log(`Parsing MaterialX XML for: ${this.rootMaterialXData.name}`);
                await state.materialXModule.readFromXmlString(doc, this.rootMaterialXData.mtlx, "");
            }

            if (debug) console.log("MaterialX root document parsed successfully");

            this.parsedDocument = doc;
            return doc;
        })();

        return this.documentParsePromise;
    }

    loadMaterial(materialIndex: number): Promise<Material> | null {
        const materialDef = this.parser.json.materials?.[materialIndex];
        if (!materialDef?.extensions?.[this.name]) {
            return null;
        }

        // Wrap the async implementation
        return this._loadMaterialAsync(materialIndex);
    }

    private async _loadMaterialAsync(materialIndex: number): Promise<Material> {
        // Initialize root data first
        this.initializeRootData();

        const materialDef = this.parser.json.materials?.[materialIndex];
        if (debug) console.log("MaterialX extension found in material:", materialDef.extensions[this.name]);

        // Handle different types of MaterialX data
        const dataIndex = materialDef.extensions[this.name] as MaterialXDataIndex;

        if (dataIndex) {
            // Create a new material and process MaterialX - AWAIT THIS!
            return await this.createMaterialXMaterial(dataIndex);
        }

        // Return fallback material instead of null
        const fallbackMaterial = new MeshStandardMaterial();
        fallbackMaterial.name = "MaterialX_Fallback";
        return fallbackMaterial;
    }

    private rootDocument: Promise<any> | null = null;
    private async createMaterialXMaterial(materialXData: MaterialXDataIndex): Promise<Material> {
        try {
            if (debug) console.log(`Creating MaterialX material: ${materialXData.name}`);

            // Ensure MaterialX is initialized and document is parsed
            await initializeMaterialX();

            if (!this.rootDocument) {
                this.rootDocument = this.parseRootDocument();
            }
            const doc = await this.rootDocument;

            if (!state.materialXModule || !state.materialXGenerator || !state.materialXGenContext) {
                console.warn("MaterialX WASM module not ready, returning fallback material");
                const fallbackMaterial = new MeshStandardMaterial();
                fallbackMaterial.userData.materialX = materialXData;
                fallbackMaterial.name = `MaterialX_Fallback_${materialXData.name}`;
                return fallbackMaterial;
            }

            // Find the renderable element following MaterialX example pattern exactly
            let renderableElement = null;
            let foundRenderable = false;

            if (debug) console.log("Mtlx doc", doc);

            // Search for material nodes first (following the reference pattern)
            const materialNodes = doc.getMaterialNodes();
            if (debug) console.log(`Found ${materialNodes.length} material nodes in document`, materialNodes);

            // Handle both array and vector-like APIs
            const materialNodesLength = materialNodes.length;
            for (let i = 0; i < materialNodesLength; ++i) {
                const materialNode = materialNodes[i];
                if (materialNode) {
                    const materialName = materialNode.getNamePath();
                    if (debug) console.log('Scan material: ', i, materialName);

                    // Find the matching material
                    if (materialName == materialXData.name) {
                        renderableElement = materialNode;
                        foundRenderable = true;
                        if (debug) console.log('-- add material: ', materialName);
                        break;
                    }
                }
            }

            /*
            // If no material nodes found, search nodeGraphs
            if (!foundRenderable) {
                const nodeGraphs = doc.getNodeGraphs();
                console.log(`Found ${nodeGraphs.length} node graphs in document`);
                const nodeGraphsLength = nodeGraphs.length;
                for (let i = 0; i < nodeGraphsLength; ++i) {
                    const nodeGraph = nodeGraphs[i];
                    if (nodeGraph) {
                        // Skip any nodegraph that has nodedef or sourceUri
                        if ((nodeGraph as any).hasAttribute('nodedef') || (nodeGraph as any).hasSourceUri()) {
                            continue;
                        }
                        // Skip any nodegraph that is connected to something downstream
                        if ((nodeGraph as any).getDownstreamPorts().length > 0) {
                            continue;
                        }
                        const outputs = (nodeGraph as any).getOutputs();
                        for (let j = 0; j < outputs.length; ++j) {
                            const output = outputs[j];
                            if (output && !foundRenderable) {
                                renderableElement = output;
                                foundRenderable = true;
                                break;
                            }
                        }
                        if (foundRenderable) break;
                    }
                }
            }
            
            // If still no element found, search document outputs
            if (!foundRenderable) {
                const outputs = doc.getOutputs();
                console.log(`Found ${outputs.length} output nodes in document`);
                const outputsLength = outputs.length;
                for (let i = 0; i < outputsLength; ++i) {
                    const output = outputs[i];
                    if (output && !foundRenderable) {
                        renderableElement = output;
                        foundRenderable = true;
                        break;
                    }
                }
            }
            */

            if (!renderableElement) {
                console.warn("No renderable element found in MaterialX document");
                const fallbackMaterial = new MeshStandardMaterial();
                fallbackMaterial.userData.materialX = materialXData;
                fallbackMaterial.name = `MaterialX_NoRenderable_${materialXData.name}`;
                return fallbackMaterial;
            }

            if (debug) console.log("Using renderable element for shader generation");

            // Check transparency and set context options like the reference
            let isTransparent = state.materialXModule.isTransparentSurface(renderableElement, state.materialXGenerator.getTarget());
            state.materialXGenContext.getOptions().hwTransparency = isTransparent;
            state.materialXGenContext.getOptions().shaderInterfaceType = state.materialXModule.ShaderInterfaceType.SHADER_INTERFACE_COMPLETE;
            state.materialXGenContext.getOptions().hwSrgbEncodeOutput = true; // Like the reference

            // Generate shaders using the element's name path
            if (debug) console.log("Generating MaterialX shaders...");
            const elementName = (renderableElement as any).getNamePath ? (renderableElement as any).getNamePath() : (renderableElement as any).getName();

            const shader = state.materialXGenerator.generate(elementName, renderableElement, state.materialXGenContext);

            // Get vertex and fragment shader source
            // Remove #version directive for newer js. It's added by RawShaderMaterial glslVersion.
            let vertexShader = shader.getSourceCode("vertex").replace(/^#version.*$/gm, '').trim();
            let fragmentShader = shader.getSourceCode("pixel").replace(/^#version.*$/gm, '').trim();

            // MaterialX uses different attribute names than js defaults,
            // so we patch the MaterialX shaders to match the js standard names.
            // Otherwise, we'd have to modify the mesh attributes (see below).
            vertexShader = vertexShader.replace(/\bi_position\b/g, 'position');
            vertexShader = vertexShader.replace(/\bi_normal\b/g, 'normal');
            vertexShader = vertexShader.replace(/\bi_texcoord_0\b/g, 'uv');
            vertexShader = vertexShader.replace(/\bi_texcoord_1\b/g, 'uv1');
            vertexShader = vertexShader.replace(/\bi_tangent\b/g, 'tangent');
            vertexShader = vertexShader.replace(/\bi_color_0\b/g, 'color');

            fragmentShader = fragmentShader.replace(/\bi_position\b/g, 'position');
            fragmentShader = fragmentShader.replace(/\bi_normal\b/g, 'normal');
            fragmentShader = fragmentShader.replace(/\bi_texcoord_0\b/g, 'uv');
            fragmentShader = fragmentShader.replace(/\bi_texcoord_1\b/g, 'uv1');
            fragmentShader = fragmentShader.replace(/\bi_tangent\b/g, 'tangent');
            fragmentShader = fragmentShader.replace(/\bi_color_0\b/g, 'color');

            // From the original code:

            /* PATCHING THE SHADER IS BETTER – THIS IS JUST FOR REFERENCE 
            // Use default MaterialX naming convention.
            var startStreamTime = performance.now();
            child.geometry.attributes.i_position = child.geometry.attributes.position;
            child.geometry.attributes.i_normal = child.geometry.attributes.normal;
            child.geometry.attributes.i_tangent = child.geometry.attributes.tangent;
            child.geometry.attributes.i_texcoord_0 = child.geometry.attributes.uv;
            */

            if (debug) {
                console.group("Material: ", materialXData.name);
                console.log("Vertex shader length:", vertexShader.length, vertexShader);
                console.log("Fragment shader length:", fragmentShader.length, fragmentShader);
                console.groupEnd();
            }

            // Extract uniforms from both vertex and pixel stages
            const loadingManager = new LoadingManager();
            loadingManager.setURLModifier((url: string) => {
                if (debug) console.log("Loading URL:", url);
                // Find the texture in the textures in the parser and load it from there
                return url;
            });
            let textureLoader = new TextureLoader(loadingManager);

            // Override the textureLoader.load method to use the parser's loadTexture directly,
            // since we want to load the textures from the glTF document and not from disk.
            // That loader does its own caching too.
            textureLoader.load = (url: string, onLoad?: (texture: Texture) => void, onProgress?: (event: ProgressEvent) => void, onError?: (err: Error) => void): Texture => {

                // Return a checkerboard texture for now
                const checkerboardTexture = new Texture();
                checkerboardTexture.image = new Image();
                checkerboardTexture.image.src = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAIAAAACCAIAAAD91JpzAAAAG0lEQVR4nGJqWH9q9e8XjA/VrL8UfQIEAAD//zn2CCX5UcsdAAAAAElFTkSuQmCC";
                checkerboardTexture.needsUpdate = true;
                // Pixelated filtering
                checkerboardTexture.magFilter = NearestFilter;
                checkerboardTexture.minFilter = NearestFilter;

                new Promise(() => {
                    // Find the index of the texture in the parser
                    const filenameWithoutExt = url.split('/').pop()?.split('.').shift() || '';

                    // Find in the root extension; there's the textures dictionary:
                    // a mapping from the texture name used in MaterialX to the glTF texture index
                    const ext = this.parser.json.extensions?.[this.name];
                    const textures = ext?.textures || [];

                    const index = textures.findIndex(tex => {
                        if (debug) console.log("Checking texture:", tex.name, "against URL:", filenameWithoutExt);
                        return tex.name === filenameWithoutExt;
                    });

                    if (index < 0) {
                        console.warn("Texture not found in parser:", filenameWithoutExt, this.parser.json);
                        onError?.(new Error(`Texture not found: ${filenameWithoutExt}`));
                        return;
                    }
                    this.parser.getDependency("texture", index).then(tex => {
                        if (debug) console.log("Texture loaded:", tex);
                        // update the checkerboard texture with the loaded texture
                        checkerboardTexture.image = tex.image;
                        checkerboardTexture.needsUpdate = true;
                        onLoad?.(tex);
                    }).catch(err => {
                        onError?.(err);
                    });
                });

                return checkerboardTexture;
            };

            const searchPath = ""; // Could be derived from the asset path if needed
            const flipV = false; // Set based on your geometry requirements

            let uniforms = {
                ...getUniformValues(shader.getStage('vertex'), textureLoader, searchPath, flipV),
                ...getUniformValues(shader.getStage('pixel'), textureLoader, searchPath, flipV),
            };

            // Get lighting and environment data from MaterialX environment
            const environment = state.materialXEnvironment;
            // const lights = environment.getLights() || [];
            const lightData = environment.getLightData() || null;
            const radianceTexture = environment.getRadianceTexture() || null;
            const irradianceTexture = environment.getIrradianceTexture() || null;
            const getLightRotation = () => {
                // Placeholder for light rotation logic, if needed
                return new Matrix4();
            };

            if (debug) console.log({ lightData, radianceTexture, irradianceTexture });

            Object.assign(uniforms, {
                u_numActiveLightSources: { value: lightData?.length || 0, type: 'i' },
            });

            if (lightData?.length) {
                Object.assign(uniforms, {
                    u_lightData: { value: lightData, type: 'v4v' },
                });
            }

            const mips = Math.trunc(Math.log2(Math.max(radianceTexture?.width ?? 0, radianceTexture?.height ?? 0))) + 1;
            if (debug) console.log("Radiance texture mips:", mips, "for texture size:", radianceTexture?.width, "x", radianceTexture?.height);
            Object.assign(uniforms, {
                u_envMatrix: { value: getLightRotation() },
                u_envRadiance: { value: radianceTexture, type: 't' },
                u_envRadianceMips: { value: 8, type: 'i' },
                // TODO we need to figure out how we can set a PMREM here... doing many texture samples is prohibitively expensive
                u_envRadianceSamples: { value: 8, type: 'i' },
                u_envIrradiance: { value: irradianceTexture, type: 't' },
                u_refractionEnv: { value: true },
            });

            // console.log("NEW MATERIAL UNIFORMS", uniforms);

            // Debug: Log the actual shaders to see what we're working with
            // console.log("Generated vertex shader:", vertexShader.substring(0, 500) + "...");
            // console.log("Generated fragment shader:", fragmentShader.substring(0, 500) + "...");

            // Create js RawShaderMaterial (with GLSL3 for MaterialX compatibility)
            const shaderMaterial = new RawShaderMaterial({
                uniforms: uniforms,
                vertexShader: vertexShader,
                fragmentShader: fragmentShader,
                glslVersion: GLSL3,
                transparent: isTransparent,
                // Only set blend settings if actually transparent
                ...(isTransparent && {
                    blendEquation: AddEquation,
                    blendSrc: OneMinusSrcAlphaFactor,
                    blendDst: SrcAlphaFactor,
                }),
                side: DoubleSide,
                // Add some debug settings
                depthTest: true,
                depthWrite: !isTransparent
            });

            // Add helper matrices for uniform updates (similar to MaterialX example)
            const normalMat = new Matrix3();
            const viewProjMat = new Matrix4();
            const worldViewPos = new Vector3();

            // Store helper objects on the material for uniform updates
            shaderMaterial.userData.materialX = materialXData;
            shaderMaterial.userData.updateUniforms = (object: Object3D, camera: Camera) => {
                const uniforms = shaderMaterial.uniforms;

                if (!uniforms) return;

                // TODO remove. Not sure why this is needed, but without it
                // we currently get some "swimming" where matrices are not up to date.
                camera.updateMatrixWorld(true);

                // Update standard transformation matrices
                if (uniforms.u_worldMatrix) {
                    if (!uniforms.u_worldMatrix.value?.isMatrix4) uniforms.u_worldMatrix.value = new Matrix4();
                    uniforms.u_worldMatrix.value = object.matrixWorld;
                }

                if (uniforms.u_viewProjectionMatrix) {
                    if (!uniforms.u_viewProjectionMatrix.value?.isMatrix4) uniforms.u_viewProjectionMatrix.value = new Matrix4();
                    uniforms.u_viewProjectionMatrix.value.multiplyMatrices(camera.projectionMatrix, camera.matrixWorldInverse);
                }

                if (uniforms.u_viewPosition) {
                    if (!uniforms.u_viewPosition.value?.isVector3) uniforms.u_viewPosition.value = new Vector3();
                    uniforms.u_viewPosition.value.copy(camera.getWorldPosition(worldViewPos));
                }

                if (uniforms.u_worldInverseTransposeMatrix) {
                    if (!uniforms.u_worldInverseTransposeMatrix.value?.isMatrix4) uniforms.u_worldInverseTransposeMatrix.value = new Matrix4();
                    uniforms.u_worldInverseTransposeMatrix.value.setFromMatrix3(normalMat.getNormalMatrix(object.matrixWorld));
                }

                // Update time uniforms
                const time = performance.now() / 1000.0;
                if (uniforms.u_time) {
                    uniforms.u_time.value = time;
                }
                if (uniforms.u_frame) {
                    uniforms.u_frame.value = time;
                }

                // Mark uniforms as needing update
                shaderMaterial.uniformsNeedUpdate = true;
            };

            shaderMaterial.name = `MaterialX_Generated_${materialXData.name}`;

            // Add debugging to see if the material compiles correctly
            if (debug) console.log("MaterialX material created successfully:", shaderMaterial.name);
            // if (debug) console.log("Material uniforms keys:", Object.keys(shaderMaterial.uniforms || {}));
            // if (debug) console.log("Material transparent:", shaderMaterial.transparent);
            // if (debug) console.log("Material side:", shaderMaterial.side);

            // Track this material for later lighting updates
            this.generatedMaterials.push(shaderMaterial);

            return shaderMaterial;

        } catch (error) {
            // This is a wasm error (an int) that we need to resolve
            console.error("Error creating MaterialX material:", error);
            // Return a fallback material with stored MaterialX data
            const fallbackMaterial = new MeshStandardMaterial();
            fallbackMaterial.userData.materialX = materialXData;
            fallbackMaterial.name = `MaterialX_Error_${materialXData.name}`;
            return fallbackMaterial;
        }
    }

    private identityMatrix: Matrix4 = new Matrix4().identity();
    // Update lighting uniforms for all generated MaterialX materials
    updateLightingFromEnvironment(environment: MaterialXEnvironment): void {

        // Get lighting data from environment
        // const lights = environment.getLights() || [];
        const lightData = environment.getLightData() || null;
        const radianceTexture = environment.getRadianceTexture() || null;
        const irradianceTexture = environment.getIrradianceTexture() || null;

        if (debug) console.log(`Updating lighting for ${this.generatedMaterials.length} MaterialX materials`, {
            lightData, radianceTexture, irradianceTexture,
        });

        // Update each generated material's lighting uniforms
        this.generatedMaterials.forEach((material, _index) => {
            if (!material.uniforms) return;

            // Update light count
            if (material.uniforms.u_numActiveLightSources && lightData) {
                material.uniforms.u_numActiveLightSources.value = lightData.length;
            }

            // Update light data if we have lights
            if (lightData) {
                if (!material.uniforms.u_lightData) {
                    material.uniforms.u_lightData = { value: null };
                }
                material.uniforms.u_lightData.value = lightData;
            }

            // Update environment uniforms
            if (material.uniforms.u_envMatrix) {
                material.uniforms.u_envMatrix.value = this.identityMatrix;
            }
            if (material.uniforms.u_envRadiance) {
                material.uniforms.u_envRadiance.value = radianceTexture;
            }
            if (material.uniforms.u_envRadianceMips) {
                material.uniforms.u_envRadianceMips.value = Math.trunc(Math.log2(Math.max(radianceTexture?.source.data.width ?? 0, radianceTexture?.source.data.height ?? 0))) + 1;
            }
            if (material.uniforms.u_envIrradiance) {
                material.uniforms.u_envIrradiance.value = irradianceTexture;
            }

            // Mark uniforms as needing update
            // console.log("Light data in uniforms", material.uniforms, material.fragmentShader);
            material.uniformsNeedUpdate = true;
        });

        /* TODO not working yet
        this.generatedMaterials.forEach((material, index) => {
            console.log("Regenerating shaders for MaterialX material:", index, material.name);
            material.userData.regenerateShaders();
            resetShaders(this.context);
        });
        */
    }
}
