import { Material, MeshStandardMaterial, DoubleSide, FrontSide } from "three";
import { GLTFLoaderPlugin, GLTFParser } from "three/examples/jsm/loaders/GLTFLoader.js";
import { ready, state, MaterialXContext } from "../materialx.js";
import { debug } from "../utils.js";
import { MaterialXMaterial } from "../materialx.material.js";

// TypeScript interfaces matching the C# data structures
export type MaterialX_root_extension = {
    /** e.g. 1.39 */
    version: string;
    /** e.g. "Material" */
    name: string;
    /** MaterialX xml content */
    mtlx: string;
    /** MaterialX texture pointers */
    textures: Array<{ name: string, pointer: string }>;
}

export type MaterialX_material_extension = {
    name: string; // Material name reference
}

type MaterialDefinition = {
    name?: string; // Optional name for the material
    doubleSided?: boolean; // Whether the material is double-sided
    extensions?: {
        [key: string]: any; // Extensions for the material, including MaterialX
    },
}

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

    private readonly _generatedMaterials: MaterialXMaterial[] = [];

    private _documentReadyPromise: Promise<any> | null = null;

    get materialX_root_data() {
        return this.parser.json.extensions?.[this.name] as MaterialX_root_extension | null;
    }

    /** Generated materialX materials */
    get materials() {
        return this._generatedMaterials;
    }

    /**
     * MaterialXLoader constructor
     * @param parser The GLTFParser instance
     * @param url The URL of the GLTF file
     * @param context The context for the GLTF loading process
     */
    constructor(private parser: GLTFParser, private url: string, private context: MaterialXContext) {
        if (debug) console.log("MaterialXLoader created for parser");
        // Start loading of MaterialX environment if the root extension exists
        if (this.materialX_root_data) {
            ready();
        }
    }

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

    // Parse the MaterialX document once and cache it
    private async _materialXDocumentReady(): Promise<any> {
        if (this._documentReadyPromise) {
            return this._documentReadyPromise;
        }
        return this._documentReadyPromise = (async () => {
            if (debug) console.log("[MaterialX] Parsing MaterialX root document...");

            // Ensure MaterialX is initialized
            await ready();

            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
            const root = this.materialX_root_data
            if (root) {
                if (debug) console.log(`[MaterialX] Parsing XML for: ${root.name}`);
                await state.materialXModule.readFromXmlString(doc, root.mtlx, "");
            }

            if (debug) console.log("[MaterialX] root document parsed successfully");
            return doc;
        })();
    }

    private async _loadMaterialAsync(materialIndex: number): Promise<Material> {

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

        // Handle different types of MaterialX data
        const ext = materialDef.extensions?.[this.name] as MaterialX_material_extension;

        if (ext) {
            return this._createMaterialXMaterial(materialDef, ext);
        }

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

    private async _createMaterialXMaterial(material_def: MaterialDefinition, material_extension: MaterialX_material_extension): Promise<Material> {
        try {
            if (debug) console.log(`Creating MaterialX material: ${material_extension.name}`);

            const doc = await this._materialXDocumentReady();

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

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

            if (debug) console.log("[MaterialX] document", doc);

            // Search for material nodes first (following the reference pattern)
            const materialNodes = doc.getMaterialNodes();
            if (debug) console.log(`[MaterialX] 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('[MaterialX] Scan material: ', i, materialName);

                    // Find the matching material
                    if (materialName == material_extension.name) {
                        renderableElement = materialNode;
                        foundRenderable = true;
                        if (debug) console.log('[MaterialX] -- 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(`[MaterialX] No renderable element found in MaterialX document (${material_extension.name})`);
                const fallbackMaterial = new MeshStandardMaterial();
                fallbackMaterial.color.set(0xff00ff);
                fallbackMaterial.userData.materialX = material_extension;
                fallbackMaterial.name = `MaterialX_NoRenderable_${material_extension.name}`;
                return fallbackMaterial;
            }

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

            // Check transparency and set context options like the reference
            const isTransparent = state.materialXModule.isTransparentSurface(renderableElement, state.materialXGenerator.getTarget());
            state.materialXGenContext.getOptions().hwTransparency = isTransparent;

            // Generate shaders using the element's name path
            if (debug) console.log("[MaterialX] 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);
            const shaderMaterial = new MaterialXMaterial({
                name: material_extension.name,
                shader,
                transparent: isTransparent,
                side: material_def.doubleSided ? DoubleSide : FrontSide,
                context: this.context,
                loaders: {
                    cacheKey: this.url,
                    getTexture: async url => {
                        // Find the index of the texture in the parser
                        const filenameWithoutExt = url.split('/').pop()?.split('.').shift() || '';

                        // Resolve the texture from the MaterialX root extension
                        const ext = this.materialX_root_data;
                        if (ext) {
                            const textures = ext.textures || [];
                            let index = -1;
                            for (const texture of textures) {
                                // Find the texture by name and use the pointer string to get the index
                                if (texture.name === filenameWithoutExt) {
                                    const ptr = texture.pointer;
                                    const indexStr = ptr.substring("/textures/".length);
                                    index = parseInt(indexStr);

                                    if (isNaN(index) || index < 0) {
                                        console.error("[MaterialX] Invalid texture index in pointer:", ptr);
                                        return;
                                    }
                                    else {
                                        if (debug) console.log("[MaterialX] Texture index found:", index, "for", filenameWithoutExt);
                                    }
                                }
                            }

                            if (index < 0) {
                                console.error("[MaterialX] Texture not found in parser:", filenameWithoutExt, this.parser.json);
                                return;
                            }
                            return this.parser.getDependency("texture", index).then(tex => {
                                if (debug) console.log("[MaterialX] Texture loaded:" + tex.name, tex);
                                return tex;
                            });
                        }
                        return null;
                    }
                }
            });
            // Track this material for later lighting updates
            this._generatedMaterials.push(shaderMaterial);

            // Add debugging to see if the material compiles correctly
            if (debug) console.log("[MaterialX]  material created:", shaderMaterial.name);
            return shaderMaterial;

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