//
// Copyright Contributors to the MaterialX Project
// SPDX-License-Identifier: Apache-2.0
//
import * as THREE from 'three';
import { debug, debugUpdate } from './utils.js';
import { MaterialX } from './materialx.types.js';

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
 */
function fromMatrix(matrix: MaterialX.Matrix, dimension: MaterialX.Matrix["size"]) {
    const vec = new Array(dimension);
    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;
}


export type Loaders = {
    /**
     * Cache key for the loaders, used to identify and reuse textures
     */
    readonly cacheKey: string;
    /**
     * Get a texture by path
     * @param {string} path - The path to the texture
     * @return {Promise<THREE.Texture>} - A promise that resolves to the texture
     */
    readonly getTexture: (path: string) => Promise<THREE.Texture>;
}

const defaultTexture = new THREE.Texture();
defaultTexture.needsUpdate = true;
defaultTexture.image = new Image();
defaultTexture.image.src = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD+wSzIAAAAAXNSR0IB2cksfwAAAAlwSFlzAAALEwAACxMBAJqcGAAAAANQTFRFr6+vGqg52AAAAAxJREFUeJxjZGBEgQAAWAAJLpjsTQAAAABJRU5ErkJggg=="
// defaultTexture.image.src = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQBAMAAADt3eJSAAAAAXNSR0IB2cksfwAAAAlwSFlzAAALEwAACxMBAJqcGAAAAB5QTFRFAAAABAQEw8PD////v7+/vb29Xl5eQEBA+/v7PDw8GPBYkgAAAB1JREFUeJxjZGBgYFQSABIUMlxgDGMGBtaIAnIZAKwQCSDYUEZEAAAAAElFTkSuQmCC";
// defaultTexture.wrapS = THREE.RepeatWrapping;
// defaultTexture.wrapT = THREE.RepeatWrapping;
// defaultTexture.minFilter = THREE.NearestFilter;
// defaultTexture.magFilter = THREE.NearestFilter;
// defaultTexture.repeat = new THREE.Vector2(100, 100);


const defaultNormalTexture = new THREE.Texture();
defaultNormalTexture.needsUpdate = true;
defaultNormalTexture.image = new Image();
defaultNormalTexture.image.src = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIBAMAAAA2IaO4AAAAAXNSR0IB2cksfwAAAAlwSFlzAAALEwAACxMBAJqcGAAAABJQTFRFgYH4gIH4gYH3gIH3gIH5gID4m94ORAAAADFJREFUeJxjZBBkfMdo9P/BB0aBj/8FGB0ufghgFGT4r8wo+P8rD2Pgo3sMjIz8jAwAMLoN0ZjS5hgAAAAASUVORK5CYII=";


function tryGetFromCache(key: string): any {
    const wasEnabled = THREE.Cache.enabled;
    THREE.Cache.enabled = true;
    const value = THREE.Cache.get(key);
    THREE.Cache.enabled = wasEnabled;
    return value;
}
function addToCache(key: string, value: any): void {
    const wasEnabled = THREE.Cache.enabled;
    THREE.Cache.enabled = true;
    THREE.Cache.add(key, value);
    THREE.Cache.enabled = wasEnabled;
    if (debug) console.log('[MaterialX] Added to cache:', key, value);
}

/**
 * Get Three uniform from MaterialX value
 */
