//
// Copyright Contributors to the MaterialX Project
// SPDX-License-Identifier: Apache-2.0
//

import { getParam } from '@needle-tools/engine';
import * as THREE from 'three';

const debug = getParam("debugmaterialx");

const IMAGE_PROPERTY_SEPARATOR = "_";
const UADDRESS_MODE_SUFFIX = IMAGE_PROPERTY_SEPARATOR + "uaddressmode";
const VADDRESS_MODE_SUFFIX = IMAGE_PROPERTY_SEPARATOR + "vaddressmode";
const FILTER_TYPE_SUFFIX = IMAGE_PROPERTY_SEPARATOR + "filtertype";
const IMAGE_PATH_SEPARATOR = "/";

/**
 * Initialized the environment texture as MaterialX expects it
 * @param {THREE.Texture} texture
 * @param {Object} capabilities
 * @returns {THREE.Texture}
 */
export function prepareEnvTexture(texture, capabilities)
{
    let newTexture = new THREE.DataTexture(texture.image.data, texture.image.width, texture.image.height, texture.format, texture.type);
    newTexture.wrapS = THREE.RepeatWrapping;
    newTexture.anisotropy = capabilities.getMaxAnisotropy();
    newTexture.minFilter = THREE.LinearMipmapLinearFilter;
    newTexture.magFilter = THREE.LinearFilter;
    newTexture.generateMipmaps = true;
    newTexture.needsUpdate = true;

    return newTexture;
}

/**
 * Get Three uniform from MaterialX vector
 * @param {any} value
 * @param {any} dimension
 * @returns {THREE.Uniform}
 */
function fromVector(value, dimension)
{
    let outValue;
    if (value)
    {
        outValue = [...value.data()];
    }
    else
    {
        outValue = [];
        for (let i = 0; i < dimension; ++i)
            outValue.push(0.0);
    }

    return outValue;
}

/**
 * Get Three uniform from MaterialX matrix
 * @param {mx.matrix} matrix
 * @param {mx.matrix.size} dimension
 */
function fromMatrix(matrix, dimension)
{
    let vec = [];
    if (matrix)
    {
        for (let i = 0; i < matrix.numRows(); ++i)
        {
            for (let k = 0; k < matrix.numColumns(); ++k)
            {
                vec.push(matrix.getItem(i, k));
            }
        }
    } else
    {
        for (let i = 0; i < dimension; ++i)
            vec.push(0.0);
    }

    return vec;
}

/**
 * Get Three uniform from MaterialX value
 * @param {mx.Uniform.type} type
 * @param {mx.Uniform.value} value
 * @param {mx.Uniform.name} name
 * @param {mx.Uniforms} uniforms
 * @param {THREE.textureLoader} textureLoader
 * @param {string} searchPath
 * @param {boolean} flipY
 */
function toThreeUniform(type, value, name, uniforms, textureLoader, searchPath, flipY)
{
    let outValue = null;
    switch (type)
    {
        case 'float':
        case 'integer':
        case 'boolean':
            outValue = value;
            break;
        case 'vector2':
            outValue = fromVector(value, 2);
            break;
        case 'vector3':
        case 'color3':
            outValue = fromVector(value, 3);
            break;
        case 'vector4':
        case 'color4':
            outValue = fromVector(value, 4);
            break;
        case 'matrix33':
            outValue = fromMatrix(value, 9);
            break;
        case 'matrix44':
            outValue = fromMatrix(value, 16);
            break;
        case 'filename':
            if (value)
            {
                // Cache / reuse texture to avoid reload overhead.
                // Note: that data blobs and embedded data textures are not cached as they are transient data.
                let checkCache = false;
                let texturePath = searchPath + IMAGE_PATH_SEPARATOR + value;
                if (value.startsWith('blob:'))   
                {
                    texturePath = value;
                    if (debug) console.log('Load blob URL:', texturePath);
                    checkCache = false;
                }
                else if (value.startsWith('http'))
                {
                    texturePath = value;
                    if (debug) console.log('Load HTTP URL:', texturePath);
                }
                else if (value.startsWith('data:'))
                {
                    texturePath = value;
                    checkCache = false;
                    if (debug) console.log('Load data URL:', texturePath);
                }
                const cachedTexture = checkCache && THREE.Cache.get(texturePath);
                if (cachedTexture)
                {
                    // Get texture from cache
                    outValue = cachedTexture;
                    if (debug) console.log('Use cached texture: ', texturePath, outValue);
                }
                else
                {
                    outValue = textureLoader.load(
                        texturePath, 
                        function (texture) {
                            if (debug) console.log('Load new texture: ' + texturePath, texture);
                            outValue = texture;

                            // Add texture to ThreeJS cache
                            if (checkCache)
                                THREE.Cache.add(texturePath, texture);
                        },
                        undefined,
                        function (error) {
                            console.error('Error loading texture: ', error);
                        });     

                    // Set address & filtering mode
                    if (outValue)
                        setTextureParameters(outValue, name, uniforms, flipY);
                }
            }
            break;
        case 'samplerCube':
        case 'string':
            break;
        default:
            const key = type + ':' + name;
            if (!valueTypeWarningMap.has(key))
            {
                valueTypeWarningMap.set(key, true);
                console.warn('MaterialX: Unsupported uniform type: ' + type + ' for uniform: ' + name, value);
            }
            outValue = null;
    }

    return outValue;
}

