import { CubeUVReflectionMapping, ShaderMaterial, Texture } from "three";
import { GroundedSkybox as GroundProjection } from 'three/examples/jsm/objects/GroundedSkybox.js';

import { Gizmos } from "../engine/engine_gizmos.js";
import { serializable } from "../engine/engine_serialization_decorator.js";
import { getBoundingBox, getTempVector, getWorldScale, Graphics, setVisibleInCustomShadowRendering, setWorldPosition } from "../engine/engine_three_utils.js";
import { delayForFrames, getParam, Watch as Watch } from "../engine/engine_utils.js";
// Type-only imports for TSDoc @see links
import type { Camera } from "./Camera.js";
import { Behaviour } from "./Component.js";
import type { ContactShadows } from "./ContactShadows.js";

const debug = getParam("debuggroundprojection");

type GroundProjectionMaterial = GroundProjection["material"] & {
    defines?: Record<string, string | number>;
};

type GroundProjectionShaderUniforms = {
    needleGroundProjectionBlurriness: { value: number };
    needleGroundProjectionBlending: { value: number };
    needleGroundProjectionAlphaFactor: { value: number };
    needleGroundProjectionBackgroundIntensity: { value: number };
};

const needleCubeUvMapVarying = /* glsl */`
#ifdef NEEDLE_USE_CUBE_UV_MAP
varying vec3 vNeedleGroundProjectionWorldDirection;
#endif
`;

const needleGroundProjectionFragmentPars = /* glsl */`
${needleCubeUvMapVarying}
uniform float needleGroundProjectionBlurriness;
uniform float needleGroundProjectionBlending;
uniform float needleGroundProjectionAlphaFactor;
uniform float needleGroundProjectionBackgroundIntensity;

float needleGroundProjectionSmoothstep(float edge0, float edge1, float x) {
    float t = clamp((x - edge0) / max(edge1 - edge0, 0.000001), 0.0, 1.0);
    return t * t * (3.0 - 2.0 * t);
}

float needleGroundProjectionDistance() {
    return length(vec2(0.0, vMapUv.y));
}

float needleGroundProjectionBlurFactor(float needleGroundProjectionDistanceValue) {
    return clamp(needleGroundProjectionSmoothstep(0.5, 1.0, needleGroundProjectionDistanceValue * 2.0), 0.0, 1.0);
}
`;

const needleCubeUvMapFragment = /* glsl */`
#ifdef USE_MAP

    float needleGroundProjectionDistanceValue = needleGroundProjectionDistance();
    float needleGroundProjectionBlurFactorValue = needleGroundProjectionBlurFactor(needleGroundProjectionDistanceValue);
    vec4 sampledDiffuseColor;

    #ifdef NEEDLE_USE_CUBE_UV_MAP
        sampledDiffuseColor = textureCubeUV(
            map,
            normalize( vNeedleGroundProjectionWorldDirection ),
            needleGroundProjectionBlurriness * needleGroundProjectionBlurFactorValue
        );
    #else
        #ifdef USE_MIPMAP_BIAS
            sampledDiffuseColor = texture2D( map, vMapUv, mipmapBias );
        #else
            sampledDiffuseColor = texture2D( map, vMapUv );
        #endif
    #endif

    #ifdef DECODE_VIDEO_TEXTURE

        // use inline sRGB decode until browsers properly support SRGB8_ALPHA8 with video textures (#26516)

        sampledDiffuseColor = vec4( mix( pow( sampledDiffuseColor.rgb * 0.9478672986 + vec3( 0.0521327014 ), vec3( 2.4 ) ), sampledDiffuseColor.rgb * 0.0773993808, vec3( lessThanEqual( sampledDiffuseColor.rgb, vec3( 0.04045 ) ) ) ), sampledDiffuseColor.w );

    #endif

    sampledDiffuseColor.rgb *= mix(1.0, needleGroundProjectionBackgroundIntensity, needleGroundProjectionBlurFactorValue);
    diffuseColor *= sampledDiffuseColor;

#endif
`;

