import { BufferGeometry, Camera, FrontSide, GLSL3, Group, Matrix3, Matrix4, Object3D, Scene, ShaderMaterial, Texture, Vector3, WebGLRenderer } from "three";
import { debug, getFrame, getTime } from "./utils.js";
import { MaterialXEnvironment } from "./materialx.js";
import { generateMaterialPropertiesForUniforms, getUniformValues } from "./materialx.helper.js";
import { cloneUniforms, cloneUniformsGroups } from "three/src/renderers/shaders/UniformsUtils.js";

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

/**
 * @typedef {Object} MaterialXMaterialInitParameters
 * @property {string} name
 * @property {string | null} [shaderName] - Optional name of the shader
 * @property {any} shader
 * @property {import('./materialx.helper.js').Callbacks} loaders
 * @property {import('./materialx.js').MaterialXContext} context
 * @property {import('three').MaterialParameters} [parameters] - Optional parameters
 * @property {boolean} [debug] - Debug flag
 */

/**
 * @typedef {"highp" | "mediump" | "lowp"} Precision
 */

// @dont-generate-component
export class MaterialXMaterial extends ShaderMaterial {

    /** The original name of the shader 
     * @type {string | null} */
    shaderName = null;

    /**
     * @param {MaterialXMaterial} source
     * @returns {this}
     */
    copy(source) {
        super.copy(source);
        this._context = source._context;
        this._shader = source._shader;
        this.uniforms = cloneUniforms(source.uniforms);
        this.uniformsGroups = cloneUniformsGroups(source.uniformsGroups);
        this.envMapIntensity = source.envMapIntensity;
        this.envMap = source.envMap;
        generateMaterialPropertiesForUniforms(this, this._shader.getStage('pixel'));
        generateMaterialPropertiesForUniforms(this, this._shader.getStage('vertex'));
        this.needsUpdate = true;
        return this;
    }

    /** @type {import('./materialx.js').MaterialXContext | null} */
    _context = null;
    /** @type {any} */
    _shader = null;
    /** @type {boolean} */
    _needsTangents = false;

    /**
     * @param {MaterialXMaterialInitParameters} [init]
     */
    constructor(init) {

        // TODO: we need to properly copy the uniforms and other properties from the source material
        if (!init) {
            super();
            return;
        }

        // Get vertex and fragment shader source, and remove #version directive for newer js. 
        // It's added by three.js glslVersion.
        let vertexShader = init.shader.getSourceCode("vertex");
        let fragmentShader = init.shader.getSourceCode("pixel");

        vertexShader = vertexShader.replace(/^#version.*$/gm, '').trim();
        fragmentShader = fragmentShader.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 original MaterialX for reference).

        // Patch vertexShader
        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_texcoord_2\b/g, 'uv2');
        vertexShader = vertexShader.replace(/\bi_texcoord_3\b/g, 'uv3');
        vertexShader = vertexShader.replace(/\bi_tangent\b/g, 'tangent');
        vertexShader = vertexShader.replace(/\bi_color_0\b/g, 'color');