function toThreeUniform(uniforms: any, type: string, value: any, name: string, loaders: Loaders, searchPath: string): THREE.Uniform {

    const uniform = new THREE.Uniform<any>(null);

    switch (type) {
        case 'float':
        case 'integer':
        case 'boolean':
            uniform.value = value;
            break;
        case 'vector2':
            uniform.value = fromVector(value, 2);
            break;
        case 'vector3':
        case 'color3':
            uniform.value = fromVector(value, 3);
            break;
        case 'vector4':
        case 'color4':
            uniform.value = fromVector(value, 4);
            break;
        case 'matrix33':
            uniform.value = fromMatrix(value, 9);
            break;
        case 'matrix44':
            uniform.value = 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 = true;
                let texturePath = searchPath + IMAGE_PATH_SEPARATOR + value;
                if (value.startsWith('blob:')) {
                    texturePath = value;
                    checkCache = false;
                }
                else if (value.startsWith('data:')) {
                    texturePath = value;
                    checkCache = false;
                }
                else if (value.startsWith('http')) {
                    texturePath = value;
                    checkCache = true;
                }

                const cacheKey = `${loaders.cacheKey}-${texturePath}`;
                const cacheValue = checkCache ? tryGetFromCache(cacheKey) : null;
                if (cacheValue) {
                    if (debug) console.log('[MaterialX] Use cached texture: ', cacheKey, cacheValue);
                    if (cacheValue instanceof Promise) {
                        cacheValue.then(res => {
                            if (res) uniform.value = res;
                            else console.warn(`[MaterialX] Failed to load texture ${name} '${texturePath}'`);
                        });
                    }
                    else {
                        uniform.value = cacheValue;
                    }
                }
                else {
                    if (name.toLowerCase().includes("normal")) {
                        uniform.value = defaultNormalTexture;
                    }
                    else {
                        uniform.value = defaultTexture;
                    }

                    if (debug) console.log('[MaterialX] Load texture:', texturePath);
                    // Save the loading promise in the cache
                    const promise = loaders.getTexture(texturePath).then(res => {
                        if (res) {
                            res = res.clone(); // we need to clone the texture once to avoid colorSpace issues with other materials
                            res.colorSpace = THREE.LinearSRGBColorSpace;
                            setTextureParameters(res, name, uniforms);
                        }
                        return res;
                    });
                    if (checkCache) {
                        addToCache(cacheKey, promise);
                    }
                    promise.then(res => {
                        if (res) uniform.value = res;
                        else console.warn(`[MaterialX] Failed to load texture ${name} '${texturePath}'`);
                    });
                }
            }
            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}`);
            }
            break;
    }

    return uniform;
}

const valueTypeWarningMap = new Map<string, boolean>();

/**
 * Get Three wrapping mode
 */
function getWrapping(mode: number): THREE.Wrapping {
    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;
}


/**
 * Set Three texture parameters
 */
function setTextureParameters(texture: THREE.Texture, name: string, uniforms: any, generateMipmaps = true) {
    const idx = name.lastIndexOf(IMAGE_PROPERTY_SEPARATOR);
    const base = name.substring(0, idx) || name;

    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 mxFilterType = uniforms.find(base + FILTER_TYPE_SUFFIX) ? uniforms.get(base + FILTER_TYPE_SUFFIX).value : -1;
    let minFilter: THREE.TextureFilter = generateMipmaps ? THREE.LinearMipMapLinearFilter : THREE.LinearFilter;
    if (mxFilterType === 0) {
        minFilter = generateMipmaps ? THREE.NearestMipMapNearestFilter : THREE.NearestFilter;
    }
    texture.minFilter = minFilter;
}

/**
 * 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: MaterialX.Document) {
    let lights = new Array<any>;
    for (let node of doc.getNodes()) {
        if (node.getType() === "lightshader")
            lights.push(node);
    }
    return lights;
}

let lightTypesBound = {};

/**
 * Register lights in shader generation context
 * @param {MaterialX.MODULE} mx MaterialX Module
 * @param {mx.GenContext} genContext Shader generation context
 */
export async function registerLights(mx: MaterialX.MODULE, genContext: any): Promise<void> {
    lightTypesBound = {};
    const maxLightCount = genContext.getOptions().hwMaxActiveLightSources;
    mx.HwShaderGenerator.unbindLightShaders(genContext);
    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);

    // Loading a document seems to reset this option for some reason, so we set it again
    genContext.getOptions().hwMaxActiveLightSources = maxLightCount;

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

// Converts Three.js light type to MaterialX node name
function 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
    }
};

export type LightData = {
    type: number, // Light type ID
    position: THREE.Vector3, // Position in world space
    direction: THREE.Vector3, // Direction in world space
    color: THREE.Color, // Color of the light
    intensity: number, // Intensity of the light
    decay_rate: number, // Decay rate for point and spot lights
    inner_angle: number, // Inner angle for spot lights
    outer_angle: number, // Outer angle for spot lights
}

/**
 * Update light data for shader uniforms
 */
export function getLightData(lights: Array<THREE.Light>, genContext: any): { lightData: LightData[], lightCount: number } {
    const lightData = new Array();
    const maxLightCount = genContext.getOptions().hwMaxActiveLightSources;

    // 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])
            console.error("MaterialX: Light type not registered in context. Make sure to register light types before using them.", lightDefinitionName);

        const wp = light.getWorldPosition(new THREE.Vector3());
        const wq = light.getWorldQuaternion(new THREE.Quaternion());
        const wd = new THREE.Vector3(0, 0, -1).applyQuaternion(wq);

        // Shader math from the generated MaterialX shader:
        // float low = min(light.inner_angle, light.outer_angle);
        // float high = light.inner_angle;
        // float cosDir = dot(result.direction, -light.direction);
        // float spotAttenuation = smoothstep(low, high, cosDir);

        const outerAngleRad = (light as THREE.SpotLight).angle;
        const innerAngleRad = outerAngleRad * (1 - (light as THREE.SpotLight).penumbra);
        const inner_angle = Math.cos(innerAngleRad);
        const outer_angle = Math.cos(outerAngleRad);

        lightData.push({
            type: lightTypesBound[lightDefinitionName],
            position: wp.clone(),
            direction: wd.clone(),
            color: new THREE.Color().fromArray(light.color.toArray()),
            // Luminous efficacy for converting radiant power in watts (W) to luminous flux in lumens (lm) at a wavelength of 555 nm.
            // Also, three.js lights don't have PI scale baked in, but MaterialX does, so we need to divide by PI for point and spot lights.
            intensity: light.intensity * ((light as THREE.PointLight).isPointLight ? 683.0 / 3.1415 : (light as THREE.SpotLight).isSpotLight ? 683.0 / 3.1415 : 1.0),
            decay_rate: 2.0,
            // Approximations for testing – the relevant light has 61.57986...129.4445 as inner/outer spot angle
            inner_angle: inner_angle,
            outer_angle: outer_angle,
        });
    }

    // Count the number of lights that are not empty
    const lightCount = lightData.length;

    // If we don't have enough entries in lightData, fill with empty lights
    while (lightData.length < maxLightCount) {
        const emptyLight = {
            type: 0, // Default light type
            position: new THREE.Vector3(0, 0, 0),
            direction: new THREE.Vector3(0, 0, -1),
            color: new THREE.Color(0, 0, 0),
            intensity: 0.0,
            decay_rate: 2.0,
            inner_angle: 0.0,
            outer_angle: 0.0,
        };
        lightData.push(emptyLight);
    }

    if (debugUpdate) console.log("Registered lights in MaterialX context", lightTypesBound, lightData);

    return { lightData, lightCount };
}

/**
 * Get uniform values for a shader
 */
export function getUniformValues(shaderStage: MaterialX.ShaderStage, loaders: Loaders, searchPath: string) {
    const threeUniforms = {};

    const uniformBlocks = shaderStage.getUniformBlocks()
    for (const [blockName, uniforms] of Object.entries(uniformBlocks)) {
        // Seems struct uniforms (like in LightData) end up here as well, we should filter those out.
        if (blockName === "LightData") continue;

        if (!uniforms.empty()) {
            for (let i = 0; i < uniforms.size(); ++i) {
                const variable = uniforms.get(i);
                const value = variable.getValue()?.getData();
                const uniformName = variable.getVariable();
                const type = variable.getType().getName();
                if (debug) console.log("Adding uniform", { path: variable.getPath(), name: uniformName, value: value, type: type });
                threeUniforms[uniformName] = toThreeUniform(uniforms, type, value, uniformName, loaders, searchPath);
            }
        }
    }

    return threeUniforms;
}

export function generateMaterialPropertiesForUniforms(material: THREE.ShaderMaterial, shaderStage: MaterialX.ShaderStage) {

    const uniformBlocks = shaderStage.getUniformBlocks()
    for (const [blockName, uniforms] of Object.entries(uniformBlocks)) {
        // Seems struct uniforms (like in LightData) end up here as well, we should filter those out.
        if (blockName === "LightData") continue;

        if (!uniforms.empty()) {
            for (let i = 0; i < uniforms.size(); ++i) {
                const variable = uniforms.get(i);
                const uniformName = variable.getVariable();
                let key = variable.getPath().split('/').pop();
                switch (key) {
                    case "_Color":
                        key = "color";
                        break;
                    case "_Roughness":
                        key = "roughness";
                        break;
                    case "_Metallic":
                        key = "metalness";
                        break;
                }
                if (key) {
                    Object.defineProperty(material, key, {
                        get: function () {
                            return this.uniforms?.[uniformName].value
                        },
                        set: function (v) {
                            const uniforms = this.uniforms;
                            if (!uniforms || !uniforms[uniformName]) {
                                console.warn(`[MaterialX] Uniform ${uniformName} not found in ${this.name} uniforms`);
                                return;
                            }
                            this.uniforms[uniformName].value = v;
                            this.uniformsNeedUpdate = true;
                        }
                    });
                }
            }
        }
    }
}