const valueTypeWarningMap = new Map();

/**
 * Get Three wrapping mode
 * @param {mx.TextureFilter.wrap} mode
 * @returns {THREE.Wrapping}
 */
function getWrapping(mode)
{
    let wrap;
    switch (mode)
    {
        case 1:
            wrap = THREE.ClampToEdgeWrapping;
            break;
        case 2:
            wrap = THREE.RepeatWrapping;
            break;
        case 3:
            wrap = THREE.MirroredRepeatWrapping;
            break;
        default:
            wrap = THREE.RepeatWrapping;
            break;
    }
    return wrap;
}

/**
 * Get Three minification filter
 * @param {mx.TextureFilter.minFilter} type
 * @param {mx.TextureFilter.generateMipmaps} generateMipmaps
 */
function getMinFilter(type, generateMipmaps)
{
    /** @type {THREE.TextureFilter} */
    let filterType = generateMipmaps ? THREE.LinearMipMapLinearFilter : THREE.LinearFilter;
    if (type === 0)
    {
        filterType = generateMipmaps ? THREE.NearestMipMapNearestFilter : THREE.NearestFilter;
    }
    return filterType;
}

/**
 * Set Three texture parameters
 * @param {THREE.Texture} texture
 * @param {mx.Uniform.name} name
 * @param {mx.Uniforms} uniforms
 * @param {mx.TextureFilter.generateMipmaps} generateMipmaps
 */
function setTextureParameters(texture, name, uniforms, flipY = true, generateMipmaps = true)
{
    const idx = name.lastIndexOf(IMAGE_PROPERTY_SEPARATOR);
    const base = name.substring(0, idx) || name;

    texture.generateMipmaps = generateMipmaps;
    texture.wrapS = THREE.RepeatWrapping;
    texture.wrapT = THREE.RepeatWrapping;
    texture.magFilter = THREE.LinearFilter;
    texture.flipY = flipY;

    if (uniforms.find(base + UADDRESS_MODE_SUFFIX))
    {
        const uaddressmode = uniforms.find(base + UADDRESS_MODE_SUFFIX).getValue().getData();
        texture.wrapS = getWrapping(uaddressmode);
    }

    if (uniforms.find(base + VADDRESS_MODE_SUFFIX))
    {
        const vaddressmode = uniforms.find(base + VADDRESS_MODE_SUFFIX).getValue().getData();
        texture.wrapT = getWrapping(vaddressmode);
    }

    const filterType = uniforms.find(base + FILTER_TYPE_SUFFIX) ? uniforms.get(base + FILTER_TYPE_SUFFIX).value : -1;
    texture.minFilter = getMinFilter(filterType, generateMipmaps);
}

/**
 * Return the global light rotation matrix
 */
export function getLightRotation()
{
    return new THREE.Matrix4().makeRotationY(Math.PI / 2);
}

/**
 * Returns all lights nodes in a MaterialX document
 * @param {mx.Document} doc 
 * @returns {Array.<mx.Node>}
 */
export function findLights(doc)
{
    let lights = [];
    for (let node of doc.getNodes())
    {
        if (node.getType() === "lightshader")
            lights.push(node);
    }
    return lights;
}

/**
 * Register lights in shader generation context
 * @param {Object} mx MaterialX Module
 * @param {Array.<mx.Node>} lights Light nodes
 * @param {mx.GenContext} genContext Shader generation context
 * @returns {Array.<mx.Node>}
 */