const needleGroundProjectionAlphaFragment = /* glsl */`
#ifdef USE_MAP
    if (needleGroundProjectionBlending > 0.000001) {
        float needleGroundProjectionBrightness = dot(diffuseColor.rgb, vec3(0.299, 0.587, 0.114));
        float needleGroundProjectionStepFactor = needleGroundProjectionBlending - needleGroundProjectionBrightness * 0.1;
        diffuseColor.a *= pow(
            1.0 - needleGroundProjectionBlending * needleGroundProjectionSmoothstep(
                0.35 * needleGroundProjectionStepFactor,
                0.45 * needleGroundProjectionStepFactor,
                needleGroundProjectionDistanceValue
            ),
            5.0
        );
    }
#endif
diffuseColor.a *= needleGroundProjectionAlphaFactor;
`;

function getCubeUvSize(texture: Texture) {
    const imageHeight = texture.image?.height;
    if (!imageHeight) return null;

    const maxMip = Math.log2(imageHeight) - 2;
    const texelHeight = 1 / imageHeight;
    const texelWidth = 1 / (3 * Math.max(Math.pow(2, maxMip), 7 * 16));
    return { texelWidth, texelHeight, maxMip };
}

function getGroundProjectionShaderUniforms(material: GroundProjectionMaterial): GroundProjectionShaderUniforms {
    const userData = material.userData as GroundProjectionMaterial["userData"] & {
        needleGroundProjectionUniforms?: GroundProjectionShaderUniforms;
    };

    return userData.needleGroundProjectionUniforms ??= {
        needleGroundProjectionBlurriness: { value: 0 },
        needleGroundProjectionBlending: { value: 0 },
        needleGroundProjectionAlphaFactor: { value: 1 },
        needleGroundProjectionBackgroundIntensity: { value: 1 }
    };
}

function configureGroundProjectionMaterial(material: GroundProjectionMaterial, texture: Texture) {
    const projectionUniforms = getGroundProjectionShaderUniforms(material);
    material.onBeforeCompile = shader => {
        shader.uniforms.needleGroundProjectionBlurriness = projectionUniforms.needleGroundProjectionBlurriness;
        shader.uniforms.needleGroundProjectionBlending = projectionUniforms.needleGroundProjectionBlending;
        shader.uniforms.needleGroundProjectionAlphaFactor = projectionUniforms.needleGroundProjectionAlphaFactor;
        shader.uniforms.needleGroundProjectionBackgroundIntensity = projectionUniforms.needleGroundProjectionBackgroundIntensity;

        shader.vertexShader = shader.vertexShader
            .replace("#include <uv_pars_vertex>", `#include <uv_pars_vertex>\n${needleCubeUvMapVarying}`)
            .replace(
                "#include <worldpos_vertex>",
                `#include <worldpos_vertex>
#ifdef NEEDLE_USE_CUBE_UV_MAP
	// GroundedSkybox mirrors geometry on Z, so undo that before deriving the sampling direction.
	vNeedleGroundProjectionWorldDirection = transformDirection( vec3( position.x, position.y, -position.z ), modelMatrix );
#endif`
            );

        shader.fragmentShader = shader.fragmentShader
            .replace(
                "#include <map_pars_fragment>",
                `#include <map_pars_fragment>
${needleGroundProjectionFragmentPars}
#include <cube_uv_reflection_fragment>`
            )
            .replace("#include <map_fragment>", needleCubeUvMapFragment)
            .replace("#include <opaque_fragment>", `${needleGroundProjectionAlphaFragment}\n#include <opaque_fragment>`);
    };

    const defines = material.defines ??= {};
    const prevDefineState = JSON.stringify(defines);
    const cubeUvSize = texture.mapping === CubeUVReflectionMapping ? getCubeUvSize(texture) : null;

    if (cubeUvSize) {
        defines.NEEDLE_USE_CUBE_UV_MAP = 1;
        defines.ENVMAP_TYPE_CUBE_UV = 1;
        defines.CUBEUV_TEXEL_WIDTH = cubeUvSize.texelWidth;
        defines.CUBEUV_TEXEL_HEIGHT = cubeUvSize.texelHeight;
        defines.CUBEUV_MAX_MIP = `${cubeUvSize.maxMip}.0`;
    }
    else {
        delete defines.NEEDLE_USE_CUBE_UV_MAP;
        delete defines.ENVMAP_TYPE_CUBE_UV;
        delete defines.CUBEUV_TEXEL_WIDTH;
        delete defines.CUBEUV_TEXEL_HEIGHT;
        delete defines.CUBEUV_MAX_MIP;
    }

    if (prevDefineState !== JSON.stringify(defines)) {
        material.needsUpdate = true;
    }
}