        // Patch fragmentShader
        const precision = init.parameters?.precision || "highp";
        vertexShader = vertexShader.replace(/precision mediump float;/g, `precision ${precision} float;`);
        vertexShader = vertexShader.replace(/#define M_FLOAT_EPS 1e-8/g, precision === "highp" ? `#define M_FLOAT_EPS 1e-8` : `#define M_FLOAT_EPS 1e-3`);
        fragmentShader = fragmentShader.replace(/precision mediump float;/g, `precision ${precision} float;`);
        fragmentShader = fragmentShader.replace(/#define M_FLOAT_EPS 1e-8/g, precision === "highp" ? `#define M_FLOAT_EPS 1e-8` : `#define M_FLOAT_EPS 1e-3`);

        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_texcoord_2\b/g, 'uv2');
        fragmentShader = fragmentShader.replace(/\bi_texcoord_3\b/g, 'uv3');
        fragmentShader = fragmentShader.replace(/\bi_tangent\b/g, 'tangent');
        fragmentShader = fragmentShader.replace(/\bi_color_0\b/g, 'color');

        // Capture some vertex shader properties
        const uv_is_vec2 = vertexShader.includes('in vec2 uv;'); // check if uv is vec2; e.g. https://matlib.gpuopen.com/main/materials/all?material=da6ec531-f5c1-4790-ac14-8a5c51d0314e
        const uv1_is_vec2 = vertexShader.includes('in vec2 uv1;');
        const uv2_is_vec2 = vertexShader.includes('in vec2 uv2;');
        const uv3_is_vec2 = vertexShader.includes('in vec2 uv3;');

        // Remove `in vec3 position;` and so on since they're already declared by ShaderMaterial
        vertexShader = vertexShader.replace(/in\s+vec3\s+position;/g, '');
        vertexShader = vertexShader.replace(/in\s+vec3\s+normal;/g, '');
        vertexShader = vertexShader.replace(/in\s+vec2\s+uv;/g, '');
        vertexShader = vertexShader.replace(/in\s+vec3\s+uv;/g, '');
        var hasUv1 = vertexShader.includes('in vec3 uv1;');
        vertexShader = vertexShader.replace(/in\s+vec3\s+uv1;/g, '');
        var hasUv2 = vertexShader.includes('in vec3 uv2;');
        vertexShader = vertexShader.replace(/in\s+vec3\s+uv2;/g, '');
        var hasUv3 = vertexShader.includes('in vec3 uv3;');
        vertexShader = vertexShader.replace(/in\s+vec3\s+uv3;/g, '');
        var hasTangent = vertexShader.includes('in vec4 tangent;');
        vertexShader = vertexShader.replace(/in\s+vec4\s+tangent;/g, '');
        var hasColor = vertexShader.includes('in vec4 color;');
        vertexShader = vertexShader.replace(/in\s+vec4\s+color;/g, '');

        // Patch uv 2-component to 3-component (`texcoord_0 = uv;` needs to be replaced with `texcoord_0 = vec3(uv, 0.0);`)
        // TODO what if we actually have a 3-component UV? Not sure what three.js does then
        if (!uv_is_vec2) vertexShader = vertexShader.replace(/texcoord_0 = uv;/g, 'texcoord_0 = vec3(uv, 0.0);');
        if (!uv1_is_vec2) vertexShader = vertexShader.replace(/texcoord_1 = uv1;/g, 'texcoord_1 = vec3(uv1, 0.0);');
        if (!uv2_is_vec2) vertexShader = vertexShader.replace(/texcoord_2 = uv2;/g, 'texcoord_2 = vec3(uv2, 0.0);');
        if (!uv3_is_vec2) vertexShader = vertexShader.replace(/texcoord_3 = uv3;/g, 'texcoord_3 = vec3(uv3, 0.0);');

        // Patch units – seems MaterialX uses different units and we end up with wrong light values?
        // result.direction = light.position - position;
        fragmentShader = fragmentShader.replace(
            /result\.direction\s*=\s*light\.position\s*-\s*position;/g,
            'result.direction = (light.position - position) * 10.0 / 1.0;');

        // Add tonemapping and colorspace handling
        // Replace `out vec4 out1;` with `out vec4 gl_FragColor;`
        fragmentShader = fragmentShader.replace(
            /out\s+vec4\s+out1;/,
            'layout(location = 0) out vec4 pc_fragColor;\n#define gl_FragColor pc_fragColor');

        // Replace `out1 = vec4(<CAPTURE>)` with `gl_FragColor = vec4(<CAPTURE>)` and tonemapping/colorspace handling
        fragmentShader = fragmentShader.replace(/^\s*out1\s*=\s*vec4\((.*)\);/gm, `
    gl_FragColor = vec4($1);
    #include <tonemapping_fragment>
    #include <colorspace_fragment>`);

        const defines = {};
        if (hasUv1) defines['USE_UV1'] = '';
        if (hasUv2) defines['USE_UV2'] = '';
        if (hasUv3) defines['USE_UV3'] = '';
        if (hasTangent) defines['USE_TANGENT'] = '';
        if (hasColor) defines['USE_COLOR'] = '';

        const searchPath = ""; // Could be derived from the asset path if needed
        const isTransparent = init.parameters?.transparent ?? false;
        super({
            name: init.name,
            uniforms: {},
            vertexShader: vertexShader,
            fragmentShader: fragmentShader,
            glslVersion: GLSL3,
            depthTest: true,
            depthWrite: !isTransparent,
            defines: defines,
            ...init.parameters, // Spread any additional parameters passed to the material
        });
        this.shaderName = init.shaderName || null;
        this._context = init.context;
        this._shader = init.shader;
        this._needsTangents = vertexShader.includes('in vec4 tangent;') || vertexShader.includes('in vec3 tangent;');

        Object.assign(this.uniforms, {
            ...getUniformValues(init.shader.getStage('vertex'), init.loaders, searchPath),
            ...getUniformValues(init.shader.getStage('pixel'), init.loaders, searchPath),

            u_worldMatrix: { value: new Matrix4() },
            u_viewProjectionMatrix: { value: new Matrix4() },
            u_viewPosition: { value: new Vector3() },
            u_worldInverseTransposeMatrix: { value: new Matrix4() },

            // u_shadowMatrix: { value: new Matrix4() },
            // u_shadowMap: { value: null, type: 't' }, // Shadow map

            u_envMatrix: { value: new Matrix4() },
            u_envRadiance: { value: null, 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: null, type: 't' },
            u_refractionEnv: { value: true },
            u_numActiveLightSources: { value: 0 },
            u_lightData: { value: [], needsUpdate: false }, // Array of light data. We need to set needsUpdate to false until we actually update it
        });

        generateMaterialPropertiesForUniforms(this, init.shader.getStage('pixel'));
        generateMaterialPropertiesForUniforms(this, init.shader.getStage('vertex'));


        if (debug || init.debug) {
            // Get lighting and environment data from MaterialX environment
            console.group("[MaterialX]: ", this.name);
            console.log(`Vertex shader length: ${vertexShader.length}\n`, vertexShader);
            console.log(`Fragment shader length: ${fragmentShader.length}\n`, fragmentShader);
            console.groupEnd();
        }
    }