export async function registerLights(mx, lights, genContext)
{
    mx.HwShaderGenerator.unbindLightShaders(genContext);

    const lightTypesBound = {};
    const lightData = [];
    let lightId = 1;

    // All light types so that we have NodeDefs for them
    const defaultLightRigXml = `<?xml version="1.0"?>
<materialx version="1.39">
    <directional_light name="default_directional_light" type="lightshader">
    </directional_light>
    <point_light name="default_point_light" type="lightshader">
    </point_light>
    <spot_light name="default_spot_light" type="lightshader">
    </spot_light>
    <!--
    <area_light name="default_area_light" type="lightshader">
    </area_light>
    -->
</materialx>`;

    // Load default light rig XML to ensure we have all light types available
    const lightRigDoc = mx.createDocument();
    await mx.readFromXmlString(lightRigDoc, defaultLightRigXml);
    const document = mx.createDocument();
    const stdlib = mx.loadStandardLibraries(genContext);
    document.setDataLibrary(stdlib);
    document.importLibrary(lightRigDoc);
    const defaultLights = findLights(document);
    if (debug) console.log("Default lights in MaterialX document", defaultLights);

    // Register types only – we get these from the default light rig XML above
    // This is needed to ensure that the light shaders are bound for each light type
    for (let light of defaultLights)
    {
        const lightDef = light.getNodeDef();
        if (debug) console.log("Default light node definition", lightDef);
        if (!lightDef) continue;

        const lightName = lightDef.getName();
        if (debug) console.log("Registering default light", { lightName, lightDef });
        if (!lightTypesBound[lightName])
        {
            // TODO check if we need to bind light shader for each three.js light instead of once per type
            if (debug) console.log("Bind light shader for node", { lightName, lightId, lightDef });
            lightTypesBound[lightName] = lightId;
            mx.HwShaderGenerator.bindLightShader(lightDef, lightId++, genContext);
        }
    }

    if (debug) console.log("Light types bound in MaterialX context", lightTypesBound);

    // MaterialX light nodes
    for (let light of lights)
    {
        // Skip if light does not have a node definition
        if (!("getNodeDef" in light)) continue;

        let nodeDef = light.getNodeDef();
        let nodeName = nodeDef.getName();
        if (!lightTypesBound[nodeName])
        {
            if (debug) console.log("bind light shader for node", { nodeName, lightId, nodeDef });
            lightTypesBound[nodeName] = lightId;
            mx.HwShaderGenerator.bindLightShader(nodeDef, lightId++, genContext);
        }

        const lightDirection = light.getValueElement("direction").getValue().getData().data();
        const lightColor = light.getValueElement("color").getValue().getData().data();
        const lightIntensity = light.getValueElement("intensity").getValue().getData();

        let rotatedLightDirection = new THREE.Vector3(...lightDirection)
        rotatedLightDirection.transformDirection(getLightRotation())

        lightData.push({
            type: lightTypesBound[nodeName],
            direction: rotatedLightDirection,
            color: new THREE.Vector3(...lightColor),
            intensity: lightIntensity,
        });
    }

    const threeLightTypeToMaterialXNodeName = (threeLightType) => {
        switch (threeLightType) {
            case 'PointLight':
                return 'ND_point_light';
            case 'DirectionalLight':
                return 'ND_directional_light';
            case 'SpotLight':
                return 'ND_spot_light';
            default:
                console.warn('MaterialX: Unsupported light type: ' + threeLightType);
                return 'ND_point_light'; // Default to point light
        }
    };

    if (debug) console.log("Registering lights in MaterialX context", lights, lightData);

    // Three.js lights
    for (let light of lights) {
        // Skip if light is not a Three.js light
        if (!light.isLight) continue;

        // Types in MaterialX: point_light, directional_light, spot_light

        const lightDefinitionName = threeLightTypeToMaterialXNodeName(light.type);

        if (!lightTypesBound[lightDefinitionName])
        {
            lightTypesBound[lightDefinitionName] = lightId;
            const nodeDef = null;
            mx.HwShaderGenerator.bindLightShader(nodeDef, lightId++, genContext);
        }

        lightData.push({
            type: lightTypesBound[lightDefinitionName],
            direction: light.direction?.clone() || new THREE.Vector3(0, -1, 0),
            color: new THREE.Vector3().fromArray(light.color.toArray()),
            intensity: light.intensity,
        });
    }

    // Make sure max light count is large enough
    genContext.getOptions().hwMaxActiveLightSources = Math.max(genContext.getOptions().hwMaxActiveLightSources, lightData.length);

    return lightData;
}

/**
 * Get uniform values for a shader
 * @param {mx.shaderStage} shaderStage
 * @param {THREE.TextureLoader} textureLoader
 */
export function getUniformValues(shaderStage, textureLoader, searchPath, flipY)
{
    let threeUniforms = {};

    const uniformBlocks = Object.values(shaderStage.getUniformBlocks());
    uniformBlocks.forEach(uniforms =>
    {
        if (!uniforms.empty())
        {
            for (let i = 0; i < uniforms.size(); ++i)
            {
                const variable = uniforms.get(i);
                const value = variable.getValue()?.getData();
                const name = variable.getVariable();
                threeUniforms[name] = new THREE.Uniform(toThreeUniform(variable.getType().getName(), value, name, uniforms,
                    textureLoader, searchPath, flipY));
            }
        }
    });

    return threeUniforms;
}