/**
 * The [GroundProjectedEnv](https://engine.needle.tools/docs/api/GroundProjectedEnv) projects the environment map onto a virtual ground plane.
 * Creates a realistic floor from 360° panoramas/HDRIs by deforming the skybox
 * into a hemisphere with a beveled floor.  
 *  
 * 
 * [![](https://cloud.needle.tools/-/media/8LDMd4TnGxVIj1XOfxIUIA.gif)](https://engine.needle.tools/samples/ground-projection)
 *
 * **Key properties:**  
 * - `radius` - Size of the projection sphere (keep camera inside)
 * - `height` - How high the original photo was taken (affects floor magnification)
 * - `autoFit` - Automatically center and position at ground level
 * - `arBlending` - Blend with real-world in AR (0=hidden, 1=visible)
 *
 * **Debug:** Use `?debuggroundprojection` URL parameter.  
 *
 * @example Apply ground projection
 * ```ts
 * const ground = myObject.getComponent(GroundProjectedEnv);
 * ground.radius = 100;
 * ground.height = 2;
 * ground.apply();
 * ```
 *
 * @summary Projects the environment map onto the ground
 * @category Rendering
 * @group Components
 * @see {@link Camera} for environment/skybox settings
 * @see {@link ContactShadows} for ground shadows
 * @link https://engine.needle.tools/samples/ground-projection for a demo of ground projection
 */
export class GroundProjectedEnv extends Behaviour {

    /**
     * If true the projection will be created on awake and onEnable
     * @default false
     */
    @serializable()
    applyOnAwake: boolean = false;

    /**
     * When enabled the position of the projected environment will be adjusted to be centered in the scene (and ground level).
     * @default true
     */
    @serializable()
    autoFit: boolean = true;

    /**
     * Radius of the projection sphere. Set it large enough so the camera stays inside (make sure the far plane is also large enough)
     * @default 50
     */
    @serializable()
    set radius(val: number) {
        this._radius = val;
        if (this._projection) this.updateProjection();
    }
    get radius(): number { return this._radius; }
    private _radius: number = 50;

    /**
     * How far the camera that took the photo was above the ground. A larger value will magnify the downward part of the image.
     * @default 3
     */
    @serializable()
    set height(val: number) {
        this._height = val;
        if (this._projection) this.updateProjection();
    }
    get height(): number { return this._height; }
    private _height: number = 3;

    /**
     * Blending factor for the AR projection being blended with the scene background.    
     * 0 = not visible in AR - 1 = blended with real world background.    
     * Values between 0 and 1 control the smoothness of the blend while lower values result in smoother blending.
     * @default 0
     */
    @serializable()
    set arBlending(val: number) {
        this._arblending = val;
        this._needsTextureUpdate = true;
    }
    get arBlending(): number { return this._arblending; }
    private _arblending = 0;

    private _lastBackground?: Texture;
    private _lastRadius?: number;
    private _lastHeight?: number;
    private _projection?: GroundProjection;
    private _watcher?: Watch;


    /** @internal */
    awake() {
        if (this.applyOnAwake)
            this.updateAndCreate();
    }
    /** @internal */
    onEnable() {
        // TODO: if we do this in the first frame we can not disable it again. Something buggy with the watch?!
        if (this.context.time.frameCount > 0) {
            if (this.applyOnAwake)
                this.updateAndCreate();
        }
        if (!this._watcher) {
            this._watcher = new Watch(this.context.scene, "background");
            this._watcher.subscribeWrite(_ => {
                if (debug) console.log("Background changed", this.context.scene.background);
                this._needsTextureUpdate = true;
            });
        }
    }
    /** @internal */
    onDisable() {
        this._watcher?.revoke();
        this._projection?.removeFromParent();
    }
    /** @internal */
    onEnterXR(): void {
        if (!this.activeAndEnabled) return;
        this._needsTextureUpdate = true;
        this.updateProjection();
    }
    /** @internal */
    async onLeaveXR() {
        if (!this.activeAndEnabled) return;
        await delayForFrames(1);
        this.updateProjection();
    }
    /** @internal */
    onBeforeRender(): void {
        if (this._projection && this.scene.backgroundRotation) {
            this._projection.rotation.copy(this.scene.backgroundRotation);
        }

        if (this._projection && this.context.scene.background instanceof Texture) {
            const blurriness = this.context.scene.backgroundBlurriness ?? 0;
            const blurrinessChanged = this._lastBlurriness !== blurriness;
            this.updateProjectionMaterial(this.context.scene.background, blurrinessChanged || this._needsTextureUpdate);
        }
    }