    /** @type {boolean} */
    _missingTangentsWarned = false;

    /**
     * @param {WebGLRenderer} renderer
     * @param {Scene} _scene
     * @param {Camera} camera
     * @param {BufferGeometry} geometry
     * @param {Object3D} object
     * @param {Group} _group
     */
    onBeforeRender(renderer, _scene, camera, geometry, object, _group) {
        if (this._needsTangents && !geometry.attributes.tangent) {
            if (!this._missingTangentsWarned) {
                this._missingTangentsWarned = true;
                console.warn(`[MaterialX] Tangents are required for this material (${this.name}) but not present in the geometry.`);
                // TODO: can we compute tangents here?
            }
        }
        const time = this._context?.getTime?.() || getTime();
        const frame = this._context?.getFrame?.() || getFrame();
        const env = MaterialXEnvironment.get(_scene);
        if (env) {
            env.update(frame, _scene, renderer);
            this.updateUniforms(env, renderer, object, camera, time, frame);
        }
    }

    /** @type {number} */
    envMapIntensity = 1.0; // Default intensity for environment map
    /** @type {Texture | null} */
    envMap = null; // Environment map texture, can be set externally

    /**
     * @param {MaterialXEnvironment} environment
     * @param {WebGLRenderer} _renderer
     * @param {Object3D} object
     * @param {Camera} camera
     * @param {number} [time]
     * @param {number} [frame]
     */
    updateUniforms = (environment, _renderer, object, camera, time, frame) => {

        const uniforms = this.uniforms;

        // Update standard transformation matrices
        if (uniforms.u_worldMatrix) {
            uniforms.u_worldMatrix.value = object.matrixWorld;
            // @ts-ignore
            uniforms.u_worldMatrix.needsUpdate = true;
        }

        if (uniforms.u_viewProjectionMatrix) {
            uniforms.u_viewProjectionMatrix.value.multiplyMatrices(camera.projectionMatrix, camera.matrixWorldInverse);
            // @ts-ignore
            uniforms.u_viewProjectionMatrix.needsUpdate = true;
        }

        if (uniforms.u_worldInverseTransposeMatrix) {
            uniforms.u_worldInverseTransposeMatrix.value.setFromMatrix3(normalMat.getNormalMatrix(object.matrixWorld));
            // @ts-ignore
            uniforms.u_worldInverseTransposeMatrix.needsUpdate = true;
        }

        if (uniforms.u_viewPosition) {
            uniforms.u_viewPosition.value.copy(camera.getWorldPosition(worldViewPos));
            // @ts-ignore
            uniforms.u_viewPosition.needsUpdate = true;
        }

        // if (uniforms.u_shadowMap) {
        //     const light = environment.lights?.[2] || null;
        //     uniforms.u_shadowMatrix.value = light?.shadow?.matrix.clone().premultiply(object.matrixWorld.clone()).invert();
        //     uniforms.u_shadowMap.value = light.shadow?.map || null;
        //     uniforms.u_shadowMap.needsUpdate = true;
        //     console.log("[MaterialX] Renderer shadow map updated", light);
        // }

        // Update time uniforms
        if (uniforms.u_time) {
            if (time === undefined) time = getTime();
            uniforms.u_time.value = time;
        }
        if (uniforms.u_frame) {
            if (frame === undefined) frame = getFrame();
            uniforms.u_frame.value = frame;
        }

        this.uniformsNeedUpdate = true;

        // Update light uniforms
        this.updateEnvironmentUniforms(environment);
    }

