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

/**
 * @import { MaterialX_root_extension, MaterialX_material_extension, MaterialXLoaderOptions } from "./loader.three.d.ts"
 */

/**
 * @typedef {Object} MaterialDefinition
 * @property {string} [name] - Optional name for the material
 * @property {boolean} [doubleSided] - Whether the material is double-sided
 * @property {Object<string, any>} [extensions] - Extensions for the material, including MaterialX
 */

/**
 * @typedef {Object} MaterialXMaterialOptions
 * @property {import('three').MaterialParameters} [parameters]
 */

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

    /** @type {MaterialXMaterial[]} */
    _generatedMaterials = [];

    /** @type {Promise<any> | null} */
    _documentReadyPromise = null;

    /**
     * @returns {MaterialX_root_extension | null}
     */
    get materialX_root_data() {
        const ext = this.parser.json.extensions?.[this.name];
        let result = null;
        if ("documents" in ext && Array.isArray(ext.documents))
            result = ext.documents;
        else
            result = [ext];
        return result;
    }

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

    /**
     * MaterialXLoader constructor
     * @param {import('three/examples/jsm/loaders/GLTFLoader.js').GLTFParser} parser - The GLTFParser instance
     * @param {MaterialXLoaderOptions} options - The loader options
     * @param {import('../materialx.js').MaterialXContext} context - The context for the GLTF loading process
     */
    constructor(parser, options, context) {
        this.parser = parser;
        this.options = options;
        this.context = context;
        
        if (debug) console.log("MaterialXLoader created for parser");
        // Start loading of MaterialX environment if the root extension exists
        if (this.materialX_root_data) {
            ready();
        }
    }

    /**
     * @param {number} materialIndex
     * @returns {Promise<Material> | null}
     */
    loadMaterial(materialIndex) {
        const materialDef = this.parser.json.materials?.[materialIndex];
        if (!materialDef?.extensions?.[this.name]) {
            return null;
        }
        // Wrap the async implementation
        return this._loadMaterialAsync(materialIndex);
    }

    /**
     * @private
     * @param {number} materialIndex
     * @returns {Promise<Material>}
     */
    async _loadMaterialAsync(materialIndex) {

        /** @type {MaterialDefinition} */
        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
        /** @type {MaterialX_material_extension} */
        const ext = materialDef.extensions?.[this.name];
        const documentIndex = ext.document || 0;
        const materialX_root_data = this.materialX_root_data?.[documentIndex];
        const mtlx = materialX_root_data.mtlx || null;

        if (ext && mtlx) {

            /** @type {MaterialXMaterialOptions} */
            const materialOptions = {
                ...this.options,
            }

            if (!materialOptions.parameters) materialOptions.parameters = {};

            if (materialOptions.parameters?.side === undefined && materialDef.doubleSided !== undefined) {
                materialOptions.parameters.side = materialDef.doubleSided ? DoubleSide : FrontSide;
            }

            return createMaterialXMaterial(mtlx, ext.name, {
                cacheKey: this.options.cacheKey || "",
                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
                    if (materialX_root_data) {
                        const textures = materialX_root_data.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);
                    }
                    return null;
                }
            }, materialOptions, this.context)
                // Cache and return the generated material
                .then(mat => {
                    if (mat instanceof MaterialXMaterial) this._generatedMaterials.push(mat);
                    return mat;
                })
        }

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

/**
 * Add the MaterialXLoader to the GLTFLoader instance.
 * @param {GLTFLoader} loader
 * @param {MaterialXLoaderOptions} [options]
 * @param {import('../materialx.js').MaterialXContext} [context]
 */
export function useNeedleMaterialX(loader, options, context) {
    loader.register(p => {
        const loader = new MaterialXLoader(p, options || {}, context || {});
        return loader;
    });
}

/**
 * Parse the MaterialX document once and cache it
 * @param {string} mtlx
 * @returns {Promise<any>}
 */
async function load(mtlx) {
    // 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
    await state.materialXModule.readFromXmlString(doc, mtlx, "");
    if (debug) console.log("[MaterialX] root document parsed successfully");
    return doc;
}

/**
 * @param {string} mtlx
 * @param {string} materialNodeName
 * @param {import('../materialx.helper.js').Callbacks} loaders
 * @param {MaterialXLoaderOptions} [options]
 * @param {import('../materialx.js').MaterialXContext} [context]
 * @returns {Promise<Material>}
 */
export async function createMaterialXMaterial(mtlx, materialNodeName, loaders, options, context) {
    try {
        if (debug) console.log(`Creating MaterialX material: ${materialNodeName}`);

        const doc = await load(mtlx);

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

        // Find the renderable element following MaterialX example pattern exactly
        let renderableElement = 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
        for (let i = 0; i < materialNodes.length; ++i) {
            const materialNode = materialNodes[i];
            if (materialNode) {
                const name = materialNode.getNamePath();
                if (debug) console.log(`[MaterialX] Scan material[${i}]: ${name}`);

                // Find the matching material
                if (materialNodes.length === 1 || name == materialNodeName) {
                    materialNodeName = name;
                    renderableElement = materialNode;
                    foundRenderable = true;
                    if (debug) console.log(`[MaterialX] Use material node: '${name}'`);
                    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 (${materialNodeName})`);
            const fallbackMaterial = new MeshStandardMaterial();
            fallbackMaterial.color.set(0xff00ff);
            fallbackMaterial.name = `MaterialX_NoRenderable_${materialNodeName}`;
            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.getNamePath ? renderableElement.getNamePath() : renderableElement.getName();

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

        const shaderMaterial = new MaterialXMaterial({
            name: materialNodeName,
            shaderName: null, //shaderInfo?.originalName || shaderInfo?.name || null,
            shader,
            context: context || {},
            parameters: {
                transparent: isTransparent,
                ...options?.parameters,
            },
            loaders: loaders,
        });

        // 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 (${materialNodeName}):`, error);
        // Return a fallback material with stored MaterialX data
        const fallbackMaterial = new MeshStandardMaterial();
        fallbackMaterial.color.set(0xff00ff);
        fallbackMaterial.name = `MaterialX_Error_${materialNodeName}`;
        return fallbackMaterial;
    }
}