    private updateAndCreate() {
        this.updateProjection();
        this._watcher?.apply();
    }


    private _needsTextureUpdate = false;

    /**
     * Updates the ground projection. This is called automatically when the environment or settings change.
     */
    updateProjection() {
        if (!this.context.scene.background) {
            this._projection?.removeFromParent();
            return;
        }
        const backgroundTexture = this.context.scene.background;
        if (!(backgroundTexture instanceof Texture)) {
            this._projection?.removeFromParent();
            return;
        }

        if (this.context.xr?.isPassThrough || this.context.xr?.isAR) {
            if (this.arBlending === 0) {
                this._projection?.removeFromParent();
                return;
            }
        }
        // If this is called from a setter during initialization
        if (!this.gameObject || this.destroyed) {
            return;
        }

        let needsNewAutoFit = true;
        // offset here must be zero (and not .01) because the plane occlusion (when mesh tracking is active) is otherwise not correct
        const offset = 0;

        const hasChanged = backgroundTexture !== this._lastBackground || this._height !== this._lastHeight || this._radius !== this._lastRadius;

        if (!this._projection || hasChanged) {
            if (debug) console.log("Create/Update Ground Projection", backgroundTexture.name);

            this._projection?.removeFromParent();

            try {
                this._projection = new GroundProjection(backgroundTexture, this._height, this._radius, 64);
                configureGroundProjectionMaterial(this._projection.material, backgroundTexture);
            }
            catch (e) {
                console.error("Error creating three GroundProjection", e);
                return;
            }
            this._projection.position.y = this._height - offset;
            this._projection.name = "GroundProjection";
            setVisibleInCustomShadowRendering(this._projection, false);
        }
        else {
            needsNewAutoFit = false;
        }

        if (!this._projection.parent)
            this.gameObject.add(this._projection);

        if (this.autoFit && needsNewAutoFit) {
            // TODO: should also update the radius (?)
            this._projection.updateWorldMatrix(true, true);
            const box = getBoundingBox(this.context.scene.children, [this._projection]);

            const floor_y = box.min.y;
            if (floor_y < Infinity) {
                const wp = getTempVector();
                wp.x = box.min.x + (box.max.x - box.min.x) * .5;
                const scale = getWorldScale(this.gameObject).x;
                wp.y = floor_y + (this._height * scale) - offset;
                wp.z = box.min.z + (box.max.z - box.min.z) * .5;
                setWorldPosition(this._projection, wp);
            }

            if (debug) Gizmos.DrawWireBox3(box, 0x00ff00, 5);
        }

        /* TODO realtime adjustments aren't possible anymore with GroundedSkybox (mesh generation)
        this.env.scale.setScalar(this._scale);
        this.env.radius = this._radius;
        this.env.height = this._height;
        */

        this.updateProjectionMaterial(backgroundTexture, true);

        this._lastBackground = backgroundTexture;
        this._lastHeight = this._height;
        this._lastRadius = this._radius;
        this._needsTextureUpdate = false;
    }

    private _blurrynessShader: ShaderMaterial | null = null;
    private _lastBlurriness: number = -1;
    
    private updateProjectionMaterial(texture: Texture, forceTextureUpdate = false) {
        if (!this._projection) {
            return;
        }

        const blurriness = this.context.scene.backgroundBlurriness ?? 0;
        const useCubeUvBlur = texture.mapping === CubeUVReflectionMapping;

        let targetTexture = texture;
        if (!useCubeUvBlur && blurriness > 0.001) {
            const hasBlurredTextureAssigned = !!this._projection.material.map && this._projection.material.map !== texture;
            if (forceTextureUpdate || !hasBlurredTextureAssigned) {
                targetTexture = this.updateBlurriness(texture, blurriness);
            }
            else if (this._projection.material.map) {
                targetTexture = this._projection.material.map;
            }
        }

        if (this._projection.material.map !== targetTexture) {
            this._projection.material.map = targetTexture;
        }

        const appliedTexture = this._projection.material.map ?? texture;
        appliedTexture.mapping = texture.mapping;
        configureGroundProjectionMaterial(this._projection.material, appliedTexture);

        const shaderUniforms = getGroundProjectionShaderUniforms(this._projection.material);
        shaderUniforms.needleGroundProjectionBlurriness.value = useCubeUvBlur ? blurriness : 0;
        shaderUniforms.needleGroundProjectionBackgroundIntensity.value = this.context.scene.backgroundIntensity ?? 1;

        const wasTransparent = this._projection.material.transparent;
        this._projection.material.transparent = (this.context.xr?.isAR === true && this.arBlending > 0.000001) ?? false;
        shaderUniforms.needleGroundProjectionBlending.value = this._projection.material.transparent ? this.arBlending : 0;
        shaderUniforms.needleGroundProjectionAlphaFactor.value = this.context.isInPassThrough ? 0.95 : 1;

        if (wasTransparent !== this._projection.material.transparent) {
            this._projection.material.needsUpdate = true;
        }

        this._projection.material.depthTest = true;
        this._projection.material.depthWrite = false;
        this._lastBlurriness = blurriness;
        this._needsTextureUpdate = false;
    }