    /**
     * @private
     * @param {MaterialXEnvironment} environment
     */
    updateEnvironmentUniforms = (environment) => {

        const uniforms = this.uniforms;

        // Get lighting data from environment
        const lightData = environment.lightData || null;
        const lightCount = environment.lightCount || 0;
        const textures = environment.getTextures(this) || null;

        // Update light count
        if (uniforms.u_numActiveLightSources && lightCount >= 0) {
            uniforms.u_numActiveLightSources.value = lightCount;
        }

        // Update light data
        if (lightData?.length) {
            uniforms.u_lightData.value = lightData;
            if ("needsUpdate" in uniforms.u_lightData && uniforms.u_lightData.needsUpdate === false) {
                if (debug) console.debug(`[MaterialX] LightData assigned (${this.name}, ${this.uuid})`, lightData);
                uniforms.u_lightData.needsUpdate = undefined;
            }
        }

        // Update environment uniforms
        if (uniforms.u_envRadiance) {
            const prev = uniforms.u_envRadiance.value;
            uniforms.u_envRadiance.value = textures.radianceTexture;
            // @ts-ignore
            if (prev != textures.radianceTexture) uniforms.u_envRadiance.needsUpdate = true;
        }
        if (uniforms.u_envRadianceMips) {
            uniforms.u_envRadianceMips.value = Math.trunc(Math.log2(Math.max(textures.radianceTexture?.source.data.width ?? 0, textures.radianceTexture?.source.data.height ?? 0))) + 1;
        }
        if (uniforms.u_envIrradiance) {
            const prev = uniforms.u_envIrradiance.value;
            uniforms.u_envIrradiance.value = textures.irradianceTexture;
            // @ts-ignore
            if (prev != textures.irradianceTexture) uniforms.u_envIrradiance.needsUpdate = true;
        }

        this.uniformsNeedUpdate = true;
    }
}