    private updateBlurriness(texture: Texture, blurriness: number): Texture {
        if (debug) console.log("Update Blurriness", blurriness);
        this._blurrynessShader ??= new ShaderMaterial({
            name: "GroundProjectionBlurriness",
            uniforms: {
                map: { value: texture },
                blurriness: { value: blurriness }
            },
            vertexShader: blurVertexShader,
            fragmentShader: blurFragmentShader
        });
        this._blurrynessShader.depthWrite = false;
        this._blurrynessShader.uniforms.map.value = texture;
        this._blurrynessShader.uniforms.blurriness.value = blurriness;
        texture.needsUpdate = true;
        const blurredTexture = Graphics.copyTexture(texture, this._blurrynessShader);
        blurredTexture.mapping = texture.mapping;
        return blurredTexture;
    }

}


const blurVertexShader = `
  varying vec2 vUv;

  void main() {
    vUv = uv;
    gl_Position = projectionMatrix * modelViewMatrix * vec4(position, 1.0);
  }
`;

// Fragment Shader
const blurFragmentShader = `
  uniform sampler2D map;
  uniform float blurriness;
  varying vec2 vUv;

  const float PI = 3.14159265359;

  // Gaussian function
  float gaussian(float x, float sigma) {
    return exp(-(x * x) / (2.0 * sigma * sigma)) / (sqrt(2.0 * PI) * sigma);
  }

  // Custom smoothstep function for desired falloff
  float customSmoothstep(float edge0, float edge1, float x) {
    float t = clamp((x - edge0) / (edge1 - edge0), 0.0, 1.0);
    return t * t * (3.0 - 2.0 * t);
  }

  void main() {
    vec2 center = vec2(0.0, 0.0);
    vec2 pos = vUv;
    pos.x = 0.0; // Only consider vertical distance
    float distance = length(pos - center);
    
    // Calculate blur amount based on custom falloff
    float blurAmount = customSmoothstep(0.5, 1.0, distance * 2.0);
    blurAmount = clamp(blurAmount, 0.0, 1.0); // Ensure blur amount is within valid range

    // Gaussian blur
    vec2 pixelSize = 1.0 / vec2(textureSize(map, 0));
    vec4 color = vec4(0.0);
    float totalWeight = 0.0;
    int blurSize = int(60.0 * min(1.0, blurriness) * blurAmount); // Adjust blur size based on distance and blurriness
    if (blurSize <= 0) {
      gl_FragColor = texture2D(map, vUv);
      return;
    }
    float lodLevel = log2(float(blurSize)) * 0.5; // Compute LOD level

    for (int x = -blurSize; x <= blurSize; x++) {
        for (int y = -blurSize; y <= blurSize; y++) {
            vec2 offset = vec2(float(x), float(y)) * pixelSize * blurAmount;
            float weight = gaussian(length(vec2(float(x), float(y))), 1000.0 * blurAmount); // Use a fixed sigma value
            color += textureLod(map, vUv + offset, lodLevel) * weight;
            totalWeight += weight;
        }
    }

    color = totalWeight > 0.0 ? color / totalWeight : texture2D(map, vUv);

    gl_FragColor = color;

    // #include <tonemapping_fragment>
    // #include <colorspace_fragment>
    
    // Uncomment to visualize blur amount
    // gl_FragColor = vec4(blurAmount, 0.0, 0.0, 1.0);
  }
`;